11institutetext: University of Maryland Medical Intelligent Imaging (UM2ii)
Center, Baltimore, MD
11email<EMAIL_ADDRESS>
# High-Throughput AI Inference for Medical Image Classification and
Segmentation using
Intelligent Streaming

Pranav Kulkarni  Sean Garin  Adway Kanhere
Eliot Siegel  Paul H. Yi  Vishwa S. Parekh

###### Abstract

As the adoption of AI systems within the clinical setup grows, limitations in
bandwidth could create communication bottlenecks when streaming imaging data,
leading to delays in patient diagnosis and treatment. As such, healthcare
providers and AI vendors will require greater computational infrastructure,
therefore dramatically increasing costs. To that end, we developed intelligent
streaming, a state-of-the-art framework to enable accelerated, cost-effective,
bandwidth-optimized, and computationally efficient AI inference for clinical
decision making at scale111Patent Pending. For classification, intelligent
streaming reduced the data transmission by 99.01% and decoding time by 98.58%,
while increasing throughput by 27.43x. For segmentation, our framework reduced
data transmission by 90.32%, decoding time by 90.26%, while increasing
throughput by 4.20x. Our work demonstrates that intelligent streaming results
in faster turnaround times, and reduced overall cost of data and transmission,
without negatively impacting clinical decision making using AI systems.

###### Keywords:

Deep learning Intelligent streaming Compression Progressive decoding HTJ2K
Classification Segmentation

##  1 Introduction

The utility of artificial intelligence (AI) systems in medical imaging has
grown rapidly due to the curation of large-scale datasets containing 2D images
and 3D volumes across image modalities (e.g., X-Rays, MRI, CT, and more).
Using such datasets, AI systems can learn complex relationships between
different tissue types and provide the capability to detect/predict the
presence of disease. As a result, numerous companies have begun to develop and
deploy intelligent imaging solutions in hospital systems across a series of
different clinical applications, such as Amazon’s HealthLake [1] and Google’s
Medical Imaging Suite [7].

Deploying intelligent imaging solutions typically involves healthcare
providers streaming large volumes of medical image data to AI vendors for
analysis. However, limitations in bandwidth create communication bottlenecks
when streaming medical images over the internet, leading to delays in patient
diagnosis and treatment. Furthermore, as the adoption of AI systems in the
clinical setup continues to increase, healthcare providers and AI vendors will
require greater computational infrastructure, therefore dramatically
increasing costs. As smaller family practices and rural clinics without high-
quality computational infrastructure will struggle to keep pace with larger
hospital systems, health disparities in rural populations will undoubtedly
widen [4]. To that end, we developed an intelligent streaming framework based
on progressive encoding for high throughput, bandwidth optimized, and cost
effective AI inference for clinical decision making at scale.

Progressive encoding has revolutionized the way we consume media content over
the internet. Prior to the advent of progressive encoding, users had to
download entire media files before being able to view them. This process was
often slow and inefficient, especially for users with slow internet
connections or limited storage and memory on their devices. Progressive
encoding works by encoding the imaging data into a series of progressive
scans, each with an increasing level of detail. These techniques allow for the
efficient delivery of large datasets in real-time and have proven useful in
enhancing user experience in several applications, including video streaming
and web applications. However, the utility of progressive encoding in
optimizing the AI inference workflow for clinical decision making at scale has
not yet been explored. Therefore, we developed an intelligent streaming
framework to determine and stream an optimal subset of the image byte stream
that can deliver the same performance as the full image byte stream while
enabling faster turnaround times with reduced overall cost of data storage and
transmission. We evaluated the intelligent streaming framework for optimal
streaming of 2D and 3D medical imaging datasets for the tasks of
classification and segmentation across different organs and imaging
modalities.

##  2 Methods

###  2.1 High-Throughput JPEG 2000 (HTJ2K)

High-Throughput JPEG 2000 (HTJ2K) is an extension to the existing JPEG 2000
(J2K) codec, that – unlike the original JPEG standard – enables faster
decoding, higher quality, lossless encoding with high scalability, and
progressive decoding of images [5, 16]. The HTJ2K codec utilizes new
lightweight and efficient block-coder algorithm that yields up to 10-30x
increase in throughput for lossless coding and a 6-10% increase in efficiency
when compared to J2K [17]. Furthermore, HTJ2K’s lower computational complexity
and support for parallelism results in a throughput on par, or even better,
than the heavily optimized original JPEG standard in single and multi-threaded
applications [16]. There has been a push to incorporate HTJ2K into the
clinical workflow by extending the DICOM transfer syntax standard to add
support for HTJ2K [8]. More recently, Amazon Web Services’ HealthLake Imaging,
a cloud-based medical imaging storage service, utilized the advancements with
HTJ2K as part of their compression pipeline to reduce storage costs and image
retrieval times for imaging data [1].

Figure 1: A visualization of the image formed by 25% of entire byte stream
using the original JPEG and HTJ2K standards respectively.

Imaging data can be progressively encoded using lossless HTJ2K as $N$ scans
(also known as decompositions or subsets of the byte stream), each with a
native resolution as a factor of $2^{N}$ of the full image resolution. The
native resolution of a scan refers to the proportion of pixels of the full
resolution image transmitted prior to any interpolation. Unlike the sequential
streams of the original JPEG standard, HTJ2K streams are transmitted as a
sequence of these $N$ scans, as visualized in Figure 1. For example, a 2D
chest x-ray (CXR) with a resolution of 256x256 can be progressively encoded as
3 scans with native resolutions 64x64, 128x128, and 256x256. This allows HTJ2K
to not just achieve a high compression ratio, but also enables quick access to
smaller resolution versions of the image, thus minimizing data duplication.
While our intelligent streaming framework is compatible with both codecs, in
this work, we utilize HTJ2K to progressively encode medical imaging datasets
and evaluate our framework across classification and segmentation tasks.

###  2.2 Intelligent Streaming

The intelligent streaming framework consists of three primary components, as
visualized in Figure 2.

Progressive Encoder: The role of the encoder is to progressively encode 2D and
3D imaging data across different imaging modalities and imaging resolutions
into a sequence of $N$ scans. Suppose the original image data consists of $M$
images with native resolutions $X_{i}\ \forall i\in\\{1,2,..,M\\}$. We impose
a condition that the native resolution of the first scan must exceed
$X_{\text{min}}$. We arbitrarily set $X_{\text{min}}$ as 32x32 to recover
sufficient information for our analysis from the first scan. In aggregate, the
number of scans $N$ was determined using the following equation:

$N=\log_{2}\left(\frac{\sum_{i=1}^{M}X_{i}}{X_{\text{min}}M}\right)-1$ (1)

Stream optimizer: The role of the stream optimizer is to determine the optimal
subset of bytes that can be streamed without negatively impacting the
performance of the AI model. Using an internal validation set, the stream
optimizer determines the lowest progressive scan for which there is no
statistically significant difference in the AI model performance between the
full resolution scan and the progressive scan.

Progressive Decoder: The role of the progressive decoder is to decode a byte
stream into a progressive scan with corresponding native resolution. For
example, if the stream optimizer streams only $2\%$ of the complete byte
stream, the progressive decoder, would decode that bytes stream and
reconstruct a complete image, as demonstrated in Figure 1.

Figure 2: An overview of the intelligent streaming framework.

###  2.3 Experimental Setup

We evaluated our intelligent streaming framework in two experimental setups.
The first experiment employs a 2D classification model to detect abnormalities
in the thorax using CXR data, while the second experiment employs a 3D
segmentation model to segment the liver and spleen using computed tomography
(CT) data. In both cases, models were trained using local data, validated
using held-out internal test sets, and evaluated on unseen data transmitted by
the hospital system. Furthermore, in both cases, the data transmitted, decode
time, and throughput (number of images processed by the AI pipeline per
second) were measured as data metrics and compared to the original non-HTJ2K
versions of the data.

Classification: We trained a multi-label DenseNet121 model architecture [9],
initialized with ImageNet weights, on the NIH Chest X-Ray (‘NIH’) dataset [18]
using transfer learning. The NIH dataset consists of 14 disease labels with
112,120 frontal-view CXRs from 30,805 patients. The dataset was randomly
divided into training (70%), validation (10%), and testing (20%, $N=22,330$)
splits while ensuring no patient appears in more than one split. All images
were downsampled to 224x224 and normalized between 0 and 1. Images within the
held-out NIH test set were originally encoded as PNG with a normalized
resolution of 1024x1024. We progressively encoded the NIH test set using
lossless HTJ2K (containing 5 scans) to yield a maximum compression ratio of
256:1. We determined the optimal scan by evaluating the model across each scan
of HTJ2K encoded and original PNG encoded data.

The optimal task-specific byte stream scan obtained using the NIH data was
evaluated on an external dataset obtained by randomly sampling $N=10,000$ CXRs
from the MIMIC-CXR-JPG dataset [6, 10]. All uncertain labels were treated as
negatives. Unlike the NIH test set, images within the MIMIC test set were
originally encoded using the original, and heavily optimized, JPEG standard
with resolutions ranging from 1470x1461 to 4280x3520. The MIMIC test set was
progressively encoded using lossless HTJ2K (consisting of 7 scans) to yield an
maximum compression ratio of 4096:1. Model performance was evaluated across
each scan of the HTJ2K encoded and original JPEG data using the binary cross
entropy loss (BCE), area under the receiver operating characteristic (AUROC)
and area under the precision recall curve (AUPRC). In a multi-label scenario,
the AUROC and AUPRC scores are defined as the mean score across all disease
labels. We also calculated the 95% confidence interval (CI) for the AUROC
scores and compared them using bootstrapping and a paired two-tailed t-test.
Statistical significance was defined as $p<0.05$.

Segmentation:  We trained a 3D-UNet [3] using the MONAI framework [2] for
segmentation of liver and spleen using the corresponding datasets from the
Medical Segmentation Decathlon (MSD) training sets. For this study, liver
tumor annotations were discarded. We sampled volumes for the held-out MSD
liver test set ($N=25$) and spleen test set ($N=9$) prior to training. The
model architecture consisted of 5 two-convolutional residual units with
down/up-sampling by a factor of 2 at each unit. All image volumes were resized
to 256x256x128 and voxel values normalized between 0 and 1. Random foreground
patches of size 128x128x32 were extracted from each volume such that the
center voxel of each patch belonged to either the foreground or background
class.

Images within the held-out MSD liver and spleen test sets were originally
encoded as NifTI volumes with a normalized resolution of 512x512x$N$. One
limitation of the HTJ2K standard is its inability to encode 3D imaging data.
As a workaround, for all volumes, we progressively encoded each 2D slice using
lossless HTJ2K (consisting of 4 scans) to yield a maximum compression ratio of
64:1. All slices were decoded individually, reconstructed into a volume, and
evaluated with the model, along with the original NifTI volumes. We further
evaluated the optimal byte stream scan obtained from internal validation on an
external dataset consisting $N=30$ volumes from the Beyond the Cranial Vault
(BTCV) dataset [12]. All volumes were similarly encoded as NifTI with a
normalized resolution of 512x512x$N$, and thus were processed and evaluated
similarly. The mean DICE score between the ground truth and model segmentation
was measured as the model metric and compared using a paired t-test.
Statistical significance was again defined as $p<0.05$.

##  3 Results

###  3.1 Classification

On the held-out NIH test set, we observed that the classification model
performance for all progressive scans was comparable to using full resolution
PNG images, while needing only a fraction of the amount of data transmitted
and decoding time (apart from scan 1), as shown in Table 1. For example, scan
2 yielded a comparable mean AUROC of 0.81 ($p=0.74$), while transmitting
97.25% less data, taking 95.60% less time to decode, and a 9.90x increase in
throughput. Despite PNG coding being heavily optimized due to its wide
adoption, we observed that decoding the full resolution images from the
complete byte stream required 4.41% less data to be transmitted and took 7.40%
less time to decode.

Table 1: Comparing model and data metrics for 2D classification on the
original PNG and intelligently streamed HTJ2K versions of the held-out NIH
test set. 95% CI for mean AUROC scores provided in brackets, along with
p-values (in grey).

On the MIMIC external test set, apart from the first two scans, we observed
similar, if not better, model performance for all progressive scans, as shown
in Table 2. For example, scan 3 yielded a slightly better mean AUROC of 0.75
($p=0.77$), while transmitting 99.01% less data, taking 98.58% less time to
decode, and a 27.43x increase in throughput. Due to the efficient and heavily
optimized nature of the original JPEG standard, the final scan required 1.78x
more data to be transmitted and took 1.16x more time to decode, thus impacting
the overall throughput.

Table 2: Comparing model and data metrics for 2D classification on the
original JPEG and intelligently streamed HTJ2K versions of the MIMIC external
test set. 95% CI for mean AUROC scores provided in brackets, along with
p-values (in grey).

###  3.2 Segmentation

For segmentation tasks, we observed similar findings as classification. On the
held-out MSD liver and spleen test sets, the both segmentation models yielded
high DICE scores across all scans and achieved a comparable DICE score of
$0.94\pm 0.02$ from the scans 3 and 4, as shown in Table 3(b). However, scan 1
for spleen segmentation performed poorly. Across both tasks, scan 3
demonstrates the superiority of HTJ2K over NifTI compression when transmitting
3D imaging data, as shown in Table 3(b). On the MSD liver test set, we
observed that scan 3 had no impact on AI inference, while transmitting 91.6%
less data, 87.92% less time to decode, and a 5.45x increase in throughput. On
the MSD spleen test set, we also observed that scan 3 had no impact on AI
inference, while transmitting 90.70% less data, 92.34% less time to decode,
and a 2.76x increase in throughput.

When it comes to transmitting full-resolution lossless imaging data, we yet
again observe the superiority of HTJ2K, as shown in Table 3(b). On the MSD
liver test set, we observed that lossless HTJ2K data required 76.94% less data
to be transmitted, 76.33% less time to decode, and a 3.55x increase in
throughput. Similarly, the HTJ2K encoded MSD spleen test set required 67.66%
less data to transmit, 86.56% less time to decode, and a 2.40x increase in
throughput.

Table 3: Comparing data and model metrics for 3D segmentation of liver and
spleen on the original NifTI and intelligently streamed HTJ2K versions of the
held-out MSD liver and spleen test sets.

(a) Liver

(b) Spleen

On the BTCV external test set, we observed similar results for liver
segmentation, but we a significant impact to the performance of spleen
segmentation. In short, liver segmentation performed well for volumes with
higher compression ratios while spleen segmentation was severely impacted by
it, as shown Table 4. However, the scan 3 yielded excellent DICE scores of
$0.88\pm 0.09$ and $0.81\pm 0.17$ on both liver and spleen segmentation tasks
respectively, while requiring 90.32% less data to transmit, 90.26% less time
to decode, and a 4.20x increase in throughput. Furthermore, we again
demonstrated that lossless HTJ2K volumes required 73.12% less data to transmit
and 81.20% less time to decode when compared to the original NifTI compressed
volumes.

Table 4: Comparing data and model metrics for 3D segmentation of liver and
spleen on the original NifTI and intelligently streamed HTJ2K versions of the
BTCV external test sets.

##  4 Discussion

As the adoption of AI systems within the clinical setup grows, limitations in
bandwidth could create bottlenecks and lead to delays in patient diagnosis and
treatment. Our work demonstrates that intelligent streaming has the potential
to significantly reduce storage and bandwidth requirements and time to decode
imaging data, while increasing overall throughput, in medical imaging tasks.
Due to the robustness of deep learning models for CXR classification [11, 14,
15], we achieved not only high compression ratios, but a non-significant
impact to AI model performance. Despite excellent performance on liver
segmentation, our model struggled to correctly segment the spleen on the
external BTCV test set. Prior literature has also suggested segmentation tasks
are more sensitive to compression [13]. Apart from being a much harder task
due to the spleen’s smaller size compared to the liver, we believe that the
performance of spleen segmentation on the external BTCV dataset was degraded
due to the model’s poor generalizability to varying intensity values [19]. We
observed that volumes within the intensity range of the original training data
correctly segmented the spleen. Despite these results, our study demonstrates
the superiority of HTJ2K for transmitting image volumes by significantly
reducing transmission and decoding overheads, while increasing overall
throughput. Furthermore, we demonstrate that intelligent streaming will result
in faster turnaround times, and reduced overall cost of data storage and
transmission, without negatively impacting clinical decision making.

## References

  * [1] Amazon Healthlake. https://aws.amazon.com/healthlake/, accessed: 2023-02-27
  * [2] Cardoso, M.J., Li, W., Brown, R., Ma, N., Kerfoot, E., Wang, Y., Murrey, B., Myronenko, A., Zhao, C., Yang, D., et al.: Monai: An open-source framework for deep learning in healthcare. arXiv preprint arXiv:2211.02701 (2022)
  * [3] Çiçek, Ö., Abdulkadir, A., Lienkamp, S.S., Brox, T., Ronneberger, O.: 3d u-net: learning dense volumetric segmentation from sparse annotation. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19. pp. 424–432. Springer (2016)
  * [4] Douthit, N., Kiv, S., Dwolatzky, T., Biswas, S.: Exposing some important barriers to health care access in the rural usa. Public health 129(6), 611–620 (2015)
  * [5] Foos, D.H., Muka, E., Slone, R.M., Erickson, B.J., Flynn, M.J., Clunie, D.A., Hildebrand, L., Kohm, K.S., Young, S.S.: Jpeg 2000 compression of medical imagery. In: Medical Imaging 2000: PACS Design and Evaluation: Engineering and Clinical Issues. vol. 3980, pp. 85–96. SPIE (2000)
  * [6] Goldberger, A.L., Amaral, L.A., Glass, L., Hausdorff, J.M., Ivanov, P.C., Mark, R.G., Mietus, J.E., Moody, G.B., Peng, C.K., Stanley, H.E.: Physiobank, physiotoolkit, and physionet: components of a new research resource for complex physiologic signals. circulation 101(23), e215–e220 (2000)
  * [7] Google medical imaging suite. https://cloud.google.com/medical-imaging, accessed: 2023-02-27
  * [8] Hafey, C., Wallace, B.: HTJ2K transfer syntax support. https://www.dicomstandard.org/docs/librariesprovider2/default-document-library/htj2k-transfer-syntax-support.docx (2022)
  * [9] Huang, G., Liu, Z., Van Der Maaten, L., Weinberger, K.Q.: Densely connected convolutional networks. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 4700–4708 (2017)
  * [10] Johnson, A.E., Pollard, T.J., Greenbaum, N.R., Lungren, M.P., Deng, C.y., Peng, Y., Lu, Z., Mark, R.G., Berkowitz, S.J., Horng, S.: Mimic-cxr-jpg, a large publicly available database of labeled chest radiographs. arXiv preprint arXiv:1901.07042 (2019)
  * [11] Koff, D.A., Shulman, H.: An overview of digital compression of medical images: can we use lossy image compression in radiology? Journal-Canadian Association of Radiologists 57(4), 211 (2006)
  * [12] Landman, B., Xu, Z., Igelsias, J., Styner, M., Langerak, T., Klein, A.: Miccai multi-atlas labeling beyond the cranial vault–workshop and challenge. In: Proc. MICCAI Multi-Atlas Labeling Beyond Cranial Vault—Workshop Challenge. vol. 5, p. 12 (2015)
  * [13] Liu, Z., Xu, X., Liu, T., Liu, Q., Wang, Y., Shi, Y., Wen, W., Huang, M., Yuan, H., Zhuang, J.: Machine vision guided 3d medical image compression for efficient transmission and accurate segmentation in the clouds. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 12687–12696 (2019)
  * [14] Sabottke, C.F., Spieler, B.M.: The effect of image resolution on deep learning in radiography. Radiology: Artificial Intelligence 2(1), e190015 (2020)
  * [15] Saffor, A., Ramli, A.R., Ng, K.H.: A comparative study of image compression between jpeg and wavelet. Malaysian Journal of computer science 14(1), 39–45 (2001)
  * [16] Taubman, D., Naman, A., Mathew, R., Smith, M., Watanabe, O.: High throughput JPEG 2000 (HTJ2K): Algorithm, performance and potential. International Telecommunications Union (ITU) pp. 15444–15 (2019)
  * [17] Taubman, D., Naman, A., Mathew, R., Smith, M., Watanabe, O., Lemieux, P.A.: High throughput JPEG 2000 (HTJ2K) and the JPH file format: a primer. https://ds.jpeg.org/whitepapers/jpeg-htj2k-whitepaper.pdf (2020)
  * [18] Wang, X., Peng, Y., Lu, L., Lu, Z., Bagheri, M., Summers, R.M.: Chestx-ray8: Hospital-scale chest x-ray database and benchmarks on weakly-supervised classification and localization of common thorax diseases. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 2097–2106 (2017)
  * [19] Zhang, L., Wang, X., Yang, D., Sanford, T., Harmon, S., Turkbey, B., Wood, B.J., Roth, H., Myronenko, A., Xu, D., et al.: Generalizing deep learning for medical image segmentation to unseen domains via deep stacked transformation. IEEE transactions on medical imaging 39(7), 2531–2540 (2020)


# Petz recovery maps: Geometrical aspects and an analysis for qudit channels

Lea Lautenbacher  Institut für Theoretische Physik, Albert-Einstein-Allee 11,
Universität Ulm, D-89069 Ulm, Germany  Vinayak Jagadish  Instytut Fizyki
Teoretycznej, Uniwersytet Jagielloński, Łojasiewicza 11, 30-348 Kraków, Poland
Francesco Petruccione  School for Data Science and Computational Thinking,
Stellenbosch University, Stellenbosch 7600, South Africa  National Institute
for Theoretical and Computational Sciences (NITheCS), Stellenbosch 7600, South
Africa  Nadja K. Bernardes  Departamento de Física, Centro de Ciências Exatas
e da Natureza, Universidade Federal de Pernambuco, Recife-PE 50670-901, Brazil

###### Abstract

Using the Petz map, we investigate the potential of state recovery when
exposed to dephasing and amplitude-damping channels. Specifically, we analyze
the geometrical aspects of the Petz map for the qubit case, which is linked to
the change in the volume of accessible states. Our findings suggest that the
geometrical characterization can serve as a potent tool for understanding the
details of the recovery procedure. Furthermore, we extend our analysis to
qudit channels by devising a state-independent framework that quantifies the
ability of the Petz map to recover a state for any dimension. Under certain
conditions, the dimensionality plays a role in state recovery.

##  I Introduction

A bijective function, also known as an invertible function, refers to a linear
map between two sets such that each element of one set corresponds to exactly
one element of the other, and every element of the sets is paired. In a
physical system, a similar idea appears as a reversible process related to the
inversion of the direction of a process that maps the set of outputs into the
set of inputs. The reversibility of a process is usually connected to
thermodynamics and the variation of entropy Norton2016 . In the quantum realm,
only a unitary evolution that describes the dynamics of an isolated quantum
system is reversible. However, no quantum system is truly isolated petruccione
; RivasBook ; haroche_exploring_2006 , but is inexorably coupled to the
environment. The type of environment which affects the quantum system of
interest indeed depends on the physical setting under consideration. This
coupling to the environment makes the quantum system noisy. Therefore, to use
quantum systems for information processing tasks, one needs to combat the
effects due to the interaction with the environment. This also amounts to
studying whether or in what circumstances they could be undone. This naturally
leads to the question of quantum recovery channels. A standard example of a
recovery channel is in quantum error correction where the states in the “code
subspace” are corrected by the appropriate recovery map.

A quantum system is associated with a Hilbert space $\mathcal{H}$, such that
the set of all possible states can be defined as
$\mathcal{D}(\mathcal{H})=\\{\rho\in\mathcal{L}(\mathcal{H})\\}$, where
$\mathcal{L}(\mathcal{H})$ represents the set of linear operators acting on
$\mathcal{H}$. The most general transformations on the quantum system of
interest are characterized by a family of linear maps that are completely
positive and trace preserving (CPTP) known as quantum channels Quanta77 . For
certain quantum channels and input states satisfying an information theoretic
criterion, a universal recovery operation was proposed by Petz Petz86 . One
should keep in mind that the recovery map should preserve the existing
correlations between the system and the environment. The Petz recovery map has
been widely explored in the context of quantum data processing inequality
Petz86 ; Petz88 ; Petz03 ; Junge18 , quantum error correction Barnum02 ,
deriving fluctuation theorems Buscemi21 ; Cenxin21 , to name a few. Recently,
a physical protocol has been proposed to achieve this recovery map Kwon22 and
a quantum algorithm has been developed to implement it Gilyen22 .

Although the Petz map is a powerful recovery map, it is built to perfectly
work in the context where the quantum relative entropy does not increase in
time. There are two crucial ingredients in constructing the Petz map: the
original dynamical map (the map intended to be reverted) and a reference state
(also called a prior state). For one-qubit decoherence channels the best
strategy using the Petz map was studied in Lea21 . Other general recovery
strategies have been explored in Surace22 by extending the Bayes retrodiction
theorem.

In this work, we look into the geometrical aspects of the Petz recovery maps
for a qubit and we explore the Petz recovery map in higher dimensions. Similar
to what was done in Ref. Lea21 for the qubit case, we would like to verify the
use of this map in a general context. The input states considered are the full
set of states and not just the states for which the Petz map is intended to
work perfectly. Two paradigmatic quantum channels are analyzed: the dephasing
and amplitude-damping channels. By exploring the reference state, an optimal
Petz map is constructed. In Ref. Lea21 , it was shown that the optimal Petz
recovery map is worst than simply returning the output state, i.e. using the
identity channel as the “recovery channel”. This was spotted by analyzing the
average fidelity between the initial and recovered states. Since the average
fidelity between a target pure state and a randomly chosen state decreases
with the dimension of the systems Zyczkowski05 , we asked ourselves if the
Petz recovery map could have a better performance in this general scenario
with $d$-dimensional systems. To quantify this, we analyzed the distance
between the corresponding Choi matrices which showed a clear dependence
between this distance and the dimension.

The article is structured as follows. We review the basic concepts of quantum
channels and Petz recovery maps in Sec. II. In Sec. III we present the volume
of accessible states for a qubit under the Petz recovery channel after the
action of a unital and a non-unital decoherence channel. We explore the
behavior of the Petz map in higher dimensions and its implications on the
recovery strategy in Sec. IV. We conclude in Sec. V.

##  II Preliminaries

###  II.1 Quantum Channels

The most general transformation of a quantum system is described by completely
positive and trace-preserving maps known as _quantum channels_ Quanta77 .
Consider a linear map
$\Lambda:\mathcal{L}(\mathcal{H})\mapsto\mathcal{L}(\mathcal{H})$, where
$\mathcal{L}(\mathcal{H})$ denotes the vector space of linear operators acting
on the Hilbert space $\mathcal{H}$. We assume that ${\rm dim}\,\mathcal{H}=d$.
Here, we briefly discuss a few representations $\Lambda$ that we make use of.

Fixing an orthonormal basis $\\{|1\rangle,\ldots,|d\rangle\\}$ in
$\mathcal{H}$, one can define the operator $J(\Lambda)$, usually referred to
as the Choi matrix Choi75 ; jam72 ,

$J(\Lambda)=\sum_{i,j=1}^{d}|i\rangle\langle j|\otimes\Lambda(|i\rangle\langle
j|),$ (1)

and the corresponding $d^{2}\times d^{2}$ Hermitian matrix

$J(\Lambda)_{ij,kl}:=\langle i\otimes j|J(\Lambda)|k\otimes l\rangle=\langle
j|\Lambda(|i\rangle\langle k|)|l\rangle.$ (2)

The map $\Lambda$ is completely positive if and only if $J(\Lambda)\geq 0$.

Another useful representation is the $\mathcal{A}$ matrix obtained by
reshuffling $J(\Lambda)$ as follows:

$\mathcal{A}=J(\Lambda)^{\mathcal{R}},~{}~{}~{}\mathcal{A}_{ik,jl}=J(\Lambda)_{ij,kl}.$
(3)

If $\Lambda$ is completely positive, it can always be represented in the
operator-sum form as

$\displaystyle\Lambda(X)$ $\displaystyle=$
$\displaystyle\sum_{i}K_{i}XK_{i}^{\dagger},$
$\displaystyle\sum_{i}K_{i}^{\dagger}K_{i}$ $\displaystyle=$
$\displaystyle\mathbb{1}\thinspace(\mathrm{Trace\,Preservation}).$ (4)

The operators $K_{i}$ are usually referred to as the Kraus operators.

There exists another representation referred to as the Bloch representation.
To this end, let $\Gamma_{i}$ $(i=0,1,\ldots,d^{2}-1)$ denote a Hermitian
orthonormal basis in $\mathcal{L}(\mathcal{H})$, ${\rm
Tr}(\Gamma_{i}\Gamma_{j})=\delta_{ij}$, and let
$\Gamma_{0}=\mathbb{1}/\sqrt{d}$. The map acting on a general density matrix
$\rho$ of a qudit can be represented as

$\Lambda(\rho)=\Lambda\Big{(}\frac{\mathbb{1}+\vec{r}.\vec{\Gamma}}{d}\Big{)}=\frac{\mathbb{1}+(M\vec{r}+\vec{\tau}).\vec{\Gamma}}{d},$
(5)

where $\vec{\tau}\in\mathbb{R}^{d^{2}-1}$, $M$ is a real square matrix of
order $d^{2}-1$ and $\vec{r}$ is referred to as the Bloch-vector corresponding
to the qudit. This represents the transformation of the Bloch vector,
$\vec{r}\to(M\vec{r}+\vec{\tau})$, and is commonly known as the _affine_
representation.

A quantum channel is said to be _unital_ if the identity is preserved. i.e.,
$\Lambda(\mathbb{1})=\mathbb{1}$.

The _dual map_ , $\Lambda^{\dagger}$, in the operator sum representation is
defined as

$\Lambda(X)=\sum_{i}K_{i}^{\dagger}XK_{i}.$ (6)

Let us now briefly discuss two channels that we analyze in this work.

####  II.1.1 Dephasing Channels

The dephasing channel is one of the most common decoherence channels. It is a
unital map that destroys the relative phases between the computational basis
states, $\\{|k\rangle\\}$. The uniform dephasing map is defined as

$\Lambda_{D}(\rho)=(1-p)\rho+p\sum_{k=0}^{d-1}P_{k}\rho P_{k},$ (7)

with the projector $P_{k}=|k\rangle\langle k|$. The parameter $p\in[0,1]$
characterizes the noise strength ($p=0$, no noise; $p=1$ maximal noise). For
$p=1$, the coherences vanish completely.

####  II.1.2 Damping Channels

The amplitude damping channel $\Lambda_{AD}$ is a non-unital map that can
model the spontaneous emission of a $d$-level atom, i.e. the energy relaxation
from an excited state to a lower energy state of the system. Here, we consider
the damping to the ground state $\ket{0}$ for a qudit system of dimension $d$.
Let $\\{\ket{i}\\}_{i=0,1,2,\cdots,d-1}$ denote an orthonormal basis of the
Hilbert space of the qudit. The map is defined as

$\displaystyle\Lambda_{AD}(\rho)$ $\displaystyle=\sum_{i=0}^{d-1}K_{i}\rho
K_{i}^{\dagger},\mathrm{with}$ $\displaystyle K_{0}$
$\displaystyle=|0\rangle\langle 0|+\sqrt{1-p}\sum_{i=1}^{d-1}|i\rangle\langle
i|$ $\displaystyle K_{i=1..d-1}$ $\displaystyle=\sqrt{p}|0\rangle\langle i|.$
(8)

For simplicity, here we include the contributions from transitions from other
levels to the state $\ket{0}$ alone and ignore the transitions between other
levels. One could also have damping maps with varying $p$. An example for the
case of a qutrit is discussed in Appendix A.

###  II.2 Non-unitality measure

Another property we take into account is the non-unitality of the map.
Consider a quantum channel $\Lambda$, such that
$\Lambda(\mathbb{1}_{d}/d)=\rho^{\prime}$, with $\mathbb{1}_{d}/d$ the
maximally mixed state for a qudit with dimension $d$. We define a non-
unitality measure $\mathcal{N}$ as the deviation of the output state
$\rho^{\prime}$ from the maximally mixed state. By means of state
distinguishability, we write

$\mathcal{N}=\frac{1}{2}||\rho^{\prime}-\frac{\mathbb{1}_{d}}{d}||_{1},$ (9)

where $\mathcal{N}=0$ iff $\Lambda$ is unital and
$||A||_{1}=\text{Tr}\\{\sqrt{AA^{\dagger}}\\}$. This way of estimating the
non-unitality of the map is convenient, especially for higher dimensions.

###  II.3 Petz Recovery Maps

An important definition we use in this work is the recoverability property of
a map. We say a map is recoverable if there exists another CPTP map
$\tilde{\Lambda}$ such that

$\tilde{\Lambda}(\Lambda(\rho))=\rho,\,\forall\rho\in\mathcal{S}(\mathcal{H}).$
(10)

The map $\tilde{\Lambda}$ is known as the recovery map of $\Lambda$ for the
subset $\mathcal{S}(\mathcal{H})$, such that
$\mathcal{S}(\mathcal{H})\subseteq\mathcal{D}(\mathcal{H})$. The most well-
known class of recovery maps are the Petz recovery maps Petz86 ; Petz88 ;
Petz03 , defined as

$\mathcal{P}^{\sigma}(\rho)=\sigma^{\frac{1}{2}}\Lambda^{\dagger}\left(\Lambda(\sigma)^{-\frac{1}{2}}\;\rho\;\Lambda(\sigma)^{-\frac{1}{2}}\right)\sigma^{\frac{1}{2}},$
(11)

with $\Lambda^{\dagger}$ the dual of $\Lambda$ and $\sigma$ the reference
state, such that $\text{supp}(\rho)\subseteq\text{supp}(\sigma)$. In the
process of recovery, all states in $\mathcal{S}(\mathcal{H})$ should be
treated equally. This imposes the introduction of a fiducial reference state
$\sigma$, without focusing on a particular state that one is interested in,
say $\rho$. The reference state, therefore, becomes handy to determine whether
the map $\Lambda$ affects the relative entropy between the states $\sigma$ and
$\rho$. The subset where a certain map is recoverable is defined by the
invariability of the relative entropy after the action of the map, i.e,

$S(\rho||\sigma)=S(\Lambda(\rho)||\Lambda(\sigma)),\,\forall\rho,\sigma\in\mathcal{S}(\mathcal{H}),$
(12)

with $S(\rho||\sigma)=\text{Tr}(\rho\log\rho-\rho\log\sigma)$ being the
relative entropy between two states. If Eq. (12) is satisfied, then there
exists a state $\sigma$ such that $\mathcal{P}^{\sigma}(\Lambda(\rho))=\rho$.

Let us now represent the Petz map in the Choi matrix form. Given a $n\times n$
matrix $X$, $|X\rangle\rangle$ denotes the vectorized version of $M$ where the
matrix is row-wise folded into a column vector.

The Petz recovery map can be expressed as a composition of three operators,
$\mathcal{P}^{\sigma}=A.B.C$, with $A=C=\sigma^{\frac{1}{2}}$ and
$B=\Lambda^{\dagger}\left(\Lambda(\sigma)^{-\frac{1}{2}}\;\rho\;\Lambda(\sigma)^{-\frac{1}{2}}\right)$.
The $\mathcal{A}$ form for the Petz map can be expressed as
$\mathcal{A}_{\text{Petz}}=|A.B.C\rangle\rangle=A\otimes
C^{T}|B\rangle\rangle$. Expanding the operator $B$, we have

$\displaystyle B$
$\displaystyle=\Lambda^{\dagger}\left(\Lambda(\sigma)^{-\frac{1}{2}}\;\rho\;\Lambda(\sigma)^{-\frac{1}{2}}\right)$
$\displaystyle=\sum_{\alpha}K_{\alpha}^{\dagger}\left(\Lambda(\sigma)^{-\frac{1}{2}}\;\rho\;\Lambda(\sigma)^{-\frac{1}{2}}\right)K_{\alpha}$
$\displaystyle=\sum_{\alpha}M_{\alpha}\rho M_{\alpha}^{\dagger},$

with $M_{\alpha}=K_{\alpha}^{\dagger}.\Lambda(\sigma)^{-\frac{1}{2}}$. We
therefore obtain

$\mathcal{A}_{\text{Petz}}=\sigma^{\frac{1}{2}}\otimes(\sigma^{\frac{1}{2}})^{T}.\sum_{\alpha}M_{\alpha}\otimes
M_{\alpha}^{*}.$ (13)

The corresponding Choi matrix can be obtained by reshuffling the matrix
$\mathcal{A}_{\text{Petz}}$.

Figure 1: $(a)$ Set of accessible states under the action of the dephasing
channel $\Lambda_{D}$ (blue ball) and the corresponding Petz recovery map for
$\sigma(0.5)$ (green ball) and $\sigma(0.01)$ (orange ball) with $p=0.3$. The
sets of accessible states for the Petz map (green and orange ball) coincide
for both choices of reference states. $(b)$ Now fixing $p=0.5$. $(c)$ Volume
$\delta V$ of accessible states for the dephasing channel $\Lambda_{D}$ (blue
line) and its corresponding Petz recovery map $\mathcal{P}_{D}^{\sigma(0.01)}$
(orange dashed-dotted line) and $\mathcal{P}_{D}^{\sigma(0.5)}$ (green dashed
line) in terms of the noise strength $p$. $(d)$ Set of accessible states under
the action of the amplitude-damping channel $\Lambda_{AD}$ (blue ball) and the
corresponding Petz recovery map for $\sigma(0.5)$ (green ball) and
$\sigma(0.01)$ (orange ball) with $p=0.3$. The sets of accessible states for
the $\Lambda_{AD}$ and $\mathcal{P}_{AD}^{\sigma(0.01)}$ (green and blue ball)
coincide. $(e)$ Now fixing $p=0.5$. $(f)$ Volume $\delta V$ of accessible
states for the amplitude-damping channel $\Lambda_{AD}$ (blue line) and its
corresponding Petz recovery map $\mathcal{P}_{AD}^{\sigma(0.01)}$ (orange
dashed-dotted line) and $\mathcal{P}_{AD}^{\sigma(0.5)}$ (green dashed line)
in terms of the noise strength $p$.

##  III Geometrical characterization for the Petz recovery maps

Let us now look into certain geometrical features of the Petz map. To this
end, we restrict our analysis to qubits. We investigate how the set of qubit
states vary under the action of a quantum channel and its corresponding Petz
recovery map. The resulting sets of evolved states
$\rho^{\prime}=\\{\Lambda(\rho)|\rho\in\mathcal{D}(\mathcal{H})\\}$ and
$\rho_{\mathcal{P}}=\\{\mathcal{P}^{\sigma}(\Lambda(\rho))|\rho\in\mathcal{D}(\mathcal{H}\\}$
for the above-mentioned scenarios, is defined here by the set of accessible
states.

For qubits, the Bloch vector is 3-dimensional. The state space of all qubit
density operators is represented by a ball of unit radius about the origin of
$\mathbb{R}^{3}$, called the Bloch ball. Each point $(r_{x},r_{y},r_{z})$ of
this ball corresponds to a unique state of the qubit. From Eq. (5), one can
note that the action of a CPTP map $\Lambda$ either rotates or shrinks the
Bloch vector for unital maps along with a possible translation in the case of
non-unital maps. Such transformations induce a change in the volume of the set
of accessible states, quantified as $\delta V=|M|$. We indicate with $|M|$ the
determinant of the matrix $M$. For any positive and trace-preserving map, this
quantity decreases monotonically with the strength of the map Cirac08 .

Let us now study the geometrical aspects of the Petz recovery map for the
dephasing channel on qubits. States lying on the $z$-axis of the Bloch ball
are invariant under the action of the dephasing map. These are the well-known
fixed points of the map, $\Lambda(\rho)=\rho$. The set of accessible states is
$\\{\rho(\vec{r^{\prime}}):\vec{r^{\prime}}=[(1-p)r_{x},(1-p)r_{y},r_{z}]^{T}\\}$.
To obtain the set of accessible states for the Petz recovery map, we consider
the following parametrization for the reference state $\sigma$

$\sigma(\epsilon)=(1-\epsilon)|0\rangle\langle 0|+\epsilon|1\rangle\langle
1|,$ (14)

with $\epsilon\in[0,1]$, a parameter that determines where the reference state
lies on the $z$-axis. Choosing $\sigma$ as in Eq. (14), we build a Petz map
$\mathcal{P}_{D}^{\sigma}$ for the dephasing channel, such that the recovered
state is given by
$\rho_{\mathcal{P}}=\mathcal{P}^{\sigma}_{D}(\Lambda_{D}(\rho))$, for any
initial state $\rho$. The corresponding Bloch vector $\vec{r}_{\mathcal{P}}$
is then written as

$\begin{split}r_{\mathcal{P}_{x}}=&\,(1+p^{2}-2p)r_{x}\\\
r_{\mathcal{P}_{y}}=&\,(1+p^{2}-2p)r_{y}\\\
r_{\mathcal{P}_{z}}=&\,r_{z}.\end{split}$ (15)

As expected, $\vec{r}_{\mathcal{P}}$ is independent of $\epsilon$. This issue
was already discussed in Lea21 ; any reference state of the form of Eq. (14)
is an optimal choice for $\mathcal{P}^{\sigma}_{D}$. As can be apprehended
from Fig. 1 $(a)$ and $(b)$ there is no distinction between the sets of
recovered states of the Petz map for the different reference states (the green
and orange balls in the plot are overlapped). However, the set of accessible
states of the Petz map is always smaller than the set of accessible states of
the dephasing map. In Fig. 1 $(c)$, we observe how the volumes of accessible
states decrease with the noise strength of the map $p$. For intermediate noise
strengths the volume for $\mathcal{P}^{\sigma}_{D}$ decreases faster than for
$\Lambda_{D}$. They coincide in the extreme points when $p=0$ and $p=1$.

It is worth mentioning that the set of accessible states of the Petz map is
obtained by acting the recovery map on a subset $\mathcal{S}(\mathcal{H})$ of
the full set of one-qubit states, Eq. (11). In other words, the Petz map does
not act over the full state space $\mathcal{D}(\mathcal{H})$. The bigger the
set of accessible recovered states is, the better the access the Petz map has
on $\mathcal{D}(\mathcal{H})$. However, it is worth pointing out that the
volume of accessible states is not the only feature we should take into
consideration when talking about state recoverability. In an optimal scenario,
i.e. perfect recoverability, the set of accessible states for the Petz map
should coincide with the set of accessible states of the decoherence channel.

Moving to the non-unital case, under the action of the amplitude damping
channel, the set of accessible states is
$\\{\rho(\vec{r^{\prime}}):\vec{r^{\prime}}=[\sqrt{1-p}\,r_{x},\sqrt{1-p}\,r_{y},(1-p)\,rz+p\,\tau_{z}]^{T}\\}$.
Observe the shift in the $z$-direction, given by the translation component
$\tau_{z}$. Taking the reference state of the form of Eq. (14), we obtain the
corresponding Petz recovered states for the amplitude damping channel
$\mathcal{P}^{\sigma}_{AD}$, whose Bloch vector is

$\begin{split}r_{\mathcal{P}_{x}}=&\frac{\sqrt{1-\epsilon}\sqrt{|1-p|}}{\sqrt{1-\epsilon(1-p)}}\,r_{x},\text{
}r_{\mathcal{P}_{y}}=\frac{\sqrt{1-\epsilon}\sqrt{|1-p|}}{\sqrt{1-\epsilon(1-p)}}\,r_{y}\\\
r_{\mathcal{P}_{z}}=&\frac{1-\epsilon+\epsilon
p-p}{1-\epsilon(1-p)}\,r_{z},\text{
}\tau_{\mathcal{P}_{z}}=\frac{p(1-2\epsilon)}{1-\epsilon(1-p)}.\end{split}$
(16)

The Petz map is now a non-unital map, depending on the choice of $\epsilon$
and $p$. From Eq. (16), we see that the recovered state depends not only on
the strength of the map, given by $p$, but also on the choice of the reference
state, given by $\epsilon$, contrary to what was observed in Eq. (15) for the
dephasing channel. Again we choose two specific reference states:
$\sigma(0.5)$ corresponds to the maximally mixed state and $\sigma(0.01)$ is a
mixed state close to the fixed point of the map $|0\rangle\langle 0|$
footnote_1 . In Fig. 1 $(d)$ and $(e)$, we observe how for $\sigma(0.5)$, the
set of accessible states for $\mathcal{P}^{\sigma}_{AD}$ (orange ball) is much
smaller than the one for $\sigma(0.01)$ (green ball). It is also noticeable
that $\mathcal{P}^{\sigma(0.01)}_{AD}$ is non-unital and its corresponding set
of accessible states is shifted on the $z-$axis by
$\tau_{\mathcal{P}_{z}}=0.296$ and $\tau_{\mathcal{P}_{z}}=0.492$ as in Figs.
1 $(d)$ and $(e)$, respectively. However, independently of the noise strength
of the map, $\mathcal{P}^{\sigma(0.5)}_{AD}$ is always a unital map, from Eq.
(16), we obtain $\tau_{\mathcal{P}_{z}}=0$. The sets of accessible states for
$\Lambda_{AD}$ and $\mathcal{P}^{\sigma(0.01)}_{AD}$ are superposed, meaning
that they coincide for all the points of the ball (the blue and orange balls
coincide). In Fig. 1 $(f)$, we observe how the volume of accessible states for
the Petz map decreases faster by choosing $\sigma$ as the maximally mixed
state. On the contrary, taking $\sigma$ as a state close to the fixed point of
the map, the volume decreases at the same rate as the map itself. Both
reference states show an interesting property. The maximally mixed state is
the trivial choice when dealing with recovery strategies, however, it is not
as effective as the fixed point of the map for intermediate noise strengths
Lea21 . We provid a visual tool for the framework developed in Lea21 and by
observing these results, we conclude that the set of accessible states
$\\{\rho_{\mathcal{P}}\\}$ reflects on the performance of the Petz map. A
larger set of accessible states means a smaller compression of the domain,
which is an advantage for the recovery map. Besides, we showed that the fixed
point of the map is always a good choice for $\sigma$ when dealing with the
full state space as an input. Unfortunately, the volume of accessible states
is not the most appropriate figure of merit when moving to higher dimensions.
Besides the increase in the computational cost, a geometrical visualization of
the state space in higher dimensions turns out to be quite complicated and it
is still an open question Kimura ; Krammer . In the next section, we develop a
different framework employing the Choi matrices to explore if the dimension
plays a role in the recovery strategy via the Petz maps.

##  IV Scalability of the Petz maps

We now attempt to investigate the performance of the Petz recovery maps for
higher dimensions, via a state-independent figure of merit. Ideally, if the
Petz map is capable of recovering the initial state $\rho$, we would obtain
$\mathcal{P}^{\sigma}\Lambda(\rho)=\rho$, which implies

$\mathcal{P}^{\sigma}\Lambda=\mathcal{I},$ (17)

where $\mathcal{I}$ is the identity channel. In order to check the validity of
Eq. (17), we propose a measure of distance between maps, given by

$D(\mathcal{P}^{\sigma}\Lambda,\mathcal{I})=||J(\mathcal{P}^{\sigma}\Lambda)-J(\mathcal{I})||_{1},$
(18)

where $J(\Lambda)$ is the Choi matrix corresponding to $\Lambda$. A map
$\Lambda$ has an inverse $\Lambda^{-1}$ if
$\Lambda^{-1}(\Lambda(\rho))=\rho,\,\forall\rho\in\mathcal{D}(\mathcal{H})$.
The measure above reflects how far the Petz map is from the inverse of the map
$\Lambda^{-1}$. If $\mathcal{P}^{\sigma}=\Lambda^{-1}$,
$D(\mathcal{P}^{\sigma}\Lambda,\mathcal{I})=0$. The inverse map is usually not
completely positive unless the original map is a unitary
jagadish_measure2_2019 .

The measure above is well-behaved and bounded by Wallman

$\frac{1}{d}||\Lambda_{1}-\Lambda_{2}||_{\diamond}\leq||J(\Lambda_{1})-J(\Lambda_{2})||_{1}\leq||\Lambda_{1}-\Lambda_{2}||_{\diamond},$
(19)

with
$||\Lambda||_{\diamond}\equiv\text{max}_{\rho}||(\mathbb{1}\otimes\Lambda)\rho||_{1}$
being the diamond norm. The channels are perfectly distinguishable if
$||\Lambda_{1}-\Lambda_{2}||_{\diamond}=2$. Observe that the measure proposed
in Eq. (18) can be more easily calculated, since it does not require an
optimization as the diamond norm. This strategy provides a state-independent
framework, which can be quite useful and has a lower computational cost when
increasing the dimension. To obtain the Petz recovery map, we extend the
parametrization from Eq. (14) to the $d$-dimensional scenario

$\sigma(\epsilon)=(1-\epsilon)|0\rangle\langle
0|+\frac{\epsilon}{d-1}\sum_{n=1}^{d}|n\rangle\langle n|,$ (20)

with $\epsilon\in[0,1]$.

Figure 2: Distance between the Petz map applied to the dephasing channel and
the identity channel,
$D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})$, versus the noise
strength $p$ for different dimensions.

For the dephasing channel, we fix the reference state as the maximally mixed
one, $\sigma(1-1/d)$. A plot of the distance
$D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})$ versus $p$ for
different dimensions $d$ is shown in Fig. 2 . We observe that the displacement
between the curves decreases with increase in dimensions and is getting much
smaller as we move towards
$D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})=2$. In order to
verify if this bound is saturated for a certain dimension $d=n$, in Fig. 3 we
plot $D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})$ versus the
dimension fixing $p=1$ (maximum noise). As expected and previously discussed,
the distance between the two maps stabilizes when increasing the dimension. As
an example, taking $d=20$ we obtain
$D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})=1.90$, for $d=40$,
$D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})=1.95$ and for
$d=60$, $D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})=1.97$.

Figure 3: Distance between the Petz map applied to the dephasing channel and
the identity channel,
$D(\mathcal{P}_{D}^{\sigma(1-1/d)}\Lambda_{D},\mathcal{I})$, versus the
dimension $d$ for $p=1$ (maximum noise). Figure 4: Distance between the Petz
map applied to the amplitude-damping channel and the identity channel,
$D(\mathcal{P}_{AD}^{\sigma(0.01)}\Lambda_{AD},\mathcal{I})$, versus the noise
strength $p$ for different dimensions. Figure 5: Distance between the Petz map
applied to the amplitude-damping channel and the identity channel,
$D(\mathcal{P}_{AD}^{\sigma(0.01)}\Lambda_{AD},\mathcal{I})$, versus the
dimension $d$ for $p=1$ (maximum noise).

Figure 6: Distance between the Petz map applied to the amplitude-damping
channel and the identity channel,
$D(\mathcal{P}_{AD}^{\sigma}\Lambda_{AD},\mathcal{I})$. On the left, the
difference between the choice of the reference state was exploited for
dimensions $d=2,10,\text{and }40$ in terms of noise parameter $p$. For each
dimension, the reference state analyzed is the corresponding maximally mixed
state, $\sigma(1-1/d)$, and a reference state close to the fixed point of the
map, $\sigma(0.01)$. On the right,
$D(\mathcal{P}_{AD}^{\sigma}\Lambda_{AD},\mathcal{I})$ for $p=0.3$ versus the
dimension $d$.

For the amplitude damping channel, a similar behavior can be observed.
Initially, we decided to fix the reference state close to the fixed point,
$\sigma(0.01)$. In Fig. 4 we observe how the distance
$D(\mathcal{P}_{AD}^{\sigma(0.01)}\Lambda_{AD},\mathcal{I})$ increases with
the dimension. Again, the curves are getting closer for higher dimensions.
This can be easily seen in Fig. 5. Fixing the noise strength at $p=1$, a
plateau is achieved and the performance of the Petz recovery map will be the
same for $d\geq 12$, since
$D^{\prime}(\mathcal{P}_{AD}^{\sigma(0.01)}\Lambda_{AD},\mathcal{I})=0$, where
$D^{\prime}(\mathcal{P}_{AD}^{\sigma(0.01)}\Lambda_{AD},\mathcal{I})$ is the
derivative of the curve displayed in Fig. 5. If we consider now the reference
as being the maximally mixed state $\sigma(1-1/d)$ for intermediate noise
strengths, we observe an advantage for $d=40$ over $d=10$ in Fig. 6 (left).
The green and red solid lines cross at $p=0.75$. To understand if the
dimension plays a role in the performance of the Petz recovery map, in Fig. 6
(right), we plot the distance versus the dimension for different choices of
the reference state. We see that for an intermediate noise strength, $p=0.3$,
the advantage of choosing $\sigma$ close to the fixed point of the map is
clear. However, when choosing $\sigma(1-1/d)$, dimension comes into play and
impacts the recovery performance. We obtain for $d=10$ and $d=40$,
$D(\mathcal{P}_{AD}^{\sigma(0.5)}\Lambda_{AD},\mathcal{I})=0.73$ and
$D(\mathcal{P}_{AD}^{\sigma(0.5)}\Lambda_{AD},\mathcal{I})=0.68$,
respectively.

Figure 7: Non-unitality measure $\mathcal{N}$ in terms of the dimension,
fixing $p=0.3$, for the amplitude-damping channel (thick) and its
corresponding Petz map for $\sigma(0.01)$ (dashed) and $\sigma(1-1/d)$ (dot-
dashed).

To clarify the role of the dimension in the state recovery via the Petz map,
we also compute the non-unitality measure, $\mathcal{N}$, introduced in
Sec.II. As can be seen from Fig. 7, for $p=0.3$, the non-unitality measure
$\mathcal{N}$ for the Petz map increases with the dimension but the outcome is
different depending on the choice of the reference state. We observe that the
curves for $\mathcal{P}_{AD}^{\sigma(0.01)}$ and $\Lambda_{AD}$ are
superposed, meaning that the Petz map for the choice of $\sigma(0.01)$ shifts
the maximally mixed state in the same way as the amplitude-damping channel. On
the other hand, for $\mathcal{P}_{AD}^{\sigma(1-1/d)}$, $\mathcal{N}$ is much
smaller and it does not reproduce the same non-unital behavior as
$\Lambda_{AD}$. This has already been discussed in Sec. III. Changing the
noise strength $p$, the same behavior of $\mathcal{N}$ is observed but with
different values. For a qubit, from Fig. 7, we obtain $\mathcal{N}=0$ and
$\mathcal{N}=0.15$ for the choices of $\sigma(1-1/d)$ and $\sigma(0.01)$,
respectively. Bringing together Figs. 6 and 7, a connection can be traced.
When choosing $\sigma(1-1/d)$, the Petz map does not retrieve the non-unital
behavior of the amplitude damping, contrary to what is observed for the choice
of $\sigma(0.01)$. This is a limitation that reflects on the bad performance
of $\mathcal{P}_{AD}^{\sigma(1-1/d)}$ compared to
$\mathcal{P}_{AD}^{\sigma(0.01)}$. However, when increasing the dimension, the
measure $\mathcal{N}$ also increases, therefore decreasing the difference
between the non-unital character of $\Lambda_{AD}$ and
$\mathcal{P}_{AD}^{\sigma(1-1/d)}$. Consequently, the recovery is improved, as
can be observed in Fig. 6 (right).

##  V Conclusions

In the present work, we analyze the geometrical aspects of the Petz recovery
maps for qubit channels. We have shown how the volume of accessible states can
be directly linked to the recoverability property of the Petz map. The similar
the volumes of accessible states for the evolved states and the recovered
states are, the better the performance of the Petz map. This result is
particularly relevant when considering the choice of reference state, as
discussed also in Lea21 .

In addition, we have developed a state-independent framework for Petz maps for
qudits, which extends our analysis to arbitrary dimensions. While the complete
characterization of the optimal recovery map is a non-trivial problem in
higher dimensions, we used the distance measure based on the Choi-Jamiołkowski
isomorphism. Our findings emphasize the importance of carefully selecting the
reference state for the Petz map. Moreover, we have identified that for non-
unital maps, the non-unitality measure for the Petz map is not the same as the
one for the corresponding quantum channel. However, the Petz map and the
original map tend to be equally non-unital when increasing the dimensions.

Our work provides a novel approach to understanding and characterizing the
Petz recovery maps, that has the potential to improve quantum information
processing protocols. We anticipate that the insights gained about the
characteristics and performance of Petz recovery maps can provide valuable
guidance in designing quantum error correction protocols that are more
effective and applicable to a wide range of quantum systems.

###### Acknowledgements.

V.J. acknowledges financial support from the Foundation for Polish Science
through TEAM-NET project (contract no. POIR.04.04.00-00-17C1/18-00). N.K.B.
acknowledges financial support from CNPq Brazil (Universal Grant No.
406499/2021-7) and FAPESP (Grant 2021/06035-0). N.K.B. is part of the
Brazilian National Institute for Quantum Information (INCT Grant
465469/2014-0). F.P. acknowledges that this research is supported by the South
African Research Chair Initiative of the Department of Science and Innovation
and the National Research Foundation (NRF).

##  Appendix A Non-uniform amplitude damping channel

We also investigate if changes in the recoverability can be observed in the
case of non-uniform amplitude damping, i.e. when different damping rates are
considered. Let us take the qutrit case as an example. The Kraus operators are
known from Eq. (II.1.2), but now with a slight modification;

$\displaystyle K_{0}$ $\displaystyle=$ $\displaystyle|0\rangle\langle
0|+\sqrt{1-p_{1}}|1\rangle\langle 1|+\sqrt{1-p_{2}}|2\rangle\langle 2|$
$\displaystyle K_{1}$ $\displaystyle=$
$\displaystyle\sqrt{p_{1}}|0\rangle\langle 1|$ $\displaystyle K_{2}$
$\displaystyle=$ $\displaystyle\sqrt{p_{2}}|0\rangle\langle 2|,$ (21)

with $p_{1},p_{2}\in[0,1]$ being the decay probabilities. In the non-uniform
case, these probabilities assume different values, which allows us to deal
with a more realistic scenario. In Fig. 8 we plot the behavior of the distance
defined in Eq. (18), in terms of $p_{1}$ and $p_{2}$.

Figure 8: Distance between the Petz map applied to the amplitude-damping and
the identity channel in terms of the decay probabilities $p_{1}$ and $p_{2}$.
The red curve corresponds to the uniform case, $p_{1}=p_{2}$.

We can see that the distance for the uniform amplitude damping (diagonal red
curve) turns out to be smaller than for the non-uniform scenario. When the
decay probabilities are equal, a better recovery is observed for the Petz map.
Assuming that one of the probabilities is zero, the system reduces to a two-
level system, and we recover the result displayed in Fig. 4 of the main text
for $d=2$.

## References

  * (1) John D. Norton. _The Impossible Process: Thermodynamic Reversibility._ Stud. Hist. Philos. Sci. B - Stud. Hist. Philos. Mod. Phys. 55: 43-61 (2016)
  * (2) Heinz-Peter Breuer and Francesco Petruccione.  The Theory of Open Quantum Systems.  (Oxford University Press, Oxford, 2002).
  * (3) Ángel Rivas and Susana F. Huelga.  Open Quantum Systems: An Introduction.  (Springer Berlin, Heidelberg, 2012).
  * (4) Serge Haroche and Jean Michel Raimond.  Exploring the Quantum: Atoms, Cavities, and Photons.  (Oxford University Press, Oxford, 2006).
  * (5) Vinayak Jagadish and Francesco Petruccione. _An invitation to quantum channels._ Quanta 7(1):54–67 (2018).
  * (6) Dénes Petz. _Sufficient subalgebras and the relative entropy of states of a von Neumann algebra._ Commun. Math.Phys. 105, 123-131 (1986).
  * (7) Dénes Petz. _Sufficiency of channels over von Neumann algebras._ Quart. J. Math.Oxford Ser. (2), 39,01 (1988).
  * (8) Dénes Petz. _Monotonicity of quantum relative entropy revisited._ Rev. Math. Phys. 15,01 (2003).
  * (9) Marius Junge, Renato Renner, David Sutter, Mark M. Wilde, and Andreas Winter. _Universal Recovery Maps and Approximate Sufficiency of Quantum Relative Entropy_. Ann. Henri Poincaré 19(10):2955–2978, (2018).
  * (10) Howard Barnum and Emanuel Knill. _Reversing quantum dynamics with near-optimal quantum and classical fidelity_. J. Math. Phys. 43(5):2097–2106, (2002).
  * (11) Francesco Buscemi and Valerio Scarani. _Fluctuation theorems from Bayesian retrodiction._ Phys. Rev. E 103(5): 052111 (2021).
  * (12) Clive Cenxin Aw, Francesco Buscemi, and Valerio Scarani. _Fluctuation theorems with retrodiction rather than reverse processes._ AVS Quantum Sci. 3(4): 045601 (2021).
  * (13) Hyukjoon Kwon, Rick Mukherjee, and M.S. Kim. _Reversing Lindblad Dynamics via Continuous Petz Recovery Map._ Phys. Rev. Lett. 128, 020403 (2022).
  * (14) András Gilyén, Seth Lloyd, Iman Marvian, Yihui Quek, and Mark M. Wilde. _Quantum Algorithm for Petz Recovery Channels and Pretty Good Measurements_ Phys. Rev. Lett. 128, 220502 (2022)
  * (15) Lea Lautenbacher, Fernando de Melo, and Nadja K. Bernardes. _Approximating invertible maps by recovery channels: Optimality and an application to non-Markovian dynamics._ Phys. Rev. A 105, 042421 (2022).
  * (16) Jacopo Surace and Matteo Scandi. _State retrieval beyond Bayes’ retrodiction and reverse processes._ Quantum 7, 990 (2023).
  * (17) Karol Życzkowski and Hans-Jürgen Sommers. _Average fidelity between random quantum states._ Phys. Rev. A 71, 032313 (2005).
  * (18) Man-Duen Choi. _Completely positive linear maps on complex matrices._ Linear Algebra Appl. 10, 285 (1975).
  * (19) Andrzej Jamiołkowski. _Linear transformations which preserve trace and positive semidefiniteness of operators._ Rep. Math. Phys. 3 (4): 275–278 (1972).
  * (20) Michael M. Wolf and Ignacio J. Cirac. _Dividing Quantum Channels._ Commun. Math. Phys. 279(1):147–168 (2008).
  * (21) The choice $\epsilon=0.01$ is based on the fact that the results converge for $\epsilon<0.01$.
  * (22) Gen Kimura. _The Bloch Vector for N-Level Systems._ Phys. Lett. A 314, 339 (2003).
  * (23) Reinhold A Bertlmann and Philipp Krammer. _Bloch vectors for qudits._ J. Phys. A: Math. Theor. 41 235303 (2008).
  * (24) Vinayak Jagadish, R. Srikanth, and Francesco Petruccione. _Measure of not-completely-positive qubit maps: The general case._ Phys. Rev. A 100, 012336 (2019).
  * (25) Joel J. Wallman and Steven T. Flammia. _Randomized benchmarking with confidence._ New J. Phys. 16, 103032 (2014).


#  F-coref: Fast, Accurate and Easy to Use Coreference Resolution

Shon Otmazgin1 Arie Cattan1 Yoav Goldberg1,2
1Computer Science Department, Bar Ilan University
2Allen Institute for Artificial Intelligence
<EMAIL_ADDRESS>

###### Abstract

We introduce fastcoref, a python package for fast, accurate, and easy-to-use
English coreference resolution. The package is pip-installable, and allows two
modes: an accurate mode based on the LingMess architecture, providing state-
of-the-art coreference accuracy, and a substantially faster model, F-coref,
which is the focus of this work. F-coref allows to process 2.8K OntoNotes
documents in 25 seconds on a V100 GPU (compared to 6 minutes for the LingMess
model, and to 12 minutes of the popular AllenNLP coreference model) with only
a modest drop in accuracy. The fast speed is achieved through a combination of
distillation of a compact model from the LingMess model, and an efficient
batching implementation using a technique we call leftover
batching.111https://github.com/shon-otmazgin/fastcoref

##  1 Introduction

Coreference Resolution consists of identifying textual mentions that refer to
the same entity in a given text (Karttunen, 1969). This fundamental NLP task
can benefit various applications such as Information Extraction (Luan et al.,
2018; Li et al., 2020; Jain et al., 2020), Question Answering (Dasigi et al.,
2019; Chen and Durrett, 2021), Machine Translation (Stojanovski and Fraser,
2018; Voita et al., 2018), and Summarization (Christensen et al., 2013; Falke
et al., 2017; Pasunuru et al., 2021). However, compared to other core tasks
such as POS Tagging, named-entity recognition or syntactic parsing, existing
packages and state-of-the-art models for coreference resolution are
challenging to apply: there are few easy-to-use packages implementing state-
of-the-art models, and the available packages consume a lot of GPU memory, and
take very long to process each document. For example, the coreference model in
the popular AllenNLP package Gardner et al. (2017), implementing the model of
Joshi et al. (2020), requires 27GB of GPU memory and takes 12 minutes to
process the 2.8K documents of the OntoNotes corpus, on a V100 GPU.

In this work, we introduce F-coref, a new open source Python package for
simply running an efficient coreference model using a few lines of code.
F-coref predicts coreference clusters 29 times faster than the AllenNLP model
(processing the OntoNotes corpus in 25 seconds) and requires only 15% of its
GPU memory use, with only a small drop in performance (78.5 vs 79.6 average
F1). The package also includes LingMess Otmazgin et al. (2022), a state-of-
the-art coreference model, which is almost twice as fast as the AllenNLP
model, while being more accurate (81.4 average F1), under the same API.

To achieve F-coref’s speed, we use two additive techniques: model distillation
of the strong-but-slow LingMess model using large unlabeled data, and an
effective batching technique that reduces the number of padded tokens in a
batch.

##  2 The F-coref API

The fastcoref Python package is pip installable (pip install fastcoref) and
provides an easy and fast API for coreference information with only few lines
of code without any prepossessing steps.

The F-coref constructor initializes our pretrained model on a single device:

⬇

from fastcoref import FCoref

model = FCoref(device=’cuda:0’)

The main functionally of the package is the predict function, which accepts a
list of texts.

⬇

preds = model.predict(

texts=[’We are so happy to see you

using our coref package.

This package is very fast!’]

)

The return value of the predict function is a list of CorefResult objects,
from which one can extract the coreference clusters (either as strings or as
character indices over the original texts), as well as the logits for each
corefering entity pair:

⬇

preds[0].get_clusters(as_strings=False)

> [[(0, 2), (33, 36)],

[(33, 50), (52, 64)]

]

preds[0].get_clusters()

> [[’We’, ’our’],

[’our coref package’, ’This package’]

]

preds[0].get_logit(

span_i=(33, 50), span_j=(52, 64)

)

> 18.852894

Processing can be applied to a collection of texts of any length in a batched
and parallel fashion:

⬇

texts = [’text 1’, ’text 2’,.., ’text n’]

# control the batch size

# with max_tokens_in_batch parameter

preds = model.predict(

texts=texts, max_tokens_in_batch=100

)

The max_tokens_in_batch parameter can be used to control the speed vs. memory
consumption tradeoff, and can be tuned to maximize the utilization of the
associated hardware.

To control speed vs. accuracy tradeoff, use the larger but more accurate
LingMess model, simply import LingMessCoref instead of FCoref:

⬇

from fastcoref import LingMessCoref

model = LingMessCoref(device=’cuda:0’)

On top of the provided models, the package also provides a custom SpaCy
component that can be plugged into a SpaCy(V3) pipeline (Honnibal et al.,
2020). Additionally, the package includes a CorefTrainer for training and
distilling coreference model on your own data, opening the possibility for
fast and accurate coreference models for additional languages and domains.

To summarize, the package provides a simple API that makes predicting
coreference entities straightforward and easy-to-use. The package supports any
text length as input, and performs efficient batching. The package’s F-coref
model is 29 times faster and 4 times smaller than the popular coreference
model in the AllenNLP package, while the provided LingMess mode is twice as
fast the AllenNLP implementation, and more accurate.

##  3 Background: Neural Coreference

Lee et al. (2017) present the first end-to-end model that jointly learns
mention detection and coreference decision. Successive follow-up works kept
improving performance through the incorporation of widely popular pretrained
architectures (Lee et al., 2018; Joshi et al., 2019; Kantor and Globerson,
2019; Joshi et al., 2020). However, as the dimensionality of contextualized
encoders increases, keeping in memory all possible span representations
becomes highly costly and computationally untractable for long documents.

###  3.1 Faster Neural Coreference

Several methods have been proposed to address this memory constraint at the
cost of the computation time and a slight performance deterioration (Xia et
al., 2020; Toshniwal et al., 2020; Thirukovalluru et al., 2021). The _s2e_
model of Kirstain et al. (2021) managed to improve computation time with a
slight _increase_ in accuracy.

#### s2e

Our F-coref is based on the architecture of the _s2e_ model by Kirstain et al.
(2021). Like other neural coreference models, _s2e_ scores each pair of spans
in the text to be co-referring to each other. However, in order to achieve
lower memory footprint s2e moves to representing each span as a function of
its start and end tokens. Consequently, the model avoids holding vector
representation for each of the $O(n^{2})$ spans in memory, and instead stores
only $O(n)$ vectors. This reduced memory footprint allows it to handle longer
sequences.

The _s2e_ architecture includes three components: (1) Longformer Beltagy et
al. (2020), a contextualized encoder; (2) a parameterized _mention scoring
function_ $f_{m}$; and (3) a parameterized _pairwise antecedent scoring
function_ $f_{a}$. To score any pair of spans to be co-referring, the model
starts by encoding the text using Longformer into vectors $x_{1},...,x_{n}$.
Using these vectors, for each possible span $q=(x_{k},x_{\ell})$ the mention
scoring function $f_{m}(q)$, scores how likely $q$ (“query”) being a mention.
Then for a pair of spans $c=(x_{i},x_{j})$, $q=(x_{k},x_{\ell})$ where $c$
(“candidate”) appears before $q$, the pairwise antecedent scoring function,
$f_{a}(c,q)$, scores how likely is $c$ being an antecedent of $q$. In
practice, to avoid complexity of $\mathcal{O}(n^{4})$, the antecedent function
scores only the $\lambda T$ spans with highest mention scores (where $T$ is
the number of tokens). Finally, the final pairwise score for a coreference
link between $c$ and $q$ is composed by the score of $q$ being a mention, $c$
being a mention, and how likely is $c$ being an antecedent of $q$:

$\displaystyle
F(c,q)=\begin{cases}f_{m}(c)+f_{m}(q)+f_{a}(c,q)&c\neq\varepsilon\\\
0&c=\varepsilon\end{cases}$

where $\varepsilon$ is the null antecedent.

The computation of $f_{m}$ and $f_{q}$ for the entire sequence can be
efficiently batched.

#### Word-level coreference

Dobrovolskii (2021) proposed moving from scoring pairs of spans to scoring
pairs of words, establishing coreference relations between the words, and then
expanding each of the relevant words into their mention boundaries. This
reduces the model complexity from $\mathcal{O}(n^{4})$ to
$\mathcal{O}(n^{2})$.

###  3.2 What remains slow?

While the _s2e_ and the word-level models are considered lightweight and
efficient, and substantally improve in speed over Joshi et al. (2019), their
computation time is still dominated by their expensive contextualized encoding
stage. They also use relatively large hidden layers in their scoring functions
(the s2e model has 26 layers and 494M parameters). Thus, one avenue for
improving coreference speed is by reducing the model size. Additionally, while
batching computations can improve parallelism and thus also throughput, the
implementation of batching long documents of varying lengths is often sub-
optimal, and results in many padded tokens which translate to wasted
computation.

###  3.3 Accurate Neural Coreference

Our recent LingMess model (Otmazgin et al., 2022) improves coreference
accuracy by observing that different types of entities require different
strategies to score, and replacing the single mention-pair scorer with a set
of specialized scorers. During inference, each mention pair is
deterministically routed to one of the scorers, based on the the type of
mentions being scored. This results in state-of-the-art coreference accuracy,
while being somewhat less efficient to run and to batch than the _s2e_ model.

##  4 Method

We employ two complementary directions in order to obtain a fast and efficient
coreference model. First, we substantially reduce the size of the _s2e_ model
using knowledge distillation (§4.1) from the LingMess model. Second, our
implementation aims to maximize parallelism via batching while limiting the
number of unnecessary computations such as padded tokens (§4.2).

###  4.1 Knowledge Distillation

Knowledge Distillation is the process of learning a small student model from a
large teacher model.

#### Teacher model

We use the state-of-the-art LingMess model of Otmazgin et al. (2022) as the
teacher model.

#### Student model

we build our student model as a variant of the _s2e_ model with fewer layers
and parameters. The “expensive” Longformer Beltagy et al. (2020) encoder was
replaced with DistilRoBERTa Sanh et al. (2019), which is on average $\times 8$
faster than Longformer. The number of parameters of the mention and the
antecedent pairwise scorers was reduced by a factor of 6. This reduces the
total number of parameters from 494M to 91M. In addition, the number of
sequential layers in the network reduced from 26 layers to only 8 layers (6
encoder layers, 1 mention scorer and 1 antecedent scorer). As a result, our
student combines the strengths of the _s2e_ model by not constructing span
representation with a lightweight encoder and substantially less model
parameters.

#### Hard distillation

Traditional approaches for knowledge distillation trains the student on the
logits of the teacher model’s predictions on unlabeled data Gou et al. (2021).
However, as we will further elaborate in Section §5.1, applying such an
approach to a coreference model with all its components (i.e. encoder, mention
scorer, pruning, antecedent scorer) achieves poor performance. To remedy this
issue, we employ hard target knowledge distillation, where the teacher model
acts as an annotator for the unlabeled data and the student model learns from
these “silver” annotations.

Figure 1: Illustration of document batching using vanilla batching (left) and
our leftovers approach (right). In our leftover batching, we create two
separate batches, one for the full segments without padding (orange) and one
for the leftover segments (green), thus substantially reducing the number of
padded tokens (red) compared to the vanilla approach.

###  4.2 Maximizing Parallelism and Reducing Unnecessary Computations

#### Mention pruning

As mentioned in Section 3, the coreference model computes antecedent scores
only for the $\lambda T$ spans with highest mention scores, where $T$ is the
number of tokens. As the number of coreferring spans cannot be known in
advance, the common approach is to use a soft pruning coefficient
($\lambda=0.4$) to guarantee high mention recall, while expecting the
antecedent scorer to assign a negative score to pairs involving a wrong
mention span (Lee et al., 2017; Kirstain et al., 2021). In F-coref, we adopt a
more aggressive pruning ($\lambda=0.25$), which decreases the number of
pairwise comparisons by a factor of 2.56 without harming performance.

#### Dynamic batching

We adopt a dynamic batching approach which, given a large number of documents,
batches documents until we reach a certain maximum number of tokens. Compared
with the naive approach of batching a fixed number of documents together,
dynamic batching enables to fully exploit the available memory in our hardware
(this approach was also used by Kirstain et al. (2021) when _training_ the
_s2e_ model, but not for inference).

#### Leftovers batching

Figure 1 illustrates our document batching strategy, in comparison with the
common approach.

As mentioned in Section 3, the first step in our coreference model consists of
encoding the document using a transformer-based encoder. The common approach
for encoding long documents with transformer encoders is to split the document
into non-overlapping segments of max_length, where each segment is encoded
separately Joshi et al. (2019); Xu and Choi (2020). With that approach, for
each long document, we obtain two222A document with fewer tokens than
max_length has only one segment. types of segment lengths: (1) one or more
segments of max_length i.e. full tokens segment, (2) exact one segment $\leq$
max_length, i.e the leftovers tokens segment.

Then to batch multiple documents, a naive but popular approach consists of
padding each document’s leftovers segment to max_length. This results in a
high number of padded tokens, e.g., 34.7% of all tokens are padded when
batching 2802 OntoNotes Pradhan et al. (2012) training set documents with
max_length $=512$. Padded tokens result in unnecessary computations in each
layer of the network, as well as unnecessary memory allocation.

To avoid such unnecessary computations, we split each batch into two batches
such that the first batch is for the full segments and the second batch is for
the leftovers tokens segment of each document. Then we pad the second batch to
the max leftovers length rather than padding the leftovers segments to
max_length. Finally we run the two batches separately and combine them
afterwards. With this technique, the padded tokens in the OntoNotes training
set reduced dramatically to 0.6%. It should be noted that the aforementioned
batching technique is not specific for coreference resolution and can be
applied for other tasks that require processing long documents.

|  | MUC |  | B3 |  | CEAF$\phi_{4}$ |  |
---|---|---|---|---|---|---|---|---
|  | P | R | F1 |  | P | R | F1 |  | P | R | F1 |  | Avg. F1
Joshi et al. (2020) |  | 85.8 | 84.8 | 85.3 |  | 78.3 | 77.9 | 78.1 |  | 76.4 | 74.2 | 75.3 |  | 79.6
Kirstain et al. (2021) |  | 86.5 | 85.1 | 85.8 |  | 80.3 | 77.9 | 79.1 |  | 76.8 | 75.4 | 76.1 |  | 80.3
Dobrovolskii (2021) |  | 84.9 | 87.9 | 86.3 |  | 77.4 | 82.6 | 79.9 |  | 76.1 | 77.1 | 76.6 |  | 81.0
Otmazgin et al. (2022) (Teacher) |  | 88.1 | 85.1 | 86.6 |  | 82.7 | 78.3 | 80.5 |  | 78.5 | 76.0 | 77.3 |  | 81.4
F-coref OntoNotes only |  | 78.5 | 84.3 | 81.3 |  | 68.2 | 74.8 | 71.4 |  | 64.1 | 72.9 | 68.2 |  | 73.7
F-coref Multi-News |  | 84.8 | 82.8 | 83.4 |  | 76.8 | 73.7 | 75.2 |  | 73.8 | 72.7 | 73.2 |  | 77.4
\+ FT OntoNotes |  | 85.0 | 83.9 | 84.4 |  | 77.6 | 75.5 | 76.6 |  | 74.7 | 74.3 | 74.5 |  | 78.5

Table 1: Performance on the test set of the English OntoNotes 5.0 dataset. The
averaged F1 of MUC, B3, CEAF$\phi$ is the main evaluation metric.

| Masc | Fem | Bias | Overall
---|---|---|---|---
Otmazgin et al. (2022) (Teacher) | 91.3 | 87.8 | 0.96 | 89.6
F-coref | 87.8 | 83.5 | 0.95 | 85.7

Table 2: Performance on the test set of the GAP coreference dataset. The
reported metrics are F1 scores.

[b]  Runtime Memory Joshi et al. (2020)1 12:06 27.4 Otmazgin et al. (2022)
(Teacher) 06:43 4.6 \+ Batching2 06:00 6.6 Kirstain et al. (2021) 04:37 4.4
Dobrovolskii (2021) 03:49 3.5 F-coref 00:45 3.3 \+ Batching2 00:35 4.5 \+
Leftovers batching2 00:25 4.0

  * 1

AllenNLP package implementation.

  * 2

10K tokens in a single batch.

Table 3: The inference time(Min:Sec) and memory(GiB) for each model on 2.8K
documents. Average of 3 runs. Hardware, NVIDIA Tesla V100 SXM2.

| #Docs | #Chains | #Mentions
---|---|---|---
OntoNotes | 2.8K | 35K | 155K
Multi-News | 123K | 2M | 9M

Table 4: Coreference statistics of the training set of OntoNotes and Multi-
News.

##  5 Experiments and Results

#### Experiments setup

In our experiments, we use the Multi-news Fabbri et al. (2019) dataset to
train our teacher-student architecture. The Multi-news dataset is an open
source dataset aimed at NLP summarization. Each entry in the dataset contains
multiple documents and a summary of these documents. For our purposes, we
ignored the summaries, and train our student model on the documents, a total
of 123,227 documents in the news domain. We chose Multi-News because it
contains a large number of documents in the news genre, which would result in
a large number of coreference clusters (see Table 4 for statistics).
Furthermore, we use the English portion of the OntoNotes Pradhan et al. (2012)
dataset to evaluate the student model performance (on the test set) and to
further fine-tune the student model (on the train set).

The student training procedure includes three phases. In the first phase, we
predict coreference clusters on MultiNews using the teacher model (Otmazgin et
al., 2022). Secondly, following (Wu et al., 2020; Dobrovolskii, 2021), we pre-
train the mention scorer of the student on the output mentions of the teacher,
then train the full student model on the predicted teacher coreference
clusters. Finally, we finetune the student model on the OntoNotes training
set.

#### Accuracy

Table 1 shows F-coref’s performance on the OntoNotes test set according to the
standard evaluation metrics for coreference resolution. F-coref achieved 78.5
F1 with knowledge distillation and finetuning on OntoNotes. When trained only
on OntoNotes, F-coref achieved only 73.7 F1 ($-$4.8 F1), showing a substantial
benefit from knowledge distillation on the Multi-news dataset Fabbri et al.
(2019). In comparison with other coreference models, F-coref degrades by 1.1
point compared to the Joshi et al. (2019) model in the AllenNLP package, by
2.9 F1 points compared to LingMess (Otmazgin et al., 2022), the teacher, and
by 1.8 F1 points versus Kirstain et al. (2021), the _s2e_ model, which F-coref
is a variant of. Table 2 shows similar trends for the GAP dataset (Webster et
al., 2018). This degradation comes in favor to the model efficiency, which we
will discuss in the next paragraph.

#### Speed and Memory Usage

Table 3 summarizes the efficiency of the different coreference models. Using
the techniques described in Section §4 which reduces the model size, maximize
batching and avoids unnecessary computation, the inference time on 2.8K
documents is significantly reduced by factor of 9 from the 03:49 minutes of
Dobrovolskii (2021)—the fastest model to date—to only 25 seconds for
F-coref.333These experiments used batch sizes of 10k tokens. Increasing the
batch sizes increase memory consumption, but does not improve overall speed on
our NVIDIA Tesla V100 hardware.

Most of the speed increase (80%) is due to the smaller model size achieved
through distillation, which alone reduces the runtime to 45 seconds. However,
introducing batching further reduces the runtime by additional 22% to 35
seconds, and our novel leftover batching reduces further 29%, and gets us to
25 seconds. The more aggressive mention pruning (not shown in the table) had
only a negligible additive effect in our experiments, reducing the runtime
from 25 to 24 seconds.

Without batching, F-coref also consumes less memory then Dobrovolskii (2021),
a model that recently reduces the coreference complexity from
$\mathcal{O}(n^{4})$ to $\mathcal{O}(n^{2})$.

Finally, compared to one of the most widely used coreference models, the Joshi
et al. (2020) model available through the AllenNLP package Gardner et al.
(2017), F-coref is 29 times faster and consumes 85% less memory.

###  5.1 Further Analysis

Figure 2: The knowledge distillation learning curve. The x-axis is the number
of documents we took from the Multi-News dataset to train the student model.
The y-axis is the student F1 score for each size.

#### Effect of Unlabeled Data Size

We first analyze the effect of the amount of unlabeled data used in
distillation, by training the student model on different amounts of training
data. As Figure 2 shows the performance gain between 25K documents and 50K
documents is 1 F1 point while the gain between 100K and 125K documents
decreases to 0.4 points. This indicates that the gain margin is decreasing,
but overall, increasing the dataset size continuously improves the model
performance and there is a room for improvement with more data. Fine-tuning
the distilled model on the in-domain OntoNotes data consistently improves the
results, but is also additive with the distillation: the fine-tuned
performance also increases with more unlabaled data in distillation.

#### Effect of the Teacher Model

To estimate the effect of the teacher model, we compare the LingMess teacher
to a _s2e_ model (Kirstain et al., 2021) teacher on the same unlabeled data.
We obtain 76.6 F1 with _s2e_ (vs. 77.4 F1 with LingMess) on the OntoNotes
(test set) when training only on Multi-News and 78.3 F1 with _s2e_ (vs. 78.5
F1 with LingMess) after further fine-tuning on OntoNotes (training set). This
indicates that the student accuracy increases when more accurate models are
used in knowledge distillation, even with hard labels.

#### Soft VS. Hard Distillation

Our first attempt to transfer the coreference knowledge from the teacher model
to the student model was the traditional knowledge distillation, i.e. soft
targets knowledge distillation. For each example in the training set, we
forward it first in the teacher model and obtained the top-scoring spans
indices, and the pairwise coreference logits. Then we forward the example in
the student network (at pruning stage we use the teacher’s top-scoring spans
indices) to obtain the student coreference pairwise logits. Following common
training objective (Hinton et al., 2015; Sanh et al., 2019; Jiao et al.,
2020), we optimize the student model using the soft cross entropy loss between
the student and the teacher logits.

Our student model reached only 64 F1, while achieving 73.7 F1 without
knowledge distillation. Soft distillation presents challenges in coreference
models. The main challenge we encounter is at the pruning stage, where both
teacher and student should prune the exact same mentions from their individual
mention scorer. This forces the student model to learn from a conditional
antecedent distribution of the teacher spans indices instead of the full
antecedent distribution.

Additionally, we observe that learning to mimic the teacher logits may trouble
the training because logits can violate transitivity (e.g positive score for
the mention pairs $(a,b)$ and $(b,c)$ but a negative score for $(a,c)$) and
propagate contradictory information. Specifically, we build the coreference
clusters based on the positive pairwise scores, and verify whether all
pairwise scores within the same coreference clusters are positive, as we would
naturally expect. In fact, 53.8% of the pairwise scores within the same
coreference cluster are negative. In contrast, this undesired behavior does
not happen in hard distillation because we assign positive labels for all
coreferring antecedents for each mention.

##  6 Conclusions

We introduce the _fastcoref_ python package for coreference resolution, and
hope its speed an ease of use will facilitate work that utilizes coreference
resolution at scale.

## Acknowledgements

This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme,
grant agreement No. 802774 (iEXTRACT). Arie Cattan is partially supported by
the Data Science Institute at Bar-Ilan University.

## References

  * Beltagy et al. (2020) Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020.  Longformer: The long-document transformer.  _ArXiv_ , abs/2004.05150.
  * Chen and Durrett (2021) Jifan Chen and Greg Durrett. 2021.  Robust question answering through sub-part alignment.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 1251–1263, Online. Association for Computational Linguistics.
  * Christensen et al. (2013) Janara Christensen, Mausam, Stephen Soderland, and Oren Etzioni. 2013.  Towards coherent multi-document summarization.  In _Proceedings of the 2013 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 1163–1173, Atlanta, Georgia. Association for Computational Linguistics.
  * Dasigi et al. (2019) Pradeep Dasigi, Nelson F. Liu, Ana Marasović, Noah A. Smith, and Matt Gardner. 2019.  Quoref: A reading comprehension dataset with questions requiring coreferential reasoning.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 5925–5932, Hong Kong, China. Association for Computational Linguistics.
  * Dobrovolskii (2021) Vladimir Dobrovolskii. 2021.  Word-level coreference resolution.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 7670–7675, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
  * Fabbri et al. (2019) Alexander Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev. 2019.  Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 1074–1084, Florence, Italy. Association for Computational Linguistics.
  * Falke et al. (2017) Tobias Falke, Christian M. Meyer, and Iryna Gurevych. 2017.  Concept-map-based multi-document summarization using concept coreference resolution and global importance optimization.  In _Proceedings of the Eighth International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 801–811, Taipei, Taiwan. Asian Federation of Natural Language Processing.
  * Gardner et al. (2017) Matt Gardner, Joel Grus, Mark Neumann, Oyvind Tafjord, Pradeep Dasigi, Nelson F. Liu, Matthew Peters, Michael Schmitz, and Luke S. Zettlemoyer. 2017\.  Allennlp: A deep semantic natural language processing platform.
  * Gou et al. (2021) Jianping Gou, B. Yu, Stephen J. Maybank, and Dacheng Tao. 2021.  Knowledge distillation: A survey.  _ArXiv_ , abs/2006.05525.
  * Hinton et al. (2015) Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015.  Distilling the knowledge in a neural network.  _ArXiv_ , abs/1503.02531.
  * Honnibal et al. (2020) Matthew Honnibal, Ines Montani, Sofie Van Landeghem, and Adriane Boyn. 2020.  spaCy: Industrial-strength Natural Language Processing in Python.
  * Jain et al. (2020) Sarthak Jain, Madeleine van Zuylen, Hannaneh Hajishirzi, and Iz Beltagy. 2020.  SciREX: A challenge dataset for document-level information extraction.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 7506–7516, Online. Association for Computational Linguistics.
  * Jiao et al. (2020) Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020.  TinyBERT: Distilling BERT for natural language understanding.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pages 4163–4174, Online. Association for Computational Linguistics.
  * Joshi et al. (2020) Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel S. Weld, Luke Zettlemoyer, and Omer Levy. 2020.  SpanBERT: Improving pre-training by representing and predicting spans.  _Transactions of the Association for Computational Linguistics_ , 8:64–77.
  * Joshi et al. (2019) Mandar Joshi, Omer Levy, Luke Zettlemoyer, and Daniel Weld. 2019.  BERT for coreference resolution: Baselines and analysis.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pages 5803–5808, Hong Kong, China. Association for Computational Linguistics.
  * Kantor and Globerson (2019) Ben Kantor and Amir Globerson. 2019.  Coreference resolution with entity equalization.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 673–677, Florence, Italy. Association for Computational Linguistics.
  * Karttunen (1969) Lauri Karttunen. 1969.  Discourse referents.  In _International Conference on Computational Linguistics COLING 1969: Preprint No. 70_ , Sånga Säby, Sweden.
  * Kirstain et al. (2021) Yuval Kirstain, Ori Ram, and Omer Levy. 2021.  Coreference resolution without span representations.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 2: Short Papers)_ , pages 14–19, Online. Association for Computational Linguistics.
  * Lee et al. (2017) Kenton Lee, Luheng He, Mike Lewis, and Luke Zettlemoyer. 2017.  End-to-end neural coreference resolution.  In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing_ , pages 188–197, Copenhagen, Denmark. Association for Computational Linguistics.
  * Lee et al. (2018) Kenton Lee, Luheng He, and Luke Zettlemoyer. 2018.  Higher-order coreference resolution with coarse-to-fine inference.  In _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers)_ , pages 687–692, New Orleans, Louisiana. Association for Computational Linguistics.
  * Li et al. (2020) Manling Li, Alireza Zareian, Ying Lin, Xiaoman Pan, Spencer Whitehead, Brian Chen, Bo Wu, Heng Ji, Shih-Fu Chang, Clare Voss, Daniel Napierski, and Marjorie Freedman. 2020.  GAIA: A fine-grained multimedia knowledge extraction system.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics: System Demonstrations_ , pages 77–86, Online. Association for Computational Linguistics.
  * Luan et al. (2018) Yi Luan, Luheng He, Mari Ostendorf, and Hannaneh Hajishirzi. 2018.  Multi-task identification of entities, relations, and coreference for scientific knowledge graph construction.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 3219–3232, Brussels, Belgium. Association for Computational Linguistics.
  * Otmazgin et al. (2022) Shon Otmazgin, Arie Cattan, and Yoav Goldberg. 2022.  Lingmess: Linguistically informed multi expert scorers for coreference resolution.  _ArXiv_ , abs/2205.12644.
  * Pasunuru et al. (2021) Ramakanth Pasunuru, Mengwen Liu, Mohit Bansal, Sujith Ravi, and Markus Dreyer. 2021\.  Efficiently summarizing text and graph encodings of multi-document clusters.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 4768–4779, Online. Association for Computational Linguistics.
  * Pradhan et al. (2012) Sameer Pradhan, Alessandro Moschitti, Nianwen Xue, Olga Uryupina, and Yuchen Zhang. 2012.  CoNLL-2012 shared task: Modeling multilingual unrestricted coreference in OntoNotes.  In _Joint Conference on EMNLP and CoNLL - Shared Task_ , pages 1–40, Jeju Island, Korea. Association for Computational Linguistics.
  * Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  _ArXiv_ , abs/1910.01108.
  * Stojanovski and Fraser (2018) Dario Stojanovski and Alexander Fraser. 2018.  Coreference and coherence in neural machine translation: A study using oracle experiments.  In _Proceedings of the Third Conference on Machine Translation: Research Papers_ , pages 49–60, Brussels, Belgium. Association for Computational Linguistics.
  * Thirukovalluru et al. (2021) Raghuveer Thirukovalluru, Nicholas Monath, Kumar Shridhar, Manzil Zaheer, Mrinmaya Sachan, and Andrew McCallum. 2021.  Scaling within document coreference to long texts.  In _Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021_ , pages 3921–3931, Online. Association for Computational Linguistics.
  * Toshniwal et al. (2020) Shubham Toshniwal, Sam Wiseman, Allyson Ettinger, Karen Livescu, and Kevin Gimpel. 2020.  Learning to Ignore: Long Document Coreference with Bounded Memory Neural Networks.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 8519–8526, Online. Association for Computational Linguistics.
  * Voita et al. (2018) Elena Voita, Pavel Serdyukov, Rico Sennrich, and Ivan Titov. 2018.  Context-aware neural machine translation learns anaphora resolution.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1264–1274, Melbourne, Australia. Association for Computational Linguistics.
  * Webster et al. (2018) Kellie Webster, Marta Recasens, Vera Axelrod, and Jason Baldridge. 2018.  Mind the GAP: A balanced corpus of gendered ambiguous pronouns.  _Transactions of the Association for Computational Linguistics_ , 6:605–617.
  * Wu et al. (2020) Wei Wu, Fei Wang, Arianna Yuan, Fei Wu, and Jiwei Li. 2020.  CorefQA: Coreference resolution as query-based span prediction.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 6953–6963, Online. Association for Computational Linguistics.
  * Xia et al. (2020) Patrick Xia, João Sedoc, and Benjamin Van Durme. 2020.  Incremental neural coreference resolution in constant memory.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 8617–8624, Online. Association for Computational Linguistics.
  * Xu and Choi (2020) Liyan Xu and Jinho D. Choi. 2020.  Revealing the myth of higher-order inference in coreference resolution.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 8527–8533, Online. Association for Computational Linguistics.