# SwdFold:A Reweighting and Unfolding method based on Optimal Transport Theory

Chu-Cheng Pan  School of physics and technology, Wuhan University, Wuhan,
Hubei, China  Xiang Dong  School of physics and technology, Wuhan University,
Wuhan, Hubei, China  Yu-Chang Sun  School of physics and technology, Wuhan
University, Wuhan, Hubei, China  Ao-Yan Cheng  School of physics and
technology, Wuhan University, Wuhan, Hubei, China  Ao-Bo Wang  School of
physics and technology, Wuhan University, Wuhan, Hubei, China  Yu-Xuan Hu
School of physics and technology, Wuhan University, Wuhan, Hubei, China  Hao
Cai<EMAIL_ADDRESS>School of physics and technology, Wuhan University, Wuhan,
Hubei, China

###### Abstract

High-energy physics experiments rely heavily on precise measurements of energy
and momentum, yet face significant challenges due to detector limitations,
calibration errors, and the intrinsic nature of particle interactions.
Traditional unfolding techniques have been employed to correct for these
distortions, yet they often suffer from model dependency and stability issues.
We present a novel method, SwdFold, which utilizes the principles of optimal
transport to provide a robust, model-independent framework to estimate the
probability density ratio for data unfolding. It not only unfold the toy
experimental event by reweighted simulated data distributions closely with
true distributions but also maintains the integrity of physical features
across various observables. We can expect it can enable more reliable
predictions and comprehensive analyses as a high precision reweighting and
unfolding tool in high-energy physics.

In high-energy physics experiments, achieving precision and accuracy in
measurements is critical for conducting meticulous analyses. However,
detectors face inherent resolution limitations, calibration errors, and issues
such as the generation of secondary particles, which can skew the measurement
of energy and momentum. Unfolding is a mathematical process used to correct
distortions in data caused by instrumental effects(detector level data),
enabling the recovery of the true distribution of physical events(particle
level data). This task is complex due to the diversity and intricacy of the
physical event distributions, compounded by systematic errors and
uncertainties from the detectors. Unfolding technique is particularly crucial
in accurately determining physical observables like invariant mass
distributions, where it aids in precisely measuring mass values and widths[1].
Moreover, unfolding can handle complex physical phenomena such as jet
measurements, efficiently recovering original distributions from data
convoluted by instrumental effects[2]. Additionally, by mitigating
dependencies on specific experimental setups, unfolding enables consistent
comparisons across different experiments and facilitates the application of
advanced theoretical tools, ensuring that findings remain relevant and
comparable long after the initial analyses are conducted[3, 4].

Unfolding in high energy physics is traditionally approached by discretizing
the probability distribution into bins and transforming the unfolding problem
into matrix inversion[5]. The iterative Bayes’ theorem utilizes repeated
applications of Bayes’ theorem to invert the response matrix, with
regularization achieved by limiting the number of iterations based on sample
statistics and binning, thus avoiding overfitting despite initial conditions
set by the training truth rather than a flat distribution[6]. On the other
hand, Singular Value Decomposition (SVD), inverts the response matrix using
singular value decomposition techniques, where regularization involves
suppressing contributions from small singular values that signify high-
frequency fluctuations, thus enhancing stability and accuracy[1]. Lastly,
Tikhonov regularization, uses polynomial regularization to stabilize the
inversion of matrices when mapping multi-dimensional distributions, with the
optimal regularization parameter determined via an L-curve analysis[7].

To overcome the limitations associated with traditional methods, such as the
need for binning and the restriction to single observables, a new approach,
known as OmniFold, involve the iterative reweighting of a simulated event set
with machine learning techniques has been introduced[3]. The reweighting
process employs deep neural network-based classifiers to continuously estimate
probability ratio and adjust weights within a simulation[8]. The methodology
can be conceptualized as an EM-type algorithm, which iteratively maximizes the
likelihood estimate[9]. However, while the EM algorithm is widely used for its
convergence properties, it only guarantees convergence to a stationary point
of the likelihood function[10]. Studies indicate that the likelihood function
may have multiple local maxima, some of which can result in significantly
worse log-likelihood values compared to the global optimum[11, 12, 13].
Moreover, with random initialization, the EM algorithm tends to converge to
these suboptimal critical points with high probability, making the process
particularly sensitive to the initial parameter settings[14].

To surmount the aforementioned difficulties, this article introduces SwdFold,
which leverages the theory of optimal transport. Optimal transport theory
provides a robust mathematical foundation for comparing different data
distributions and possesses extensive applicative value across various
disciplines, including computer vision, machine learning, and statistical
inference[15, 16]. It has also been validated as an effective metric and loss
function in high-energy physics experiments for measuring the distributional
differences between diverse high-dimensional data sets[17, 18, 19, 20]. Our
SwdFold facilitates a gradual approximation to the true data distribution
without the need for data binning. This method does not rely on any prior
physical models during the unfolding process, enhancing its versatility and
applicability. On the optimization front, SwdFold offers a more stable energy
landscape for unfolding, which aids in pursuing an optimal solution more
effectively. Compared to binary weighted classification reweighting based
algorithms, it is less susceptible to becoming trapped in local optima,
thereby ensuring more accurate and robust unfolding results[21].

In exploring the estimation of probability density ratios via SwdFold, we
begin with the fundamental elements of one-dimensional probability measures,
transitioning gradually to practical computational domains. The core of this
study is the construction of the Cumulative Distribution Function (CDF) on the
set of real numbers, $\mathbb{R}$, for a probability measure $\alpha$[22]. The
mathematical definition is the integration of the measure $\alpha$ up to a
certain point $x\in\mathbb{R}$, expressed as:

$\forall
x\in\mathbb{R},\quad\mathcal{CDF}_{\alpha}(x)\stackrel{{\scriptstyle\text{
def. }}}{{=}}\int_{-\infty}^{x}\mathrm{~{}d}\alpha.$ (1)

For the 1-Wasserstein distance, its computation becomes the focal point of our
methodological exposition [23]. Against two probability density functions
$\alpha,\beta$, its expression is:

$WD(\alpha,\beta)=\int_{-\infty}^{\infty}\left|\mathcal{CDF}_{\alpha}(x)-\mathcal{CDF}_{\beta}(x)\right|\mathrm{~{}d}x$
(2)

This expression represents the essence of optimal transport, quantifying the
work required to transform one probability distribution into another.

For discrete cases, data vectors $A$ and $B$ should be sorted in ascending
order. Firstly, we then transfer their weights $w_{A}$ and $w_{B}$ into same
Range:

$w^{\prime}_{b}=\frac{w_{b}\sum w_{a}}{\sum w_{b}}$ (3)

Subsequently, we define the weighted cumulative distribution functions:

$\mathcal{CDF}_{a}(x_{i})=\sum_{j=1}^{i}w_{a}(x_{j}),\mathcal{CDF}_{b}(x_{i})=\sum_{j=1}^{i}w^{\prime}_{b}(x_{j})$
(4)

where $i=1,2,\ldots,n$ and $j=1,2,\ldots.m,n$ and $m$ are the lengths of $A$
and $B$ respectively.

Finally, in the Wasserstein distance computation phase, we calculate the
Wasserstein-1 distance as follows:

$WD\left(A,W_{A};B,W_{B}\right)=\sum_{i=0}^{n+m-1}\left|\mathcal{CDF}_{a}\left(x_{i}\right)-\mathcal{CDF}_{b}\left(x_{i}\right)\right|\cdot\left(x_{i+1}-x_{i}\right)$
(5)

where $x\in\left(\mathcal{A}\cup\mathcal{B}\right),x_{0}\leq
x_{1}\leq\ldots\leq x_{i}\leq\ldots\leq x_{m+n}$.

Thus, we complete the computation of the one-dimensional Wasserstein
distance[24].

Next, we extend to the multidimensional case. As Fig.1 shown, We adopt a
slicing approach, performing slices like radon trasform across different
feature quantities in each possible dimension
$A_{\theta_{1}},A_{\theta_{2}},\ldots$ [25, 26]. In practice, we can use huge
dimensional like $n=256$ to approximate. Thus We can get

$SWD\left(A,W_{A};B,W_{B}\right)=\frac{1}{n}\sum_{i=1}^{n}WD(A_{\theta_{i}},W_{A};B_{\theta_{i}},W_{B})$
(6)

So our optimization objective is to ensure that the estimated probability
density ratios are approximately consistent with those derived from optimal
transport theory across each possible dimension. This allows us to correctly
estimate the probability density ratio in high-dimensional cases. The
optimization function can be written as:

$f_{SWDRW}\left(A,W_{A};B,W_{B}\right)={\arg\min}_{W_{ratio}}SWD(A,W_{A}\cdot
W_{ratio};B,W_{B})$ (7)

After optimization by gradient descent using Adam optimizer[27], we can get
the probability density ratio between two weighted samples $(A,W_{A})$ and
$(B,W_{B})$ by

$f_{SWDRW}\left(A,W_{A};B,W_{B}\right)(x)=\frac{p_{(B,W_{B})}(x)}{p_{(A,W_{A})}(x)}$
(8)

where $p_{(X,W)}(x)$ means the probability density of x estimated by sample
$X$ with weight $W$. Figure 1 effectively demonstrates this convergence by
presenting both a normalized histogram and a cumulative distribution function
using two toy normal distribution’s samples with different parameters. These
visual comparisons showcase how the reweighted simulations iteratively
approach the actual data distribution. As indicated in the plots, the
successive iterations (10, 100, and 1000) yield increasingly accurate
approximations to the true data distribution. This progressive improvement,
evident in the narrowing differences between two samples, highlights the
efficacy of our optimization method in achieving precise reweighting toward
the observed distribution.

Figure 1: Example for Reweighted Simulation Data Across Iterations using
SWDRW. This figure consists of two plots that illustrate the effectiveness of
iterative reweighting of simulated data(Sim) to match actual observed
data(Data). The left plot data (Sim) and actual data (Data), alongside
reweighted simulated data at various iterations: 10 (Iter.10), 100 (Iter.100),
and 1000 (Iter.1000). The right plot is a Cumulative Distribution Function
(CDF) over the same value range, comparing the cumulative probability
distributions for the same categories of data. These plots are essential for
evaluating how effectively the iterative reweighting process adjusts the
simulated data to align with the distribution of the actual data,
demonstrating increased alignment as the number of iterations grows.

To demonstrate the flexibility and efficacy of our method, our study generate
toy Monte Carlo events that simulated the studied production and decay
processes in phase space. We simulate the particle level and detector level
event produced by electron-positron collisions, with a focus on the decay of
the $\psi(2S)$ resonance into a $\phi$ meson and a pair of $\pi^{+}$,
$\pi^{-}$, followed by the further decay of the $\phi$ meson into a pair of
$K^{+}$, $K^{-}$. We retained relevant physical information in the phase space
Monte Carlo, referred to as synthetic data: one at the particle level, termed
”Generation” representing theoretical predictions without detector effects in
phase space; and another at the simulation level, termed ”Simulation”
including simulated detector responses. A straightforward physical model was
used to assign weights to the generated data. These weights adjusted the
probability distribution of our simulated events to resemble that of another
toy Monte Carlo with a physically meaningful probability distribution.
Subsequently, we resampled the weighted generated data to produce an
unweighted dataset, which we refer to as nature. This nature dataset is
markedly different from the synthetic distribution, particularly notable in
the probability distribution of some key physical features. Similarly, nature
data at the particle level is referred to as “Truth” and at the simulation
level as “Data” The unfolding process uses the Simulation, Generation, and
Data datasets as inputs, training an algorithm capable of approximating the
Truth probability distribution. The objective is to correct experimental
effects present in the data, thereby yielding a distribution very close to
Truth. Our method allows for the simultaneous unfolding of multiple
observables related to the decay process, providing a comprehensive
understanding of the underlying physics. Additionally, having weight labels
enables us to more accurately compare the performance of various methods and
ascertain the accuracy of their estimated distributions.

In the specific iterative unfolding process of SwdFold, we split into the
reweight process(SwdRw) to learn detector level weights $w_{n}$ and particle
level model learning process to learn particle level weights $v_{n}$ [3] by

$f_{Model}(A,1;A,W_{RW})={\arg\max}_{g}w_{RW}log(g(a))$ (9)

In the first interation of reweight process, we set particle level weight
$\nu_{0}$ as uniform weights, then we can optimize by the following step:

$w_{n}(x)=\nu_{n-1}(x)f_{\text{SWDRW}}(\text{Sim},\nu_{n-1};\text{Data},1)(x)$
(10)
$\nu_{n}(y)=\nu_{n-1}(y)f_{\text{Model}}(\text{Gen},1;\text{Gen},w_{n})(y)$
(11)

where $(x,y)$ is paired event in synthetic dataset, where x is detector level
event, and y is corresponding particle level event. Then after n iterations,
we can estimate the unfolded weights as

$p_{(\text{Truth},1)}^{(n)}(y)\approx\nu_{n}(y)p_{(\text{Gen},1)}(y)$ (12)

We employed a multi-layer fully connected neural network architecture
incorporating shortcut connections akin to those found in residual network
designs. The internal representation of the network is characterized by 256
-dimensional latent space features, allowing for intricate modeling of the
underlying structure of the data. This network processes a diverse set of
input features, including kinematic properties and particle identities related
to decay events. These features encompass the Cartesian momentum components (
$\left.p_{x},p_{y},p_{z}\right)$, total energy (E), magnitude of momentum
vector $(p)$, transverse momentum ( $p_{T}$ ), and angular coordinates (
$\theta$ and $\phi$ ) for $K^{+},K^{-}$, $\pi^{+},\pi^{-},\phi$, and
$\pi^{+}\pi^{-}$ pairs. Additionally, the model is sensitive to the involved
particle types, namely $\mathrm{K}^{+},\mathrm{K}^{-},\pi^{+}$, and $\pi^{-}$.
To further enrich the feature space, the model also incorporates several key
invariant mass combinations, particularly the $M_{\phi}$ of the $\phi$ meson
and $M_{\pi^{+}\pi^{-}}$ for the $\pi^{+}\pi^{-}$ pair. Our model was trained
using the PyTorch library, and the optimization of network parameters was
conducted using the AdamW optimizer, a variant of the Adam optimizer that
includes a decoupled weight decay regularization [28].

Our comprehensive analysis underscores the efficacy of the SwdFold method,
particularly in the context of unfolding and reweighting within our
challenging particle physics datasets. The unfolding task in our study was
confronted with the intricate difficulty of subtle detector effects, which
typically do not manifest in variable distributions when we unfold all the
observable in Fig.6. This issue was most pronounced when examining the
resolution of the $\phi$ mass, where a comparison of detector-level events
(Data) and particle-level events (Truth) revealed minimal apparent
discrepancies. Nonetheless, as Fig.2 shown, SwdFold demonstrated its
exceptional capacity to refine the $\phi$ mass distribution, aligning it more
closely with the truth distribution and thereby enabling a more precise
measurement of the $\phi$ mass width.

Figure 2: Histogram Analysis of Unfolding Methodologies. The performance of
our proposed SwdFold method is shown with brown dots, compared with the
OmniFold method with black dots. Our SwdFold method exhibits a high degree of
congruence with the truth event(Truth), signifying a successful reweighting of
the generation event distribution(Gen) to estimate the truth event
distribution(Truth), as particularly evident in the sharp peak around 1.020
GeV/c² in the plot.

The reweighting process faced its own set of challenges, especially when
comparing synthetic datasets with natural datasets. In the natural dataset,
the $\pi^{+}\pi^{-}$mass distribution displayed pronounced peaks, signaling
resonance in decay-peaks that were missing from the synthetic dataset.
SwdFold’s robust algorithmic framework not only succeeded in accurately
representing these resonant peaks but also replicated the corresponding widths
observed in nature in Fig.3. These advancements in reweighting signify
SwdFold’s nuanced capability to accurately reflect complex physical phenomena,
thereby enhancing data analysis and parameter estimation within the domain of
particle physics research.

Figure 3: Histogram Analysis of Reweight Methodologies. We presents a detailed
examination of the $\pi^{+}\pi^{-}$mass distribution across three distinct
mass ranges: $(0.8,1.05),(1.1,1.4)$, and
$(1.5,1.8)\mathrm{GeV}/\mathrm{c}^{2}$. The performance of the SwdFold method,
illustrated by the black line, showcases its ability to reveal the resonance
peaks by reweighting the Generation Event(Gen), aligning closely with the
Truth Event(Truth) in both peak value and width.

Our SwdFold method not only achieves high agreement with the one-dimensional
distributions, has demonstrated remarkable success in capturing and preserving
the intricate correlations present within the multi-dimensional phase space of
particle interactions. As evidenced by the Dalitz plots presented in Fig.4,
the close alignment of the SwdFold plot with the Truth data attests to the
method’s precision in reweighting and its adeptness in handling the complex
correlations characteristic of particle physics data. This proficiency extends
beyond mere visual congruence; it indicates a deeper, quantitative concordance
across the full kinematic domain, showcasing the potential of our method for
high-precision reweighting and parameter estimation. The ability of SwdFold to
deliver physically meaningful results, positions it as a powerful tool for
enhancing the accuracy and reliability of analyses in experimental particle
physics.

Figure 4: Visualization of the Dalitz plots for the $\phi\pi$ system. The left
panel depicts the distribution from the generated dataset (Gen), which serves
as the simulated data before any reweighting. The center panel represents the
Truth data, reflecting the actual particle interactions without the influence
of detector effects or experimental noise. The right panel showcases the
distribution post the application of the SwdFold method. The SwdFold plot is
close to the Truth distribution, demonstrating its successful reweighting.

In conclusion, we presents SwdFold, a pioneering reweighting and unfolding
method grounded in the principles of optimal transport, which addresses the
longstanding challenges faced by traditional unfolding techniques in high-
energy physics. Our method significantly mitigates issues related to model
dependency, binning artifacts, and sensitivity to statistical fluctuations,
which are prevalent in conventional methods. By deploying sliced wasserstein
distance theory to approximate the probability density ratio of simulation
dataset and experimental dataset, SwdFold can reweight the particle level
simulation event to estimate the true physical result without detector effect.
The empirical validations using high-energy physics datasets, specifically the
simulated events of electron-positron collisions and the decay processes
$\psi(2S)\rightarrow\phi K^{+}K^{-},\phi\rightarrow K^{+}K^{-}$, demonstrate
that SwdFold can accurately recover the true distribution of physical events
with remarkable precision. Notably, the method excels in maintaining the
integrity of multidimensional correlations and nuanced physical features,
which are often obscured by detector effects and other experimental noise.
SwdFold’s innovative approach to data unfolding and its ability to handle
high-dimensional data without binning represent substantial advancements in
the field. These enhancements enable physicists to derive more accurate
insights from experimental data, thereby facilitating a deeper understanding
of underlying physical phenomena and contributing to the ongoing development
of theoretical models. Future work will focus on refining the algorithm’s
efficiency and exploring its application in other complex systems where data
integrity is critical. The adaptability and performance of SwdFold affirm its
potential as a valuable tool for a wide range of applications in experimental
physics and beyond, promising to significantly impact how data-driven
discoveries are made in the sciences.

Looking ahead, SwdFold paves new possibilities for directly training deep
learning models in an unsupervised manner for tasks in unfolding and
reweighting. This approach holds promising prospects not only for future
research in particle physics. To foster continued research and development in
the field, we plan to make our SwdFold code and related datasets publicly
available. This initiative will enable other researchers to apply our method
in various settings, further validating and enhancing its applicability.

## References

  * Hoecker and Kartvelishvili [1996] Andreas Hoecker and Vakhtang Kartvelishvili.  Svd approach to data unfolding.  _Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment_ , 372(3):469–481, 1996.
  * Andreev et al. [2022] Vladimir Andreev, M Arratia, A Baghdasaryan, A Baty, K Begzsuren, A Belousov, Arthur Bolz, V Boudry, G Brandt, D Britzger, et al.  Measurement of lepton-jet correlation in deep-inelastic scattering with the h1 detector using machine learning for unfolding.  _Physical review letters_ , 128(13):132002, 2022.
  * Andreassen et al. [2020] Anders Andreassen, Patrick T Komiske, Eric M Metodiev, Benjamin Nachman, and Jesse Thaler.  Omnifold: a method to simultaneously unfold all observables.  _Physical review letters_ , 124(18):182001, 2020.
  * Cowan [2002] Glen Cowan.  A survey of unfolding methods for particle physics.  In _Conf. Proc. C_ , volume 203181, pages 248–257, 2002.
  * Adye [2011] Tim Adye.  Unfolding algorithms and tests using roounfold.  _arXiv preprint arXiv:1105.1160_ , 2011.
  * D’Agostini [1995] Giulio D’Agostini.  A multidimensional unfolding method based on bayes’ theorem.  _Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment_ , 362(2-3):487–498, 1995.
  * Schmitt [2012] Stefan Schmitt.  Tunfold, an algorithm for correcting migration effects in high energy physics.  _Journal of Instrumentation_ , 7(10):T10003, 2012.
  * Andreassen and Nachman [2020] Anders Andreassen and Benjamin Nachman.  Neural networks for full phase-space reweighting and parameter tuning.  _Physical Review D_ , 101(9):091901, 2020.
  * Andreassen et al. [2021] Anders Andreassen, Patrick T Komiske, Eric M Metodiev, Benjamin Nachman, Adi Suresh, and Jesse Thaler.  Scaffolding simulations with deep learning for high-dimensional deconvolution.  _arXiv preprint arXiv:2105.04448_ , 2021.
  * Kolouri et al. [2018a] Soheil Kolouri, Gustavo K Rohde, and Heiko Hoffmann.  Sliced wasserstein distance for learning gaussian mixture models.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 3427–3436, 2018a.
  * Améndola et al. [2016] Carlos Améndola, Mathias Drton, and Bernd Sturmfels.  Maximum likelihood estimates for gaussian mixtures are transcendental.  In _Mathematical Aspects of Computer and Information Sciences: 6th International Conference, MACIS 2015, Berlin, Germany, November 11-13, 2015, Revised Selected Papers 6_ , pages 579–590. Springer, 2016.
  * Hoffmann [2005] Heiko Hoffmann.  _Unsupervised learning of visuomotor associations_ , volume 11.  Bielefeld University, 2005.
  * Jin et al. [2016] Chi Jin, Yuchen Zhang, Sivaraman Balakrishnan, Martin J Wainwright, and Michael I Jordan.  Local maxima in the likelihood of gaussian mixture models: Structural results and algorithmic consequences.  _Advances in neural information processing systems_ , 29, 2016.
  * Jian and Vemuri [2010] Bing Jian and Baba C Vemuri.  Robust point set registration using gaussian mixture models.  _IEEE transactions on pattern analysis and machine intelligence_ , 33(8):1633–1645, 2010.
  * Torres et al. [2021] Luis Caicedo Torres, Luiz Manella Pereira, and M Hadi Amini.  A survey on optimal transport for machine learning: Theory and applications.  _arXiv preprint arXiv:2106.01963_ , 2021.
  * Santambrogio [2015] Filippo Santambrogio.  Optimal transport for applied mathematicians.  _Birkäuser, NY_ , 55(58-63):94, 2015.
  * Pan et al. [2023] Chu-Cheng Pan, Xiang Dong, Yu-Chang Sun, Ao-Yan Cheng, Ao-Bo Wang, Yu-Xuan Hu, and Hao Cai.  Event generation and consistence test for physics with sliced wasserstein distance.  _arXiv preprint arXiv:2310.17897_ , 2023.
  * Howard et al. [2022] Jessica N Howard, Stephan Mandt, Daniel Whiteson, and Yibo Yang.  Learning to simulate high energy particle collisions from unlabeled data.  _Scientific Reports_ , 12(1):7567, 2022.
  * Fraser et al. [2022] Katherine Fraser, Samuel Homiller, Rashmish K Mishra, Bryan Ostdiek, and Matthew D Schwartz.  Challenges for unsupervised anomaly detection in particle physics.  _Journal of High Energy Physics_ , 2022(3):1–31, 2022.
  * Cai et al. [2020] Tianji Cai, Junyi Cheng, Nathaniel Craig, and Katy Craig.  Linearized optimal transport for collider events.  _Physical Review D_ , 102(11):116019, 2020.
  * Kolouri et al. [2018b] Soheil Kolouri, Gustavo K Rohde, and Heiko Hoffmann.  Sliced wasserstein distance for learning gaussian mixture models.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 3427–3436, 2018b.
  * Deisenroth et al. [2020] Marc Peter Deisenroth, A Aldo Faisal, and Cheng Soon Ong.  _Mathematics for machine learning_.  Cambridge University Press, 2020.
  * Kolouri et al. [2017] Soheil Kolouri, Se Rim Park, Matthew Thorpe, Dejan Slepcev, and Gustavo K Rohde.  Optimal mass transport: Signal processing and machine-learning applications.  _IEEE signal processing magazine_ , 34(4):43–59, 2017.
  * Peyré et al. [2019] Gabriel Peyré, Marco Cuturi, et al.  Computational optimal transport: With applications to data science.  _Foundations and Trends® in Machine Learning_ , 11(5-6):355–607, 2019.
  * Toft [1996] Peter Toft.  The radon transform.  _Theory and Implementation (Ph. D. Dissertation)(Copenhagen: Technical University of Denmark)_ , 1996.
  * Bonneel et al. [2014] Nicolas Bonneel, Julien Rabin, Gabriel Peyre, and Hanspeter Pfister.  Sliced and radon wasserstein barycenters of measures.  _Journal of Mathematical Imaging and Vision_ , 2014.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Loshchilov and Hutter [2017] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.

##  Appendix A Additional Figure

in 0,…, 19

in 20,…,39

Figure 6: The unfolding result for all observables


# CheXseg: Combining Expert Annotations with DNN-generated Saliency Maps for
X-ray Segmentation

Soham Gadgil<EMAIL_ADDRESS>
Mark Endo11footnotemark: 1<EMAIL_ADDRESS>
Emily Wen11footnotemark: 1<EMAIL_ADDRESS>
Andrew Y. Ng<EMAIL_ADDRESS>
Pranav Rajpurkar<EMAIL_ADDRESS>
Department of Computer Science Contributed equally Stanford University

###### Abstract

Medical image segmentation models are typically supervised by expert
annotations at the pixel-level, which can be expensive to acquire. In this
work, we propose a method that combines the high quality of pixel-level expert
annotations with the scale of coarse DNN-generated saliency maps for training
multi-label semantic segmentation models. We demonstrate the application of
our semi-supervised method, which we call CheXseg, on multi-label chest x-ray
interpretation. We find that CheXseg improves upon the performance (mIoU) of
fully-supervised methods that use only pixel-level expert annotations by 13.4%
and weakly-supervised methods that use only DNN-generated saliency maps by
91.2%. Furthermore, we implement a semi-supervised method using knowledge
distillation and find that though it is outperformed by CheXseg, it exceeds
the performance (mIoU) of the best fully-supervised method by 4.83%. Our best
method is able to match radiologist agreement on three out of ten pathologies
and reduces the overall performance gap by 71.6% as compared to weakly-
supervised methods.

###### keywords:

Semi-Supervised Segmentation, Saliency Maps, Localization Performance

##  1 Introduction

The “black box” nature of neural networks represents a barrier to physicians’
trust and model adoption in the clinical setting [Kelly et al. (2019)].
Saliency maps are a popular set of explanation methods that highlight regions
of the image that are important for disease classification, but they have been
shown to be untrustworthy for medical image interpretation [Eitel and Ritter
(2019), Crosby et al. (2020), Young et al. (2019), Arun et al. (2020)].
Segmentation models can produce more accurate pixel-level maps, but their
training is typically limited by the time-consuming process of collecting
expert annotations. The combination of saliency maps generated from widely
available classification models and a limited amount of expert annotations for
training medical image segmentation models may be able to provide higher
quality segmentations at a lower cost, but this approach remains relatively
unexplored.

In this work, we develop CheXseg, a semi-supervised method for multi-pathology
segmentation that leverages both the pixel-level expert annotations and the
saliency maps generated by image classification models. First, we find that
CheXseg achieves a mean IoU of 0.299, outperforming both fully-supervised
(mIoU of 0.263) and weakly-supervised (mIoU of 0.156) methods. Second, as an
additional comparison, we implement a semi-supervised knowledge distillation
method that utilizes a large number of unlabeled images and a fully-supervised
teacher model [Wang and Yoon (2021)]. We find that the student model improves
upon its teacher’s performance (3.42% average improvement in mIoU) but has a
6.65% worse performance (mIoU) than CheXseg. Third, we find that initializing
the segmentation encoder with weights learned from supervised classification
of the same task leads to higher performance than using a self-supervised MoCo
initialization [He et al. (2020)] or ImageNet initialization [Deng et al.
(2009)]. Finally, CheXseg reduces the overall gap to radiologist localization
performance (mIoU) by 71.6% compared to solely using DNN-generated saliency
maps. We expect this method to be broadly useful for medical image
segmentation, where classification models are widely available and expert
annotations are expensive.

##  2 Related Work

###  2.1 Weakly-Supervised Semantic Segmentation

In this work, we focus on an approach in which classification models trained
with image-level labels are used to create pixel-level pseudo-labels [Yao and
Gong (2020), Ciga and Martel (2019), Ouyang et al. (2019)]. These pseudo-
labels can then be utilized to train a segmentation model. Our paper is
closely related to Viniavskyi et al. (2020), which proposes a deep CNN-based
approach that generates pseudo-labels by applying an Inter-pixel Relation
Network (IRNet) Ahn et al. (2019) to improve Grad-CAM++ Chattopadhay et al.
(2018) generated activation maps. This approach is evaluated on the SIIM-ACR
Pneumothorax dataset. In our work, we generate pseudo-labels with IRNet and
extend the approach to the semi-supervised setting for a larger set of
pathologies.

Several other methods also propagate class activation from areas of high-
confidence to similar adjacent regions Kolesnikov and Lampert (2016); Huang et
al. (2018); Ahn and Kwak (2018). We choose to build upon Viniavskyi et al.
(2020) rather than these approaches because it has the best performance on the
PASCAL VOC 2012 Everingham et al. (2010) validation set (mIoU of 0.646) and
was shown to perform well on a medical imaging task.

###  2.2 Semi- and Fully-Supervised Semantic Segmentation

Semi-supervised methods use a combination of expert pixel-level annotations
and pseudo-labels to train semantic segmentation models. Some weakly-
supervised methods have been extended to semi-supervised methods through the
replacement of pseudo-labels. The weakly-supervised SGAN model Yao and Gong
(2020) was adapted to a semi-supervised setting by replacing a subset of the
saliency maps with the corresponding manually annotated ground truth labels.
In this work, we use a similar idea of utilizing radiologist annotated labels
in addition to saliency maps to train the segmentation model, extending
Viniavskyi et al. (2020)’s weakly-supervised model.

Though weakly- and semi-supervised methods can perform well, fully-supervised
methods are still considered the upper-bound Chan et al. (2020). Many fully-
supervised semantic segmentation approaches have been proposed for chest
X-rays Sirazitdinov et al. (2019); Jaiswal et al. (2019), but none of these
extend their work to the semi-supervised setting.

###  2.3 Knowledge Distillation

Knowledge distillation is a method used to transfer knowledge from a trained
teacher model to a student model Hinton et al. (2015); Liu et al. (2020). This
has shown promise in a self-distillation setup Xie et al. (2020), in which a
large set of unlabeled data can be used on a single model architecture which
takes on the dual role as both the teacher and student. Previous work has
shown the student model can generalize better than the teacher and outperform
the teacher on natural images Xie et al. (2020); Chen et al. (2020). In our
work, we extend this idea to self-distillation for semantic segmentation of
medical images.

##  3 Methods

###  3.1 Setup

The multi-label semantic segmentation task is to classify each pixel of a
chest X-ray image into zero or more of 10 possible pathologies: Airspace
Opacity, Atelectasis, Cardiomegaly, Consolidation, Edema, Enlarged
Cardiomediastinum, Lung Lesion, Pleural Effusion, Pneumothorax, and Support
Devices.

We utilize CheXpert Irvin et al. (2019), an existing large dataset with
224,316 chest X-rays of 65240 patients. This dataset features image-level
labels obtained using an automated labeler that detects the aforementioned
pathologies from radiology reports. A subset of the dataset is hand-annotated
by radiologists at the pixel level. In our work, we use a set of 200
radiologist-annotated chest X-rays as the training set for the fully-
supervised method. We use an additional set of 500 radiologist-annotated
images for the test set.

Models are evaluated by their average performance on the semantic segmentation
task across the ten pathologies of interest. For each of the pathologies, the
IoU (Intersection-over-Union) score is computed. We report the mIoU (mean IoU)
score, which is the average IoU score across all pathologies.

###  3.2 CheXseg

We develop CheXseg, a semi-supervised method for multi-pathology segmentation
that leverages both the pixel-level expert annotations and the saliency maps
generated by image classification models. In this method, a DenseNet121
Iandola et al. (2014) classification model, trained on the entire CheXpert
train set, is first used to generate saliency maps using Grad-CAM Selvaraju et
al. (2017). This approach uses the classification model outputs to create a
coarse localization map highlighting the image regions important for
prediction. The saliency maps are further processed to create per-pixel
segmentation masks, referred to as weak pseudo-labels, by using either a
thresholding scheme or an Inter-Pixel Relation Network (IRNet) Ahn et al.
(2019). IRNet takes these generated CAMs and tries to improve the seeds by
training two output branches, a displacement vector field and a class boundary
map. Details about these methods are provided in
appendix:weak_pseudo_label_generation.

Once the pseudo-labels have been generated, we combine them with pixel-level
expert annotations in a semi-supervised way to train semantic segmentation
models. In order to adjust the level of contribution each label type has on
model training, we alter the sampling rate of each respective label type.

For comparison, we train fully-supervised segmentation models (solely using
pixel-level annotations) and weakly-supervised segmentation models (solely
using pseudo-labels). For a semi-supervised comparison, we implement a
distillation method in which the predicted probability outputs of the fully-
supervised models on unlabeled images act as pseudo-labels to train student
models. All our methods utilize DeepLabv3+ as the core semantic segmentation
model Chen et al. (2018). We experiment with various encoder initializations
to transfer knowledge from the classification task to the segmentation task.
Our experiments utilize a ResNet encoder architecture He et al. (2016). Figure
LABEL:fig:supervision_levels shows a visual representation of the different
supervision strategies used. All training details are available in
appendix:training_details.

fig:supervision_levels

Figure 1: Workflows of the different methods analysed for chest X-ray
segmentation

##  4 Experiments

###  4.1 Combining Weak and Full Supervision

We investigate the segmentation performance of combining DNN-generated
saliency maps and expert annotations with various sampling ratios. We use 100
saliency maps in combination with 200 annotated pixel-level labels and explore
different weightings between the two types of labels. We vary the probability
of selecting an expert annotation in a single batch during training, $p\in
0,0.2,0.4,0.6,0.8.0.85,0.9$. Thus, $p$ determines the expected fraction of
images with expert annotations in a single batch. For each value of $p$, we
perform three trials containing different sets of the 100 saliency maps. The
results reported are the mIoU scores obtained by averaging across these 3
trials. The segmentation model is initialized with CheXpert encoder weights
since it performs the best as observed in experiment 4.3. We also compare the
performance of this semi-supervised model with the weakly-supervised and
fully-supervised models.

#### Results

We find that for both Grad-CAM and IRNet, there is an inverted U-shape trend
in performance as we increase $p$. There is a sharp increase in performance as
$p$ increases from 0 (mIoU score of 0.156), and then the curves remain
relatively flat before dropping off when $p=1$. Specifically, $p=0.9$
(CheXseg) gives the best mIoU performance of 0.299 and 0.294 for CAM and IRNet
respectively. This high weighting of pixel-level labels takes advantage of the
more accurate information encoded within these labels as compared to the
saliency maps. The reduced performance for the fully-supervised case ($p=1$,
mIoU score of 0.263) is likely attributed to the weak labels no longer being
utilized in training. The size of the train set shrinks, and the model does
not benefit from the variation and the scale provided by the weak pseudo-
labels. For the weakly-supervised case ($p=0$, mIoU score of 0.156), the poor
performance may be due to the absence of any pixel-level expert annotations to
guide the model predictions. Detailed results are shown in Figure
LABEL:fig:pneumo_consolid_plot.

fig:pneumo_consolid_plot

Figure 2: IoU scores of semi-supervised segmentation models (averaged across 3
trials) using either Grad-CAM or IRNet to generate weak labels. The DeepLabV3+
and ResNet18 setup is used with CheXpert encoder initializion. p is the
probability of selecting a pixel-level labeled training sample in the current
batch. (a) is the average IoU score calculated across all pathologies. (b),
(c), and (d) are IoU scores for Pleural, Consolidation, and Edema
respectively. $\textit{p}=0$ represents the weakly-supervised case while
$\textit{p}=1$ represents the fully-supervised case. Full results in
tab:weightings.

###  4.2 Comparison to Semi-Supervised Distillation

We implement a semi-supervised distillation approach to see whether it (1)
outperforms the fully-supervised methods, and (2) whether it outperforms
CheXseg. Whereas CheXseg is trained on saliency maps and expert annotations
simultaneously, for this method, we now distill the knowledge of the fully-
supervised setup using the teacher model’s predictions on CheXpert images. The
distillation method does not require the use of an existing classification
model (except for initializing the encoder with CheXpert classification
weights). We evaluate the mIoU of the self-distilled models trained on
different amounts of data and with different encoder initializations.

#### Results

The student model improves upon its teacher’s performance for each encoder
initialization. With random initialization and MoCo-CXR initialization, the
student outperforms the teacher using both 1% and 10% of the training data.
With CheXpert initialization, the student only outperforms the teacher
performance using 10% of the training data, and even then it only outperforms
the teacher marginally (1.69% improvement in mIoU score). This marginal
improvement may be explained by the fact that the pre-trained CheXpert encoder
weights have already learned task-specific information from the downstream
classification task. With MoCo-CXR initialization, the encoder weights have
also learned from the data to a lesser extent with self-supervision. The
random encoder initialization setup has the most to gain from distillation
since it has not previously seen the CheXpert images. Although the
distillation method outperforms the fully-supervised teacher models, it is
outperformed by CheXseg (6.65% difference in mIoU score). This is likely
because saliency maps are a stronger supervision source than unlabeled data.
Detailed distillation results are shown in Figure
LABEL:fig:distillation_results.

fig:distillation_results

Figure 3: mIoU of distilled models using either random initialization, MoCo-
CXR initialization, or CheXpert initialization. Dashed line denotes teacher
model performance, which is initialized with same encoder weights as student
model and fine-tuned on a small set of fully-supervised images. In all three
cases, with sufficient distillation training data ($1\%$ for random
initialization and MoCo-CXR initialization and $10\%$ for CheXpert
initialization), student improves upon teacher model performance.

###  4.3 Comparing Encoder Initializations

We investigate the impact of using various encoder initializations on
segmentation performance. In the fully-supervised and weakly-supervised
methods, we initialize the encoder weights to either a CheXpert classification
model Irvin et al. (2019), MoCo-CXR Sowrirajan et al. (2020), ImageNet Deng et
al. (2009), or random. In our distillation experiment, we only compare
CheXpert, MoCo-CXR, or random encoder initializations.

#### Results

We find that for all methods, the best models are initialized with CheXpert
encoder weights. This may be expected since CheXpert weights are learned from
supervised learning on the same dataset that we use for segmentation. The
models initialized with MoCo-CXR weights have similar performance to the
models with ImageNet encoder initialization. Figure
LABEL:fig:weakly_supervised_results shows the detailed results for fully-
supervised and weakly-supervised encoder initializations, and Figure
LABEL:fig:distillation_results shows the detailed results for distillation
encoder initializations.

fig:weakly_supervised_results[mIoU of fully-supervised segmentation models
with DeepLabV3+ and ResNet18 initialized with various weights. CheXpert
encoder initialization results in best performance (average 0.263).] [mIoU of
weakly-supervised segmentation models with DeepLabV3+ and ResNet18 setup using
either Grad-CAM or IRNet pseudo-labels and various encoder initializations.
IRNet outperforms CAMs when using CheXpert encoder initialization (0.156 vs
0.142), but underperforms when using other initializations (0.111 vs 0.136
with Random, 0.128 vs 0.139 with MoCo-CXR, and 0.124 vs 0.139 with ImageNet).
Full comparisons with confidence intervals in tab:weakresults.]

Figure 4: Performance of fully-supervised and weakly-supervised methods

###  4.4 Comparison to Radiologists

Compared to our best weakly-supervised method, CheXseg reduces the overall
performance (mIoU) gap with radiologists by 71.6%. CheXseg outperforms
radiologists in terms of IoU score on the segmentation of Atelectasis (165%
higher), Airspace Opacity (79% higher), and Pleural Effusion (26% higher),
while performing worse on the remaining pathologies. Detailed results are
shown in Figure LABEL:fig:human_comparison.

fig:human_comparison

Figure 5: IoUs of radiologists, our best semi-supervised method (CheXseg), and
our best weakly-supervised method. CheXseg uses Grad-CAM while the weakly-
supervised method uses IRNet.

##  5 Discussion

In this work, we develop CheXseg, a semi-supervised method for multi-pathology
segmentation that leverages the benefits of both available medical image
classification models and expert pixel-level annotations.

How does CheXseg performance compare to fully-supervised and weakly-supervised
methods? We find that with a weighted sampling of saliency maps and expert
annotations, our proposed method outperforms both the fully-supervised and
weakly-supervised methods alone. An expert annotation to saliency map sample
ratio of 0.9 (CheXseg) gives the best mIoU score of 0.299, compared to 0.263
for fully-supervised and 0.156 for weakly supervised. This suggests a tradeoff
between emphasizing the accurate information of expert pixel-level annotations
and incorporating additional but noisy cues from weak pseudo-labels.

How does the performance of a semi-supervised distillation method compare to
CheXseg and fully-supervised methods? We find that the implemented
distillation method has a 6.65% worse performance (mIoU) than CheXseg. That
said, we observe that self-distillation can improve learning performance for
each fully-supervised model setup. Previous work shows similar results for
natural images Xie et al. (2020); Chen et al. (2020). Future work may
investigate the result of introducing different sources of noise during
training, such as data augmentation, stochastic depth, and dropout, which may
further improve performance from teacher to student.

How do segmentation models initialized with CheXpert, ImageNet, MoCo-CXR, and
random encoder weights compare? We find that CheXpert encoder initialization
achieves the highest performance, followed by self-supervised (MoCo-CXR)
initialization and ImageNet initialization. Random encoder initialization has
the worst performance, as no transfer learning occurs. Since CheXpert weights
are pre-trained using image-level labels for the same tasks, it is expected
that this knowledge transfers well to segmentation on the same dataset and the
same set of tasks. Whereas classification models with MoCo-CXR encoder
initialization have been found to outperform classification models with
ImageNet encoder initialization Sowrirajan et al. (2020), we find that the
these two initializations have approximately the same performance for
segmentation.

In closing, our work proposes a simple semi-supervised method that combines
the benefits of the wide availability of classification models with the
quality of expert annotations. We expect our method may be broadly useful,
able to lower the cost of development and improve the performance of medical
image segmentation models.

## References

  * Ahn and Kwak (2018) Jiwoon Ahn and Suha Kwak.  Learning pixel-level semantic affinity with image-level supervision for weakly supervised semantic segmentation.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 4981–4990, 2018.
  * Ahn et al. (2019) Jiwoon Ahn, Sunghyun Cho, and Suha Kwak.  Weakly supervised learning of instance segmentation with inter-pixel relations.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 2209–2218, 2019.
  * Arun et al. (2020) Nishanth Arun, Nathan Gaw, Praveer Singh, Ken Chang, Mehak Aggarwal, Bryan Chen, Katharina Hoebel, Sharut Gupta, Jay Patel, Mishka Gidwani, et al.  Assessing the (un) trustworthiness of saliency maps for localizing abnormalities in medical imaging.  _arXiv preprint arXiv:2008.02766_ , 2020.
  * Chan et al. (2020) Lyndon Chan, Mahdi S Hosseini, and Konstantinos N Plataniotis.  A comprehensive analysis of weakly-supervised semantic segmentation in different image domains.  _International Journal of Computer Vision_ , pages 1–24, 2020.
  * Chattopadhay et al. (2018) Aditya Chattopadhay, Anirban Sarkar, Prantik Howlader, and Vineeth N Balasubramanian.  Grad-cam++: Generalized gradient-based visual explanations for deep convolutional networks.  In _2018 IEEE Winter Conference on Applications of Computer Vision (WACV)_ , pages 839–847. IEEE, 2018.
  * Chen et al. (2018) Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam.  Encoder-decoder with atrous separable convolution for semantic image segmentation.  In _Proceedings of the European conference on computer vision (ECCV)_ , pages 801–818, 2018.
  * Chen et al. (2020) Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton.  Big self-supervised models are strong semi-supervised learners, 2020.
  * Ciga and Martel (2019) Ozan Ciga and Anne L Martel.  Learning to segment images with classification labels.  _arXiv preprint arXiv:1912.12533_ , 2019.
  * Crosby et al. (2020) Jennie Crosby, Sophia Chen, Feng Li, Heber MacMahon, and Maryellen Giger.  Network output visualization to uncover limitations of deep learning detection of pneumothorax.  In Frank W. Samuelson and Sian Taylor-Phillips, editors, _Medical Imaging 2020: Image Perception, Observer Performance, and Technology Assessment_ , volume 11316, pages 125 – 128. International Society for Optics and Photonics, SPIE, 2020.  10.1117/12.2550066.  URL https://doi.org/10.1117/12.2550066.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pages 248–255. Ieee, 2009.
  * Eitel and Ritter (2019) Fabian Eitel and Kerstin Ritter.  Testing the robustness of attribution methods for convolutional neural networks in mri-based alzheimer’s disease classification, 2019.
  * Everingham et al. (2010) Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman.  The pascal visual object classes (voc) challenge.  _International journal of computer vision_ , 88(2):303–338, 2010.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * He et al. (2020) Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9729–9738, 2020.
  * Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network, 2015.
  * Huang et al. (2018) Zilong Huang, Xinggang Wang, Jiasi Wang, Wenyu Liu, and Jingdong Wang.  Weakly-supervised semantic segmentation network with deep seeded region growing.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 7014–7023, 2018.
  * Iandola et al. (2014) Forrest Iandola, Matt Moskewicz, Sergey Karayev, Ross Girshick, Trevor Darrell, and Kurt Keutzer.  Densenet: Implementing efficient convnet descriptor pyramids.  _arXiv preprint arXiv:1404.1869_ , 2014.
  * Irvin et al. (2019) Jeremy Irvin, Pranav Rajpurkar, Michael Ko, Yifan Yu, Silviana Ciurea-Ilcus, Chris Chute, Henrik Marklund, Behzad Haghgoo, Robyn Ball, Katie Shpanskaya, et al.  Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 590–597, 2019.
  * Jaiswal et al. (2019) Amit Kumar Jaiswal, Prayag Tiwari, Sachin Kumar, Deepak Gupta, Ashish Khanna, and Joel JPC Rodrigues.  Identifying pneumonia in chest x-rays: A deep learning approach.  _Measurement_ , 145:511–518, 2019.
  * Kelly et al. (2019) Christopher J Kelly, Alan Karthikesalingam, Mustafa Suleyman, Greg Corrado, and Dominic King.  Key challenges for delivering clinical impact with artificial intelligence.  _BMC medicine_ , 17(1):1–9, 2019.
  * Kolesnikov and Lampert (2016) Alexander Kolesnikov and Christoph H Lampert.  Seed, expand and constrain: Three principles for weakly-supervised image segmentation.  In _European conference on computer vision_ , pages 695–711. Springer, 2016.
  * Liu et al. (2020) Yifan Liu, Changyong Shun, Jingdong Wang, and Chunhua Shen.  Structured knowledge distillation for dense prediction, 2020.
  * Ouyang et al. (2019) Xi Ouyang, Zhong Xue, Yiqiang Zhan, Xiang Sean Zhou, Qingfeng Wang, Ying Zhou, Qian Wang, and Jie-Zhi Cheng.  Weakly supervised segmentation framework with uncertainty: A study on pneumothorax segmentation in chest x-ray.  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , pages 613–621. Springer, 2019.
  * Selvaraju et al. (2017) Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra.  Grad-cam: Visual explanations from deep networks via gradient-based localization.  In _Proceedings of the IEEE international conference on computer vision_ , pages 618–626, 2017.
  * Sirazitdinov et al. (2019) Ilyas Sirazitdinov, Maksym Kholiavchenko, Tamerlan Mustafaev, Yuan Yixuan, Ramil Kuleev, and Bulat Ibragimov.  Deep neural network ensemble for pneumonia localization from a large-scale chest x-ray database.  _Computers & Electrical Engineering_, 78:388–399, 2019\.
  * Sowrirajan et al. (2020) Hari Sowrirajan, Jingbo Yang, Andrew Y Ng, and Pranav Rajpurkar.  Moco pretraining improves representation and transferability of chest x-ray models.  _arXiv preprint arXiv:2010.05352_ , 2020.
  * Viniavskyi et al. (2020) Ostap Viniavskyi, Mariia Dobko, and Oles Dobosevych.  Weakly-supervised segmentation for disease localization in chest x-ray images.  In _International Conference on Artificial Intelligence in Medicine_ , pages 249–259. Springer, 2020.
  * Wang and Yoon (2021) Lin Wang and Kuk-Jin Yoon.  Knowledge distillation and student-teacher learning for visual intelligence: A review and new outlooks.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , page 1–1, 2021.  ISSN 1939-3539.  10.1109/tpami.2021.3055564.  URL http://dx.doi.org/10.1109/TPAMI.2021.3055564.
  * Xie et al. (2020) Qizhe Xie, Minh-Thang Luong, Eduard Hovy, and Quoc V. Le.  Self-training with noisy student improves imagenet classification, 2020\.
  * Yao and Gong (2020) Qi Yao and Xiaojin Gong.  Saliency guided self-attention network for weakly and semi-supervised semantic segmentation.  _IEEE Access_ , 8:14413–14423, 2020.
  * Young et al. (2019) Kyle Young, Gareth Booth, Becks Simpson, Reuben Dutton, and Sally Shrapnel.  Deep neural network or dermatologist?  _Lecture Notes in Computer Science_ , page 48–55, 2019.  ISSN 1611-3349.  10.1007/978-3-030-33850-3_6.  URL http://dx.doi.org/10.1007/978-3-030-33850-3_6.

##  Appendix A Methods for Generating Weak Pseudo-Labels

###  A.1 Grad-CAM

Grad-CAM is used to obtain saliency maps from the model predictions to
highlight the areas that the model focuses on during classification. It makes
use of the gradients of the output classes flowing into the last convolutional
layer to make low-resolution heatmaps ($3\times 3$ in case of ResNets).
Specifically, the gradients flowing back are global-average-pooled over the
width and height dimensions (indexed by i and j respectively) to obtain the
importance of the $k$th feature map for target class $c$, $\alpha^{c}_{k}$:

$\alpha^{c}_{k}=\frac{1}{Z}\sum_{i}\sum_{j}\frac{\partial y^{c}}{A^{k}_{ij}}$
(1)

Here, $y^{c}$ is the score (before softmax) of the class $c$, $Z$ is the
normalization factor, and $A^{k}$ is the $k$th feature map activation.

After this, a weighted sum of the final feature maps followed by a ReLU is
performed:

$L^{c}_{Grad-CAM}=ReLU\left(\sum_{k}\alpha^{c}_{k}A^{k}\right)$ (2)

A thresholding scheme is then used to convert the heatmap for each pathology
into a segmentation map to use as pseudo-labels. The probability threshold is
determined per pathology by maximizing the mIoU on the CheXpert train set.

###  A.2 Inter-Pixel Relation Network (IRNet)

We follow the method IRNet, which takes the previously generated CAMs and
tries to improve these seeds by training two output branches. The first branch
predicts a displacement vector field in which each pixel is represented by a
2D vector pointing to the centroid of the instance that the pixel is a part
of. This displacement field is converted to a class-agnostic instance map by
grouping together pixels whose vectors point to the same centroid. The second
branch is used to detect class boundaries by computing pairwise semantic
affinities, which is a confidence score for class equivalence between a pair
of pixels. The instance-wise CAMs obtained from the first branch are enhanced
by propagating their attention scores to relevant areas using the computed
affinities between neighboring pixels. Finally, pseudo-labels are generated
independently for each pathology with positive pixels being the ones with
higher class attention scores.

##  Appendix B Training Details

###  B.1 Classification Model

The pre-trained classification model used for generating the CAMs from image-
level labels is a DenseNet121. We use the Adam optimizer with default
$\beta$-parameters of $\beta$1 = 0.9, $\beta$2 = 0.999 and learning rate
$1\times 10^{-4}$ which is fixed for the duration of the training. Batches are
sampled using a fixed batch size of 16 images. We train for 3 epochs, saving
checkpoints every 4800 iterations.

###  B.2 Segmentation Model

The semantic segmentation model is trained using a class average dice loss and
Adam optimizer. We use a learning rate of 0.001 when training a small amount
of data and we decrease it to 0.0001 when training a large amount of data. We
train on up to four Nvidia GTX 1070s using a batch size of 8.

###  B.3 IRNet

The two branches of IRNet share the same ResNet50 backbone and are jointly
trained by minimizing the sum of three losses:

  1. 1.

Loss for displacement field prediction, which consists of two losses:

    1. (a)

L1 loss between the image coordinate displacement between a pair of nearby
foreground pixels $(i,j)$, denoted by $\hat{\delta}(i,j)=x_{j}-x_{i}$, and the
displacement obtained from their vector fields, $\mathcal{D}$, denoted by
$\delta(i,j)=\mathcal{D}(x_{j})-\mathcal{D}(x_{i})$:

$\mathcal{L}^{\mathcal{D}}_{\text{fg}}=\frac{1}{\mid\mathcal{P}^{+}_{\text{fg}}\mid}\sum_{(i,j)\in\mathcal{P}^{+}_{\text{fg}}}\left|\hat{\delta}(i,j)-\delta(i,j)\right|$
(3)

Here, $\mathcal{P}^{+}_{\text{fg}}$ refers to the set of neighboring
foreground pixel pairs with the same pseudo label (pixels with attention
scores larger than 0.3).

    2. (b)

Loss for background pixels as a normalized sum of their image coordinate
displacements:

$\mathcal{L}^{\mathcal{D}}_{\text{bg}}=\frac{1}{\mid\mathcal{P}^{+}_{\text{bg}}\mid}\sum_{(i,j)\in\mathcal{P}^{+}_{\text{bg}}}\left|\delta(i,j)\right|$
(4)

Here, $\mathcal{P}^{+}_{\text{bg}}$ refers to the set of neighboring
background pixel pairs with the same pseudo label (pixels with attention
scores less than 0.05).

  2. 2.

Loss for Class Boundary Detection, which makes use of the semantic affinities
between a pair of pixels $x_{i}$ and $x_{j}$, $a_{ij}$:

$a_{ij}=1-\text{max}_{k\in\Pi_{ij}}\mathcal{B}(x_{k})$ (5)

Here, $\Pi_{ij}$ is a set of pixels on the line between $x_{i}$ and $x_{j}$
and $\mathcal{B}\in\text{[0,1]}^{w\times h}$ is the output. Then, the loss is
the cross-entropy loss between the binary affinity label, with value 1 for the
same pseudo-class labels and 0 otherwise, and the predicted affinity of two
pixels:

$\mathcal{L}^{\mathcal{B}}=-\sum_{(i,j)\in\mathcal{P}^{+}_{\text{fg}}}\frac{\text{log}(a_{ij})}{2\mid\mathcal{P}^{+}_{\text{fg}}\mid}-\sum_{(i,j)\in\mathcal{P}^{+}_{\text{bg}}}\frac{\text{log}(a_{ij})}{2\mid\mathcal{P}^{+}_{\text{bg}}\mid}-\sum_{(i,j)\in\mathcal{P}^{-}}\frac{\text{log}(1-a_{ij})}{\mid\mathcal{P}^{-}\mid}$
(6)

Here, $\mathcal{P}^{-}$ represents the set of pixel pairs with different
pseudo labels.

The two branches are jointly trained by minimizing all three losses at the
same time:

$\mathcal{L}=\mathcal{L}^{\mathcal{D}}_{\text{fg}}+\mathcal{L}^{\mathcal{D}}_{\text{bg}}+\mathcal{L}^{\mathcal{B}}$
(7)

The model is trained with stochastic gradient descent using a learning rate of
0.1 with polynomial decay and a batch size of 16. The segmentation maps
created are then used in the semantic segmentation model as weak pseudo-
labels.

###  B.4 Distillation

In all experiments, we use the Adam optimizer, binary cross-entropy loss with
a sigmoid temperature $\tau=10$, and a learning rate of 0.001. For experiments
with pre-trained encoders, the teacher and student models are pre-trained with
the same encoders.

##  Appendix C Semi-Supervised Results

tab:weightings  Task Method p=0 p=0.2 p=0.4 p=0.6 p=0.8 p=0.85 p=0.9 p=1 Avg.
IoU CAM 0.163 0.269 0.274 0.273 0.290 0.292 0.299 0.263 IRNet 0.151 0.250
0.257 0.263 0.281 0.281 0.294 Airspace Opacity CAM 0.189 0.384 0.388 0.392
0.410 0.416 0.415 0.405 IRNet 0.262 0.379 0.378 0.387 0.407 0.408 0.413
Atelectasis CAM 0.176 0.319 0.313 0.308 0.332 0.327 0.324 0.328 IRNet 0.124
0.278 0.282 0.289 0.319 0.321 0.320 Cardiomegaly CAM 0.325 0.459 0.470 0.449
0.465 0.465 0.470 0.476 IRNet 0.249 0.420 0.431 0.456 0.462 0.459 0.471
Consolidation CAM 0.107 0.099 0.100 0.100 0.114 0.119 0.112 0.106 IRNet 0.0987
0.105 0.109 0.116 0.123 0.122 0.132 Edema CAM 0.176 0.239 0.253 0.246 0.260
0.263 0.258 0.248 IRNet 0.159 0.229 0.218 0.225 0.242 0.241 0.245 Enlarged
Cardiomediastinum CAM 0.165 0.522 0.521 0.522 0.538 0.539 0.547 0.549 IRNet
0.267 0.509 0.510 0.515 0.530 0.533 0.554 Lung Lesion CAM 0.059 0.072 0.063
0.059 0.066 0.075 0.135 0.003 IRNet 0.018 0.0449 0.0508 0.048 0.057 0.048
0.107 Pleural Effusion CAM 0.234 0.258 0.261 0.257 0.271 0.271 0.269 0.248
IRNet 0.190 0.207 0.227 0.231 0.248 0.247 0.252 Pneumothorax CAM 0.052 0.103
0.107 0.103 0.153 0.154 0.159 0.003 IRNet 0.014 0.089 0.117 0.116 0.165 0.163
0.167 Support Devices CAM 0.148 0.234 0.264 0.274 0.290 0.287 0.297 0.267
IRNet 0.127 0.243 0.244 0.252 0.261 0.266 0.278

Table 1: IoU scores for semi-supervised segmentation using IRNet and CAMs by
weighting of pixel-level and weak pseudo-labels. p represents the probability
of picking a pixel-level labeled training example in the current batch.

##  Appendix D Weakly-Supervised Results

tab:weakresults  Method Encoder Architecture Train Set Size Test mIoU CAM
ResNet18 100 $0.074\pm 0.00783$ 2000 $0.126\pm 0.00275$ 20000 $0.136\pm
0.00990$ ResNet18-ImageNet 100 $0.082\pm 0.00396$ 2000 $0.132\pm 0.00687$
20000 $0.139\pm 0.00095$ ResNet18-MoCo-CXR 100 $0.084\pm 0.00238$ 2000
$0.133\pm 0.00002$ 20000 $0.139\pm 0.00076$ ResNet18-CheXpert 100 $0.129\pm
0.00036$ 2000 $0.141\pm 0.00003$ 20000 $0.142\pm 0.00095$ IRNet ResNet18 100
$0.076\pm 0.00133$ 2000 $0.105\pm 0.00292$ 20000 $0.111\pm 0.00103$
ResNet18-ImageNet 100 $0.081\pm 0.00245$ 2000 $0.117\pm 0.00576$ 20000
$0.124\pm 0.00001$ ResNet18-MoCo-CXR 100 $0.082\pm 0.00007$ 2000 $0.120\pm
0.00010$ 20000 $0.128\pm 0.00001$ ResNet18-CheXpert 100 $0.125\pm 0.00031$
2000 $0.146\pm 0.00002$ 20000 $0.156\pm 0.00000$

Table 2: IoU scores of weakly-supervised segmentation models using either CAM
or IRNet pseudo-labels and varying encoder architectures and train set sizes.
Confidence intervals are calculated using a $p$-value of 0.95