# Branches of a Tree: Taking Derivatives of Programs with
Discrete and Branching Randomness in High Energy Physics

Michael Kagan  Corresponding authors contributed equally to this work:
makagan<EMAIL_ADDRESS>SLAC National Accelerator
Laboratory  Lukas Heinrich  Corresponding authors contributed equally to this
work:
makagan<EMAIL_ADDRESS>Technical University of Munich

###### Abstract

We propose to apply several gradient estimation techniques to enable the
differentiation of programs with discrete randomness in High Energy Physics.
Such programs are common in High Energy Physics due to the presence of
branching processes and clustering-based analysis. Thus differentiating such
programs can open the way for gradient based optimization in the context of
detector design optimization, simulator tuning, or data analysis and
reconstruction optimization. We discuss several possible gradient estimation
strategies, including the recent Stochastic AD method, and compare them in
simplified detector design experiments. In doing so we develop, to the best of
our knowledge, the first fully differentiable branching program.

††preprint: APS/123-QED

##  I Introduction

Gradient-based optimization methods are at the core of many modern successes
in Machine Learning (ML) and Artificial Intelligence (AI), especially Deep
Learning. The development and application of these ML methods in High Energy
Physics (HEP) has similarly enabled large performance improvements in a wide
array of tasks, such as particle reconstruction, fast simulation, and
parameter inference (for recent topical reviews, see e.g. Radovic  _et al._
(2018); K. Cranmer, U. Seljak, K. Terao, and the Particle Data Group (2022);
Calafiura  _et al._ (2022); Cranmer  _et al._ (2020); Shlomi  _et al._ (2020);
Thais  _et al._ (2022); Butter  _et al._ (2023); Adelmann  _et al._ (2022)).
Gradient-based optimization methods rely on automatic differentiation (AD)
Bücker  _et al._ (2005); Baydin  _et al._ (2018), an algorithmic way to
efficiently compute the derivatives of numerical software. AD is a general
tool that can be applied to scientific software beyond ML, such as simulators
and inference algorithms, and used for optimizing the parameters of this
software.

The broader application of AD to numerical software, potentially mixed with ML
components, is often referred to as Differentiable Programming (DP). For
instance, combining AD-enabled HEP software with ML can lead to optimizable
hybrid physics-AI models with built-in physics knowledge from the HEP
software, such as AI-augmented / AI-guided simulation and reconstruction, or
analysis-by-synthesis inference methods with simulators in the loop Heinrich
and Kagan (2023); Simpson and Heinrich (2023); Cheong  _et al._ (2022). Using
such AD-integrated HEP software in ML models can be considered a complimentary
approach to adding inductive bias into ML models through structure and
architecture, such as symmetry equivariance and relational inductive bias.

While interest is quickly growing to apply gradient-based optimization methods
to a broader set challenges in HEP, such as detector design or end-to-end
reconstruction, a major limitation has been the fact that standard AD can only
compute derivatives of deterministic continuous functions or stochastic
functions with reparametrized continuous random variables Glasserman (2013);
Kingma and Welling (2014); Rezende  _et al._ (2014). Specifically, in HEP,
many programs are both stochastic and rely on sampling _discrete_ random
variables or decisions, such as branching points in parton showers, particle-
material interactions, or clustering steps in jet building. Standard AD may
not compute the desired derivative of such programs correctly, particularly
when the discrete stochasticity depends on the parameter one aims to optimize
(and thus differentiate with respect to). Instead, more careful consideration
on how to compute the appropriate derivative is required.

There are several methods for gradient-based optimization in programs with
discrete randomness, which we explore within the context of HEP applications.
One method uses the score-based approach to estimating derivatives of
expectations values Williams (1992), which has been examined sparsely within
HEP and not for tasks such as detector design optimization. Recently, Arya et
al. Arya  _et al._ (2022) proposed a new AD method for handling and composing
programs with discrete randomness. Using these tools, we develop simplified
differentiable HEP simulators,which nonetheless exhibit the critical behaviors
which until now hampered progress, and case studies to examine the behavior of
and the variance of these gradient estimators.

A review of methods for computing derivatives of stochastic programs is found
in Sec. II. Related work is discussed in Sec. III. Sec. IV presents
applications and comparisons of different gradient estimators in HEP case
studies, with an emphasis on detector design optimization.

Contributions111Project code can be found at
https://github.com/lukasheinrich/branches_of_a_tree/: We introduce several
methods to enable differentiation through the discrete randomness to HEP
programs. We show how the score function based approach can be used for design
optimization in HEP. While the score function is often used outside of HEP for
design optimization, and has been used for other tasks within HEP, it has not
yet been explored within the context of HEP detector design optimization. We
also introduce stochastic AD Arya  _et al._ (2022) to HEP and its ability to
enable differentiable programming even in programs with discrete randomness.
We provide the first application of stochastic AD to branching processes and
in doing so we develop the first (to the best of our knowledge) differentiable
branching program. To test these methods, we provide comparisons of gradient
methods on detector design optimization case studies.

##  II Review of Differentiation of Stochastic Programs

⬇ def f(|$\theta$|):  p = 0.5  b |$\sim$| Bernoulli(p)  return g(|$\theta$|) + b width .8pt |  ⬇ def f(|$\theta$|):  p = h(|$\theta$|)  b |$\sim$| Bernoulli(p)  return g(|$\theta$|) + b
---|---
Figure 1: Assuming differentiable $g:\mathbb{R}\to\mathbb{R}$ and
$h:\mathbb{R}\to(0,1)$, Left: Toy program without $\theta$-dependence in the
discrete stochastic behavior. As such, the derivative of the expected value of
this program is $\frac{dE[f(\theta)]}{d\theta}=\frac{dg(\theta)}{d\theta}$ and
standard AD can correctly estimate this. Right: Toy program with
$\theta$-dependence in the discrete stochastic behavior through the Bernoulli
parameter $p$. Standard AD would ignore this dependence, and the resulting
derivative estimator would be the same as the program on the left. However the
correct derivative should be
$\frac{d\mathbb{E}[f(\theta)]}{d\theta}=\frac{dg(\theta)}{d\theta}+\frac{dh(\theta)}{d\theta}$.

In many HEP applications, a quantity of interest can be formulated as an
expectated value of a function $f(x,\theta)$ over a parametrized density
$p_{\theta}(x)$: $\bar{f}(\theta)=\mathbb{E}_{p_{\theta}(x)}[f(x,\theta)]$.
For optimizing such quantities one requires the gradients of these expectation
values of stochastic programs, e.g.
$\frac{d}{d\theta}\mathbb{E}_{p_{\theta}(x)}[f(x,\theta)]$. Importantly, the
expected value of such stochastic programs may be continuous and
differentiable, even when they depend on discrete randomness. For instance,
the expected value of a Bernoulli random variable $b\sim$Bernoulli($\theta$)
has the derivative $\frac{d}{d\theta}\mathbb{E}[b]=\frac{d}{d\theta}\theta=1$.
However, standard AD tools applied to such expectations may not produce the
desired result. For instance, Figure 1 shows two programs with discrete
stochasticity. On the left, the discrete stochasticity does not depend on the
parameter of differentiation $\theta$, and standard AD will produce the
correct derivative. On the right the discrete stochasticity depends on
$\theta$, standard AD will ignore this dependence and the resulting derivative
will be incorrect as it will ignore this dependence. Handling these challenges
requires more dedicated consideration on home to compute the appropriate
derivative. We briefly review several approaches to gradient estimation below
(see Ref. Mohamed  _et al._ (2020) for a detail review).

Finite Differences (Numerical Differentiation): Finite difference methods
estimate derivatives by computing the difference between forward evaluations
of a program and a perturbed version of the program, for instance:

$\frac{d}{d\theta}\mathbb{E}_{p_{\theta}(x)}[f(x)]\approx\frac{\mathbb{E}_{p_{\theta+\epsilon}(x)}[f(x)]-\mathbb{E}_{p_{\theta}(x)}[f(x)]}{\epsilon}$
(1)

Finite difference methods are prone to high variance Baydin  _et al._ (2018),
and require large numbers of program evaluations when $\theta$ is high
dimensional. Central difference methods can reduce error. One large
contributor to this large variance is that multiple independent evaluations of
the program are used to estimate this gradient, introducing separate
stochastic evaluation paths of the program.

Reparameterization Trick: In many cases, when sampling $x\sim p_{\theta}(x)$,
we can smoothly reparametrize this sampling as $z\sim p(z)$ and
$x=g(z,\theta)$, where $p(z)$ is often a simple “base" distribution and
$g(\cdot,\theta)$ provides a differentiable, $\theta$-dependent transformation
of samples from the base to the desired distribution. For example, the normal
distribution $x\sim\mathcal{N}(\mu,\sigma)$ is location-scale
reparameterizable through $z\sim\mathcal{N}(0,1)$ and $x=\sigma
z+\mu\sim\mathcal{N}(\mu,\sigma)$. This is particularly convenient for
computing derivatives of expectation values of differentiable functions
$f(\cdot)$,

$\displaystyle\frac{d}{d\theta}\mathbb{E}_{p_{\theta}(x)}[f(x)]$
$\displaystyle=$ $\displaystyle\frac{d}{d\theta}\int p(z)f(g(z,\theta))dz$ (2)
$\displaystyle=$ $\displaystyle\int
p(z)\frac{df}{dg}\frac{dg(z,\theta)}{d\theta}dz$

Many HEP simulators, which implement a structural causal model, can be
considered as a reparametrization, mapping from noise variables to random
variables $x$ with physical meaning. However if the random variables $x$ are
discrete, the map is not differentiable, which limits the applicability of the
reparametrization trick within a HEP context.

Surrogate Methods: When a reparameterization is not possible, either because
$f(\cdot)$ is non differentiable or $p_{\theta}(x)$ does not admit a smooth
reparameterization, surrogate methods can be used. In this case, an ML model
$S(z,\theta)$, with $z\sim p(z)$ a chosen distribution, is trained to mimic
the stochastic program. As such, surrogate methods try enable
reparameterization though a ML model and thus enable differentiation, for
instance:

$\displaystyle\frac{d}{d\theta}\mathbb{E}_{p_{\theta}(x)}[f(x)]$
$\displaystyle\approx$ $\displaystyle\frac{d}{d\theta}\int p(z)S(z,\theta)dz$
(3) $\displaystyle=$ $\displaystyle\int p(z)\frac{dS(z,\theta)}{d\theta}dz$

The quality of this derivative estimator will depend significantly on the
quality of the surrogate model as an approximation of the original program.
Moreover, the surrogate will learn a smooth approximation of non-
differentiable elements of the program, but how this approximation is
performed and if bias is introduced is difficult to asses.

Score function: When the parameter dependence of a differentiable distribution
$p_{\theta}(x)$ is known and differentiable with respect to the parameters,
one can compute:

$\displaystyle\frac{d}{d\theta}\mathbb{E}_{p_{\theta}(x)}[f(x)]$
$\displaystyle=$ $\displaystyle\int p_{\theta}(x)\frac{d\log
p_{\theta}(x)}{d\theta}f(x)dx$ (4) $\displaystyle=$
$\displaystyle\mathbb{E}_{p_{\theta}(x)}\Big{[}\frac{d\log
p_{\theta}(x)}{d\theta}f(x)\Big{]}$

where $\frac{d}{d\theta}\log p_{\theta}(x)$ is known as the score function.
This gradient estimator is also known as REINFORCE Williams (1992), and is
commonly used in reinforcement learning. The benefit of this approach is that
the $f(x)$ does not need to be differentiable, only $p_{\theta}(x)$ must be
differentiable with respect to $\theta$. As such, discrete random variables
can be used in the computation of $f(x)$.

Control Variates: Score function based gradient estimates often have a large
variance, which can make tasks like optimization with gradient descent slow
and more difficult. A control variate, $c(x,\theta)$, can be subtracted from
$f(x)$ in Eqn. 4 to reduce the variance of the estimator as long as it does
not bias the estimator, i.e. as long as
$\mathbb{E}_{p_{\theta}(x)}\Big{[}\frac{d\log
p_{\theta}(x)}{d\theta}c(x,\theta)\Big{]}=0$. Noting that
$\mathbb{E}_{p_{\theta}(x)}\Big{[}\frac{d\log
p_{\theta}(x)}{d\theta}\Big{]}=0$, one way to find a control variate is to
choose a $c(\theta)$ which does not depend on $x$. A common control variate,
often also called a baseline, is the function mean $\bar{f}_{\theta}=\int
p_{\theta}(x)f(x)$; in practice $\bar{f}_{\theta}$ is often estimated using
the mean of a mini-batch. We will use this baseline for the experiments in
Sec. IV. More details on variance reduction methods for score based gradient
estimation can be found in Ref. Greensmith  _et al._ (2004).

Proposal Distributions: In some case, we may not know or have access to
$p_{\theta}(x)$, for example when $g(\theta)$ is a simulator with parameters
as input and sampling is done internally within the program. One approach can
be to imagine the input to the program as a sample from a proposal
distribution $\theta\sim\pi_{\psi}(\theta)$, where $\psi$ are the parameters
of the proposal distribution. For instance one could choose a normal
distribution $N(\psi,1)$ for the proposal. One would then aim to optimize the
mean of this proposal,

$\frac{d}{d\psi}\mathbb{E}_{\pi_{\psi}(\theta)}[g(\theta)]=\mathbb{E}_{\pi_{\psi}}\Big{[}\frac{d\log\pi_{\psi}(\theta)}{d\psi}g(\theta)\Big{]}$
(5)

Stochastic AD: The stochastic derivative of the expected value of a function
$f(\cdot)$ of a discrete random variable $x\sim p_{\theta}(x)$ has the form
Arya  _et al._ (2022):

$\frac{d}{d\theta}\mathbb{E}_{p_{\theta}(x)}[f(x)]=\mathbb{E}_{p_{\theta}(x,y)}[\delta+\beta\big{(}f(y)-f(x)\big{)}]$
(6)

where $\delta$ is the standard derivative $\partial f/\partial\theta$ as
computed with AD, and the second term corresponds to the effect of a change in
$\theta$ on the sampling of the discrete random variable $x$. The weight
$\beta$ depends on the underlying sampling distribution and $y$ is an
alternative value of the discrete random variable. Conceptually, for discrete
$X$ with consecutive integer range one can understand this result through the
lens of reparameterization. One can reparamaterize the discrete distribution
via the inversion method, e.g.:

$\displaystyle\omega$ $\displaystyle\sim$ $\displaystyle U[0,1]$
$\displaystyle x$ $\displaystyle=$ $\displaystyle\left\\{X\
:\omega\in[CDF_{p_{\theta}}(X),CDF_{p_{\theta}}(X+1)\ )\ \right\\}$

For example, if $x$ is a Bernoulli random variable with parameter $\theta$,
then one can reparameterize $x=H(\omega>1-\theta)$ where $H(\cdot)$ is the
Heaviside step function. The boundaries which define the set of values of
$\omega$ that result in a value of $X$ are now dependent on the parameters
$\theta$. A change in parameters changes the boundaries, and thus changes the
probabilities of different outcomes $x$. The second term of the stochastic AD
derivative accounts for the infinitesimal change in probabilities as the
boundaries are changed as well as the alternative value of the program $y$
that would result from such a change. Importantly, one can define this
derivative at each program evaluation and at each stochastic sampling within
the program, allowing for the development of composition rules and of forward
mode automatic differentiation. For a more detailed discussion of Stochastic
AD, see Ref. Arya  _et al._ (2022).

The expectation on the right hand side of Equation 6 is taken with respect to
the joint distribution $p_{\theta}(x,y)$, which is often denoted the coupling.
The marginal distributions of this coupling must be the same as the original
sampling distribution, e.g. $\int dyp_{\theta}(x,y)=\int
dxp_{\theta}(x,y)=p_{\theta}(x)$ to ensure both the primal evaluation of the
program and the alternative proceed under the normal operation of the program.
As such, there is a distribution over possible alternative programs. In
practice, out of all possible alternatives from all of the discrete samplings
within a program, only one alternative is randomly chosen. This pruning
process treats the set of alternatives as a categorical distribution with the
probability of a given alternative being the weight of the alternative
relative to the total event weight (computed using the composition rules).
This alternative may occur in the middle of the program, and is then tracked
in parallel to the primal until completion of the alternative program. One can
then use this program alternative $y$ for computing even gradients using
Equation 6.

Variance Reduction with Random Number Reuse: While the marginals of the
coupling are fixed, one has considerable flexibility to choose the correlation
structure between $x$ and $y$. This is important because once an alternative
is determined from a discrete sampling within the program, the alternative
program will be run to completion and thus may require additional sampling of
discrete random variables. The more correlated are the evaluations of the
alternative completion to that of the primal evaluation, the lower the
variance of the gradient estimator may be. As downstream discrete samplings
also occur in the primal program, one can reuse the reparameterized sampling
in the primal program, i.e. the $\omega$ values sampled for the inversion
method. By reusing $\omega$ values, less additional variance is added to the
alternative program then if the downstream discrete random variable were
sampled independently. In this work for the experiments in Sec. IV, we use a
first-in first-out (FIFO) approach, where at each time step we store $\omega$
values in the FIFO while iterating over branches (particle) in the primal, and
then pull $\omega$ values from the FIFO while iterating over branches in the
alternative. If additional $\omega$ values are needed by the alternative, they
are sampled independently of the primal.

##  III Related Work

Automatic differentiation Bücker  _et al._ (2005); Baydin  _et al._ (2018),
and its use in gradient based optimization, is ubiquitous in ML, statistics,
and applied math. AD uses the chain rule to evaluate derivatives of a function
that is represented as a computer program. AD takes as input program code and
produces new code for evaluating the program and derivatives. AD typically
builds a computational graph, or a directed acyclic graph of mathematical
operations applied to an input. Gradients are defined for each operation of
the graph, and the total gradient can be evaluated from input to output,
called forward mode, and from output to input, called reverse mode or
backpropagation in ML. AD is backbone of ML / DP frameworks like TensorFlow
Abadi  _et al._ (2015), JAX Bradbury  _et al._ (2018), and PyTorch Paszke  _et
al._ (2019).

Significant work has been performed on developing Monte Carlo estimators for
gradients in machine learning, as discussed in the recent review Mohamed  _et
al._ (2020), and gradients of stochastic computation graphs Schulman  _et al._
(2015). More recently methods such as Stochastic AD Arya  _et al._ (2022,
2023) have been developed to target derivatives of programs with discrete
stochastic behavior in a compositional way, as well as developing specific
applications with dedicated variance reduction methods (e.g. developing
coupling for these applications). Extensions to AD have recently been proposed
for differentiating the expectations of stochastic programs Lew  _et al._
(2023); Krieken  _et al._ (2021), and to account for parametric
discontinuities Bangaru  _et al._ (2021).

Differentiable programming approaches have begun to be explored in HEP.
Examples include histogram fitting in pyHF Heinrich  _et al._ ; Heinrich  _et
al._ (2021), in analysis optimization in Neos Simpson and Heinrich (2023);
Heinrich and Simpson (2020), for modeling parton distribution functions (pdf)
used by matrix element generators Carrazza  _et al._ (2021); Ball  _et al._
(2021), and for developing differentiable matrix element simulations with
MadJax Heinrich and Kagan (2023). Related to our work, Ref. Nachman and
Prestel (2022) studies differentiating a parton shower, which focuses on the
derivative of the primal shower program but does not examine differentiation
through discrete randomness in such branching programs. Within cosmology, the
Differentiable Universe Initiative is developing a differentiable simulation
and analysis toolchain for comological data analyses Differentiable Universe
Initiative .

Within HEP, score function gradient estimators have been used within the
context of jet grooming Carrazza and Dreyer (2019) and hierarchical jet
clustering Brehmer  _et al._ (2020). To the best of our knowledge, this work
is the first application to detector design optimization in HEP.

On HEP detector design, surrogate based optimization methods have been
developed and explored for particle detectors in Ref. Shirobokov  _et al._
(2020). Surrogate methods have also been applied to neutrino detector design
Haack and Schumacher . Detector design optimization with standard AD tools and
with surrogate methods is discussed in Ref. Dorigo  _et al._ (2023). A branch-
and-bound type algorithm is explored in Ref. Gorordo  _et al._ (2022).

(a) Single particle energy loss

(b) Particle shower with splitting

Figure 2: Event Displays of a single particle energy loss (left) and a
particle shower with particle splitting (right). A primal event is shown in
purple, while an alternative event, as determined using Stochastic AD, is
shown in yellow. The material map is shown in grey.

##  IV Applications

We present a series of applications in a simplified simulation of particles
interacting with detector material. In these experiments, we examine how
different gradient estimation methods can be applied and compare their
performance in terms of the variance of the estimators.

###  IV.1 Particle Interaction Simulator

The simplified simulator in this work aims to capture the salient features of
particle physics simulators that model the traversal of particles through
matter but be simple enough allow reimplementation in a programming language
of choice for a detailed study of various gradient estimation in a self-
contained setting.

[ht!] Simplified Particle Interaction Simulator

1:$E_{0}$ : energy threshold

2:$\epsilon$ : energy loss at interaction

3:$m(x,\theta):\mathbb{R}^{3}\times\mathbb{R}^{n}\to[0,1]$ : material map

4:$P=\\{p_{i}=(\vec{x}_{i}\vec{p}_{i},E_{i})\\}$: initial particle list

5:$H=\emptyset$: list of hits

6:function Simulate($P$,$\theta$)

7: while not all $p\in P$ below threshold do

8: $P\leftarrow\emptyset$ $\triangleright$ list of surviving particles

9: for for all particles in $P$ do

10: if $E_{i}<E_{0}$ then

11: continue $\triangleright$ particle below thr.

12: end if

13: $p_{i}=\mathrm{propagate}(p_{i})$

14: $\rho_{\mathrm{E}}\leftarrow$ `sample`
$\mathrm{Ber}(\rho_{\mathrm{E}}|m_{\theta}(x_{i}))$

15: if $\rho_{E}$ then $\triangleright$ particle interacts

16: $H\leftarrow H\cup\\{x_{i}\\}$ $\triangleright$ add position to hits

17: $\rho_{\mathrm{split}}\leftarrow$ `sample`
$\mathrm{Ber}(\rho_{\mathrm{split}}|m_{\theta}(x_{i}))$

18: if $\rho_{\mathrm{split}}$ then $\triangleright$ particle splits

19: $p_{L},p_{R}\leftarrow\mathrm{split}(p_{i})$

20: $P\leftarrow P\cup\\{p_{R},p_{L}\\}$

21: else

22: $E_{i}\leftarrow E_{i}-\epsilon$ $\triangleright$ lose energy

23: $P\leftarrow P\cup\\{p_{i}\\}$

24: end if

25: else

26: $P\leftarrow P\cup\\{p_{i}\\}$

27: end if

28: end for

29: end while

30: return $H$

31:end function

The simulator $S_{\theta}(p)$ models the stochastic evolution of particles
through two main processes. A binary _splitting process_ $p_{0}\to
p_{L},p_{R}$ that splits a parent particle momentum evenly across two child
particles and a _energy loss process_ $E\to E-\epsilon$. The probability of a
particle interaction is modeled as a function of the material map $m(x)$. The
simulation is performed by fixed time steps, first propagating particles in
their direction of travel, and then querying the material map to determine if
an interaction occurs, and if so which type of interaction (i.e. splitting or
energy loss). Pseudo-code for the simulator can be found in Alg. IV.1.

The detector is simulated as a continuous material map
$m_{\theta}(x_{0},x_{1})$ which takes as input a position $(x_{0},x_{1})$ and
outputs an interaction probability. This interaction probability is dependent
on detector parameters $\theta$, and we will examine examples where
derivatives with respect to $\theta$ are sought. Only a single detector
parameter is used in the following experiments, which is the detector inner
radius which will be denoted $\theta_{R}$. With
$r=\sqrt{x_{0}^{2}+x_{1}^{2}}$, and $\phi=\arctan{\frac{x_{0}}{x_{1}}}$, the
material map is defined as:

$m_{\theta}(x,y)=\frac{1}{2}m_{start}(r,\theta_{R})\ m_{\Phi}(r,\phi)\
m_{R}(r)\ m_{end}(r,\theta_{R})$ (7)

where

$\displaystyle m_{start}(r,\theta_{R})$ $\displaystyle=$
$\displaystyle\frac{1}{1+e^{-\beta(r-\theta_{R})}}$ $\displaystyle
m_{\Phi}(r,\phi)$ $\displaystyle=$
$\displaystyle\frac{1}{1+e^{\beta\sin(\omega(\phi+2r))}}$ $\displaystyle
m_{R}(r)$ $\displaystyle=$
$\displaystyle\frac{1}{1+e^{\beta\cos(\omega(r-2))}}$ $\displaystyle
m_{end}(r,\theta_{R})$ $\displaystyle=$
$\displaystyle\frac{1}{1+e^{\beta(r-\theta_{R}-R_{max})}}\ .$

The terms $m_{start}(r,\theta_{R})$ and $m_{end}(r,\theta_{R})$ determine the
inner and outer radius of the detector, respectively. The terms
$m_{\Phi}(r,\phi)$ and $m_{R}(r)$ determine the segmentation in $\phi$ and
$r$, respectively. The constants $\beta$, $\omega$, and $R_{max}$ control the
sharpness of the smooth material map, the segmentation frequency in the
azimuthal direction, and the maximum depth of the detector, respectively. The
parameter that will be optimized is $\theta_{R}$, the inner radius of the
detector.

Example material maps can be seen in grey in the event displays of Figure 2,
where the darkness of the shade of grey indicates the strength of interaction.

###  IV.2 Single Particle Energy Loss

Figure 3: For simulations of single particle energy loss: (Left) The loss
function and various gradient estimators of the loss are shown as a function
of the detector radius parameter. Sample primal evaluations of the loss are
shown as markers, and the interquantile interval is shown in black. The red
dashed line shows the derivative of an polynomial interpolation of the mean
loss. (Middle) The mean and standard deviation of the numeric, Stochastic AD
(STAD), score function (SCORE), and score function with baseline (SCORB)
gradient estimators as a function of the detector radius parameter. The
gradient of the polynomial interpolation of the mean loss is shown in dashed
black. (Right) Box plot of the four variance estimators evaluated at parameter
value $\theta_{R}=2.5$m, with the mean shown as a dashed line.

In the single particle energy loss setting, interactions which cause splitting
are turned off, i.e. there are no showers. At each step, the particle
interacts with the detector with a probability
$p_{\mathrm{Eloss}}=m(x,\theta)$ which is dependent on material map parameters
$\theta$. This probability is large is high density regions of detector
material and small in low density regions. A Bernoulli distribution with
parameter $m(x,\theta)$ is sampled at each time step to determine if the
interaction occurs and, if so, the particle deterministically loses energy
$\epsilon=1$ GeV. All particles are set to have initial energy of 25 GeV and
when the particle energy falls below $E_{0}=0.5~{}\mathrm{GeV}$, the particle
is stopped. An example event display can be seen in Fig. 2(a), where the
primal particle trajectory is seen in purple, the alternative trajectory
determined with Stochastic AD is seen in yellow, and the material map is seen
in grey.

In this example, there is only one detector parameter $\theta\equiv R$, the
inner radius of the detector, and derivatives are computed with respect to $R$
using several methods of estimating gradients. The mean squared error between
the radial position of points of particle interaction and a target radius
$\bar{R}_{T}$ is used as a loss function that one may minimize for the
purposes of design optimization. In this example we set $\bar{R}_{T}=2$m.

The loss as a function of the detector radius parameter can be seen on the
left in Fig. 3. The loss from individual primal simulation samples can be seen
in grey, the median loss and interquartile range in black, and a polynomial
interpolation of the average loss. Even though the simulation is stochastic,
and one can see the variations of the loss at each parameter in the grey
points, there is a clear minimum to the loss function.

The gradient estimators and their standard deviations, calculated over the
5000 simulation runs, can be seen as a function of the detector radius
parameter in the middle in Fig. 3. The distributions of gradient estimators
evaluated at the parameter value $\theta_{R}=2.5$m can be seen in the box plot
on the right in Fig. 3. As expected, numerical derivatives have the largest
standard deviation, though it should be noted that this can depend highly on
the size of the finite different $\epsilon$ and the method for calculating the
numerical derivative. Similarly, score function gradient estimators without
baseline shows a high standard deviation, especially at large radius parameter
where the loss function has larger standard deviation over different
simulation samples. The score with baseline has a much reduced standard
deviation over the score function without baseline across all parameter
values. Stochastic AD shows the smallest standard deviation of all estimators,
likely owing to the ability to couple much of the alternative program
evaluation to that of the primal up to the alternative branching point. Across
all gradient estimators, the mean of the gradient estimator, shown in solid
lines, are close to the gradient of the polynomial fit of the loss (which
serves as a rough guide to the gradient of the expected loss) within one
standard deviation. A numerical comparison of gradient estimator mean and
variance can be found in Tab. 1.

###  IV.3 Branching Shower

Figure 4: For simulations of particle showers with splitting: (Left) The loss
function and various gradient estimators of the loss are shown as a function
of the detector radius parameter. Sample primal evaluations of the loss are
shown as markers, and the interquantile interval is shown in black. The red
dashed line shows the derivative of an polynomial interpolation of the mean
loss. (Middle) The mean and standard deviation of the numeric, Stochastic AD
(STAD), score function (SCORE), and score function with baseline (SCORB)
gradient estimators as a function of the detector radius parameter. The
gradient of the polynomial interpolation of the mean loss is shown in dashed
black. (Right) Box plot of the four variance estimators evaluated at parameter
value $\theta_{R}=2.5$m, with the mean shown as a dashed line.

In the branching shower example, the same material map, and thus interaction
probability, as the single particle energy loss simulation is used but if an
interaction occurs, the particle is deterministically split into two daughter
particles each with half the energy of the parent particle and with an opening
angle of 0.1 radians. Initial particles are set to have starting energy of 25
GeV and when any particle energy falls below 0.5 GeV, the particle is stopped.
The same loss function as in the single particle energy loss example is used
here. An example event display can be seen in Fig. 2(b), where the primal
particle shower is seen in purple, the alternative shower determined with
Stochastic AD is seen in yellow, and the material map is seen in grey.

The loss function and the standard deviation of the gradients, as functions of
the detector radius parameter, can be seen on the left and middle,
respectively, in Fig. 4. As in the single particle energy loss example, the
numerical gradients are found to have the largest standard deviation of
gradients, with the score function without baseline estimator having the
second largest standard deviation. Notably, the score function with baseline
estimator is found to have the smallest standard deviation, slightly smaller
than the Stochastic AD gradient estimator. Unlike the single particle example,
the splitting shower has many program branching points which can create
alternative outputs that are significantly different from the primal shower.
In turn, this leads to a reduction in the correlation between the primal and
alternative showers and ultimately to an increase in the gradient estimator
standard deviation. A comparison of the distribution of gradient estimators,
at detector parameter value $\theta_{R}=2.5$m, can be seen on the right in
Fig. 4. While the mean values (dotted lines) in each box agree well across
estimators, the variance estimates as well as the tails are significantly
better behaved for Stochastic AD and score function with baseline estimators.
Similarly, a comparison of the mean and standard deviation of the gradient
estimators evaluated at the parameter value $\theta_{R}=2.5$m for both the
single particle energy loss and the splitting shower can be found in Tab. 1.

Estimator |  $E$-loss | Shower
---|---|---
StochAD | $\mathbf{3.17\pm 4.47}$ | $2.53\pm 6.37$
Score w/ Baseline | $3.01\pm 6.59$ | $\mathbf{2.47\pm 4.42}$
Score w/o Baseline | $2.68\pm 17.18$ | $2.76\pm 12.20$
Numerical | $3.83\pm 139.96$ | $2.43\pm 74.85$
Table 1: Gradient estimator mean and standard deviation, for both the single
particle energy loss and splitting shower, evaluated at parameter value
$\theta_{R}=2.5$m and determined from 5,000 samples. The estimator with lowest
standard deviation is shown in bold.

It should be noted that there is considerable flexibility in Stochastic AD for
how to couple the randomness in the primal and alternative programs after the
point at which the alternative is produced, i.e. how to choose the join
distribution over random variables in the primal and alternative programs.
This selection of coupling can have a considerable impact on the Stochastic AD
gradient estimator variance. In this work, we have used a simple approach of
re-using random variables sampled in the primal for the alternative, without
regard for where those random variables are re-used in the alternative. We
have seen that this re-use can have a large impact; we observed that removing
the re-use of random variables in the alternative can increase the Stochastic
AD gradient estimator standard deviation by factors of 1.5 or more. More
generally, a more careful strategy of re-using of random variables may
considerably reduce the Stochastic AD gradient estimator variance.

###  IV.4 Design Optimization with Splitting Shower

We test the ability to use the various gradient estimators to perform a
gradient based optimization using of the detector radius parameter using the
aforementioned loss with a target radial shower depth of $\bar{R}_{T}=2$. Each
epoch consists of a single step of the optimization, with a mini-batch size of
only 2 simulation runs used to estimate gradients in each epoch. The Adam
optimizer Kingma and Ba (2015) is used. A learning rate of 0.01 is used for
the gradient descent parameter update. For all optimizations, the initial
detector radius parameter value is set to $\theta_{init}=3$m and the
optimization is run for 500 gradient steps. Each gradient method is used in 10
separate optimizations, and the average and standard deviation of the loss at
each optimization step is shown in Fig. 5. As expected, the score function
estimator without baseline and the numeric gradients shows the largest
standard deviation to the extent that optimization is not feasible in this
setting. We also see that the numeric and score function without baseline
estimators are significantly slower at optimizing the objective. The score
with baseline estimator and Stochastic AD estimators show similar standard
deviation and similar progress towards the loss minimum as a function of
optimization step. This suggests that Stochastic AD and score function with
baseline estimators provide significantly better gradient estimates, even with
very small sample sizes, and are likely interesting estimators for further
study of detector design optimizations in more complex settings.

Figure 5: The mean and inter-quantile range of the loss versus epoch of
detector design optimization is shown. Mean and quantiles are computed from 10
optimizations.

##  V Conclusion

In this work, we discuss several strategies for differentiating stochastic
programs, with a focus on methods capable of differentiating programs with
discrete randomness, and discuss their application to High Energy Physics
detector design optimization. We develop the first application of Stochastic
AD to branching processes and, more generally, the first differentiable
branching program capable of estimating gradients through the discrete
processes within a particle shower. We also introduce score function gradient
estimators within this HEP detector design context. We find that Stochastic AD
and score function gradient estimators, using control variates, provide the
best gradient estimators in terms of smallest standard deviation among the
gradient estimators examined within a case study of detector design. We show
that both techniques can successfully be used for gradient-based HEP detector
design on a toy detector simulator.

More broadly, we believe that the careful study and application of techniques
like Stochastic AD and score function estimation can open the way to a wide
array of new differentiable programming applications in HEP and other
sciences.

## Acknowledgements

We thank Gaurav Arya, Frank Schäfer, and Moritz Schauer for the helpful
discussions regarding Stochastic AD, and thank Gaurav Arya for the helpful
feedback on the manuscript. We thank Michael Brenner for the helpful
discussions regarding score function gradient estimators at the Aspen Center
for Physics, as this work was partially performed at the Aspen Center for
Physics, which is supported by National Science Foundation grant PHY-2210452.
We also thank the Munich Institute for Astro-, Particle and BioPhysics
(MIAPbP) which is funded by the Deutsche Forschungsgemeinschaft (DFG, German
Research Foundation) under Germany´s Excellence Strategy – EXC-2094 –
390783311, as this work was partially performed at the MIAPbP workshop on
Differentiable and Probabilistic Programming for Fundamental Physics.

MK is supported by the US Department of Energy (DOE) under grant DE-
AC02-76SF00515. LH is supported by the Excellence Cluster ORIGINS, which is
funded by the Deutsche Forschungsgemeinschaft (DFG, German Research
Foundation) under Germany’s Excellence Strategy - EXC-2094-390783311.

## References

  * Radovic  _et al._ (2018) A. Radovic _et al._ , “Machine learning at the energy and intensity frontiers of particle physics,” Nature 560, 41–48 (2018).
  * K. Cranmer, U. Seljak, K. Terao, and the Particle Data Group (2022) K. Cranmer, U. Seljak, K. Terao, and the Particle Data Group, “Review of Particle Physics,” Progress of Theoretical and Experimental Physics 2022, 083C01 (2022), chapter 41.
  * Calafiura  _et al._ (2022) P. Calafiura, D. Rousseau, and K. Terao,  _Artificial Intelligence for High Energy Physics_ (WORLD SCIENTIFIC, 2022).
  * Cranmer  _et al._ (2020) K. Cranmer, J. Brehmer, and G. Louppe, “The frontier of simulation-based inference,” Proceedings of the National Academy of Sciences 117, 30055–30062 (2020).
  * Shlomi  _et al._ (2020) J. Shlomi, P. Battaglia, and J.-R. Vlimant, “Graph neural networks in particle physics,” Machine Learning: Science and Technology 2, 021001 (2020).
  * Thais  _et al._ (2022) S. Thais _et al._ , “Graph neural networks in particle physics: Implementations, innovations, and challenges,” (2022), arXiv:2203.12852 .
  * Butter  _et al._ (2023) A. Butter _et al._ , “Machine learning and LHC event generation,” SciPost Phys. 14, 079 (2023).
  * Adelmann  _et al._ (2022) A. Adelmann _et al._ , “New directions for surrogate models and differentiable programming for high energy physics detector simulation,” (2022), arXiv:2203.08806 .
  * Bücker  _et al._ (2005) H. M. Bücker, G. F. Corliss, P. D. Hovland, U. Naumann, and B. Norris, eds.,  _Automatic Differentiation: Applications, Theory, and Implementations_, Lecture Notes in Computational Science and Engineering (Springer, New York, NY, 2005).
  * Baydin  _et al._ (2018) A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind, “Automatic differentiation in machine learning: a survey,” Journal of Machine Learning Research 18, 1–43 (2018).
  * Heinrich and Kagan (2023) L. Heinrich and M. Kagan, “Differentiable Matrix Elements with MadJax,” J. Phys. Conf. Ser. 2438, 012137 (2023), arXiv:2203.00057 .
  * Simpson and Heinrich (2023) N. Simpson and L. Heinrich, “neos: End-to-End-Optimised Summary Statistics for High Energy Physics,” J. Phys. Conf. Ser. 2438, 012105 (2023), arXiv:2203.05570 .
  * Cheong  _et al._ (2022) S. Cheong _et al._ , “Novel light field imaging device with enhanced light collection for cold atom clouds,” Journal of Instrumentation 17, P08021 (2022).
  * Glasserman (2013) P. Glasserman,  _Monte Carlo Methods in Financial Engineering_ (Springer, 2013).
  * Kingma and Welling (2014) D. P. Kingma and M. Welling, “Auto-encoding variational bayes,”  _Proceedings of the International Conference on Learning Representations_ , ICLR’14 (2014), arXiv:1312.6114 .
  * Rezende  _et al._ (2014) D. Rezende, S. Mohamed, and D. Wierstra, “Stochastic backpropagation and approximate inference in deep generative models,” in  _Proceedings of the 31st International Conference on International Conference on Machine Learning - Volume 32_ , ICML’14 (2014) p. II–1278–II–1286.
  * Williams (1992) R. Williams, “Simple statistical gradient-following algorithms for connectionist reinforcement learning,” Machine Learning 8, 229–256 (1992).
  * Arya  _et al._ (2022) G. Arya, M. Schauer, F. Schäfer, and C. Rackauckas, “Automatic Differentiation of Programs with Discrete Randomness,” in  _Advances in Neural Information Processing Systems_, Vol. 35 (Curran Associates, Inc., 2022).
  * Mohamed  _et al._ (2020) S. Mohamed, M. Rosca, M. Figurnov, and A. Mnih, “Monte carlo gradient estimation in machine learning,” (2020), arXiv:1906.10652 .
  * Greensmith  _et al._ (2004) E. Greensmith, P. L. Bartlett, and J. Baxter, “Variance reduction techniques for gradient estimates in reinforcement learning,” Journal of Machine Learning Research 5, 1471–1530 (2004).
  * Abadi  _et al._ (2015) Martín Abadi _et al._ , “TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems,”  (2015), software available from tensorflow.org.
  * Bradbury  _et al._ (2018) J. Bradbury _et al._ , “JAX: composable transformations of Python+NumPy programs,”  (2018).
  * Paszke  _et al._ (2019) Adam Paszke _et al._ , “PyTorch: An Imperative Style, High-Performance Deep Learning Library,” in  _Advances in Neural Information Processing Systems 32_ (Curran Associates, Inc., 2019) pp. 8024–8035.
  * Schulman  _et al._ (2015) J. Schulman, N. Heess, T. Weber, and P. Abbeel, “Gradient estimation using stochastic computation graphs,” in  _Advances in Neural Information Processing Systems_, Vol. 28 (Curran Associates, Inc., 2015).
  * Arya  _et al._ (2023) G. Arya _et al._ , “Differentiating metropolis-hastings to optimize intractable densities,” (2023), arXiv:2306.07961 .
  * Lew  _et al._ (2023) A. K. Lew, M. Huot, S. Staton, and V. K. Mansinghka, “ADEV: Sound automatic differentiation of expected values of probabilistic programs,” Proceedings of the ACM on Programming Languages 7, 121–153 (2023).
  * Krieken  _et al._ (2021) E. Krieken, J. Tomczak, and A. Ten Teije, “Storchastic: A framework for general stochastic automatic differentiation,” in  _Advances in Neural Information Processing Systems_, Vol. 34 (Curran Associates, Inc., 2021) pp. 7574–7587.
  * Bangaru  _et al._ (2021) S. Bangaru, J. Michel, K. Mu, G. Bernstein, T.-M. Li, and J. Ragan-Kelley, “Systematically differentiating parametric discontinuities,” ACM Trans. Graph. 40, 107:1–107:17 (2021).
  * (29) L. Heinrich, M. Feickert, and G. Stark, “pyhf: v0.6.3,” https://github.com/scikit-hep/pyhf/releases/tag/v0.6.3.
  * Heinrich  _et al._ (2021) L. Heinrich, M. Feickert, G. Stark, and K. Cranmer, “pyhf: pure-Python implementation of HistFactory statistical models,” Journal of Open Source Software 6, 2823 (2021).
  * Heinrich and Simpson (2020) L. Heinrich and N. Simpson, “pyhf/neos: initial zenodo release [Link],”  (2020).
  * Carrazza  _et al._ (2021) S. Carrazza, J. M. Cruz-Martinez, and M. Rossi, “Pdfflow: Parton distribution functions on gpu,” Computer Physics Communications 264, 107995 (2021).
  * Ball  _et al._ (2021) R. D. Ball _et al._ , “An open-source machine learning framework for global analyses of parton distributions,” (2021), arXiv:2109.02671 .
  * Nachman and Prestel (2022) B. Nachman and S. Prestel, “Morphing parton showers with event derivatives,” (2022), arXiv:2208.02274 .
  * (35) Differentiable Universe Initiative, home page & software.
  * Carrazza and Dreyer (2019) S. Carrazza and F. A. Dreyer, “Jet grooming through reinforcement learning,” Phys. Rev. D 100, 014014 (2019), arXiv:1903.09644 .
  * Brehmer  _et al._ (2020) J. Brehmer _et al._ , “Hierarchical clustering in particle physics through reinforcement learning,” (2020), arXiv:2011.08191 .
  * Shirobokov  _et al._ (2020) S. Shirobokov _et al._ , “Black-box optimization with local generative surrogates,” in  _Advances in Neural Information Processing Systems_, Vol. 33 (Curran Associates, Inc., 2020) pp. 14650–14662.
  * (39) C. Haack and L. Schumacher, “Machine-learning aided detector optimization of the pacific ocean neutrino experiment,” ICRC2023, 1059.
  * Dorigo  _et al._ (2023) T. Dorigo _et al._ (MODE), “Toward the end-to-end optimization of particle physics instruments with differentiable programming,” Rev. Phys. 10, 100085 (2023), arXiv:2203.13818 .
  * Gorordo  _et al._ (2022) T. Gorordo _et al._ , “Geometry Optimization for Long-lived Particle Detectors,” (2022), arXiv:2211.08450 .
  * Kingma and Ba (2015) D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” in  _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ (2015).


11institutetext: Department of Science and Technology, Linköping University,
Sweden
22institutetext: Center for Medical Image Science and Visualization, Linköping
University, Sweden
33institutetext: Sectra AB, Sweden
33email: {karin.stacke, gabriel.eilertsen, jonas.unger,
<EMAIL_ADDRESS>
# A Closer Look at Domain Shift for Deep Learning in Histopathology

Karin Stacke 1133 Gabriel Eilertsen 1122 Jonas Unger 1122 Claes Lundström
112233

###### Abstract

Domain shift is a significant problem in histopathology. There can be large
differences in data characteristics of whole-slide images between medical
centers and scanners, making generalization of deep learning to unseen data
difficult. To gain a better understanding of the problem, we present a study
on convolutional neural networks trained for tumor classification of H&E
stained whole-slide images. We analyze how augmentation and normalization
strategies affect performance and learned representations, and what features a
trained model respond to. Most centrally, we present a novel measure for
evaluating the distance between domains in the context of the learned
representation of a particular model. This measure can reveal how sensitive a
model is to domain variations, and can be used to detect new data that a model
will have problems generalizing to. The results show how learning is heavily
influenced by the preparation of training data, and that the latent
representation used to do classification is sensitive to changes in data
distribution, especially when training without augmentation or normalization.

###### Keywords:

histopathology domain shift cross-dataset generalization

##  1 Introduction

A fundamental part of machine learning is the problem of generalization, that
is, how to make sure that a trained model performs well on unseen data. If the
unseen data has different distribution, i.e. a domain shift exists, the
problem is significantly more difficult [15, 22] – even the smallest changes
in the statistics as compared to the training data can cause a deep
convolutional neural network (DNN) to fail completely [10].

In the case of digital pathology, domain shift challenges in whole-slide-
imaging (WSI) are typically ascribed to color, brightness and contrast
variations, caused by stain variations and scanner properties [24].
Diversifying the training data can alleviate the problem, e.g., by using
samples from different medical centers. However, there are still no guarantees
that the trained model will generalize to all situations faced by a real-world
application.

Figure 1: Visualizations of non-Gabor-like filters from the first
convolutional layer for Simple CNN (top) and Mini-GoogLeNet (bottom)
architectures, trained with histopathological data. See Section 2 for model
definitions.

Much research has been devoted to understand DNNs trained with natural images,
for example what representations the model learns, how the choice of loss
function and batch size affects the representation [7, 6, 9] and how this can
be leveraged for transfer learning, novelty detection and detection of
adversarial examples [6, 17, 14]. However, the relevance of these research
results for histopathology is unclear, as the data characteristics
substantially differ from natural images. As Raghu et al. observed in [16],
models trained on medical image datasets did not learn Gabor filters, which
were the case when trained on ImageNet [4]. The same is true for models
trained with histopathological data, as demonstrated in Figure 1. While Gabor
filters are known to be a common denominator between convolutional neural
networks (CNNs) trained for different tasks and on different data, this
counter-example motivates us to look closer at how CNN optimization is shaped
by WSI data, and what the implications are for generalization to unseen
domains.

The problem of domain shift in pathological images from different medical
centers and different scanners has been discussed in previous work [20, 11, 2,
1, 3, 21], but not with the purpose of getting a deeper understanding of the
underlying elements of the problem. In this paper, we take a closer look at
how different strategies for increasing generalization affects not only
performance, but also the representations learned by a CNN trained on
histopathology data. To this end, we look at the problem from three different
perspectives, which constitute the main contributions of this paper:

  1. 1.

We analyze the test accuracy of CNNs trained for tumor classification on H&E
stained images: how it changes with different models, test data and
augmentation strategies.

  2. 2.

We optimize for the input images which maximally activate different
convolutional filters of a trained model, to get an indication on what
particular features a model respond to.

  3. 3.

We use the filter activations of a trained model to formulate a novel
representation shift metric that captures differences in feature
representation caused by test data from a different distribution than the
training data.

The experiments show that a tumor classifier respond to very different
features depending on how training data is processed. We are also able to
demonstrate a correlation between the representation shift and the drop in
classification accuracy on images from a new domain. Together, the experiments
help us shape a better understanding on what a tumor classifier has learned
and how sensitive its representation is to changes in data distribution. We
believe that this understanding is crucial for future development of reliable
deep models that are robust to domain variations and thus possible to use in a
real-world clinical pipeline.

##  2 Experimental setup and methods

Below, we first describe the data and models used, followed by the setup of
the three experiments of our study. Finally, the details of our proposed
domain shift quantification are given.

(a) Original

(b) Color augm.

(c) Stain norm.

(d) CycleGAN

Figure 2: Five example images for each center (the three top rows are scanned
with Scanner 1, the forth row with Scanner 2 and the last row with Scanner 3).

###  2.1 Data

We use the CAMELYON17 dataset [12], consisting of H&E stained lymph node
whole-slide images, captured with three different scanners, (henceforth
denoted Scanner 1, Scanner 2 and Scanner 3) collected from three, one and one
medical centers, respectively. Three types of data transformations were
included (Fig. 2), in order to measure how well traditional and state-of-the-
art methods can overcome the domain shift between scanners:

####  2.1.1 Color and intensity augmentations

(henceforth shortened color augmentation) aim to increase image diversity in
order to make the model more robust to color variations. The augmentation was
performed randomly on each channel of the HSV color space. The amount of
augmentation was chosen for best generalization, which meant extreme and
unnatural color variations. However, this was necessary in order to force a
model to focus on other features than color.

####  2.1.2 Stain normalization

refers to first decomposing the image by stain colors (stain separation), and
then normalizing it based on a target distribution. A wide range of stain
normalization techniques has been presented. For this study, the method
presented in [23] was used. Each patch was normalized in relation to a
reference patch, taken as a representative patch (in terms of structure and
color) from one of the medical centers (Canisius-Wilhelmina Hospital, CWZ).

####  2.1.3 CycleGAN

[25] is a method for image-to-image translation, with the goal of translating
an image using a mapping $G$, from source domain $X$ to target domain $Y$. The
goal is to learn this mapping $G$, such that the distribution of $G(X)$ is
indistinguishable from $Y$. In order to achieve this, the inverse mapping
$F:Y\rightarrow X$ is introduced to force $F(G(X))\approx X$. Using this
approach, features important to the target domain is transferred to $X$, but
unimportant ones are unchanged. In our work, we use this approach to transfer
images from one medical center to another. The class label of the image is
left unchanged, only the features that are different between the centers are
changed. Images from all centers where transferred to medical center CWZ.

###  2.2 Model architecture

Two model architectures were evaluated. One simple CNN architecture consisting
of three convolutional layers and two fully connected layers (with dropout, no
batch normalization), and a small version of GoogLeNet architecture [18] (only
using the output from the first auxiliary classifier, henceforth denoted Mini-
GoogLeNet). The motivations for choice in architecture is the small training
dataset requiring a reduced size in model capacity to prevent overfitting, and
as the goal of this experiment is not to achieve state-of-the-art results,
simple and small models are to be preferred. We also tested larger and more
high-capacity models, such as ResNet [8] and Inception-v3 [19]. However, the
performance did not improve, and in most cases overfitting was a significant
problem. All models were trained with geometric augmentations (random flip,
rotation, and crop), with a patch size of 96x96px, and a resolution of 0.25
microns per pixel. The datasets were separated at slide level, making sure no
patches from the same slide were in both training and test sets. All results
shown are the average of three training sessions.

###  2.3 Experiments

The three domain shift analysis experiments were designed as follows.

####  2.3.1 Cross-dataset generalization for tumor classification

was evaluated by training CNN models as tumor classifiers, detecting tumor vs.
non-tumor patches. The training data consisted of slides scanned with Scanner
1. For test, both same-scanner data as well as unseen data from the two other
scanners was used, with results reported as patch classification accuracy.

####  2.3.2 Feature visualization

of learned model filters was employed to give a visual representation of what
image features the model respond to. Using a gradient ascent method [5, 13] we
start from a noise image and iteratively optimize towards the input image
which maximally activates a filter in the last convolution layer.

####  2.3.3 Activation difference

for filters responding to data from different domains was the third aspect in
focus. To study this, we developed a quantitative metric, representation
shift, described in the next section. The metric was applied to the learned
models of the cross-dataset generalization experiment.

###  2.4 Representation shift metric

To quantify domain shift in the context of a CNN, we are not only interested
in the statistical differences between images from different domains, but also
in how the images are processed by the CNN. For this purpose, we propose the
representation shift metric. This measures the differences in distributions of
the learned feature representation, comparing the training set to a dataset
from a different domain. We denote by $p_{r,c_{i}}$ the continuous
distribution of mean activations from the convolutional filter $c_{i}$ in the
last convolutional layer $\\{c_{1},...,c_{L}\\}$, computed using input data
$X_{r}=\\{x_{r1},...,x_{rn}\\}$. A second dataset, $X_{\theta}=\\{x_{\theta
1},...,x_{\theta m}\\}$ similarly generates $p_{\theta,c_{i}}$. If
$\pi(p_{r,c_{i}},p_{\theta,c_{i}})$ are the joined distributions with margins
$p_{r,c_{i}}$ and $p_{\theta,c_{i}}$, the Wasserstein distance between the
distributions is given by:

$W(p_{r,c_{i}},p_{\theta,c_{i}})=\inf_{\pi\in\Gamma(p_{r,c_{i}},p_{\theta,c_{i}})}\int_{RxR}|(x-y)|d\pi(x,y).$
(1)

We define the representation shift $R$ as the mean distance between the
distributions over all filters $c_{i}$,

$R(p_{r},p_{\theta})=\frac{1}{L}\sum_{i=1}^{L}W(p_{r,c_{i}},p_{\theta,c_{i}}).$
(2)

If $X_{r}$ and $X_{\theta}$ are statistically similar in the image domain, or
if the CNN maps the datasets to similar representations, filter responses
should be similar, and $R(p_{r},p_{\theta})$ small. We can then expect a
similar classification accuracy of the datasets. If $R(p_{r},p_{\theta})$ is
large, then filter responses has changed between the datasets, resulting in a
higher risk of generalization error. The method does not aim to do out-of-
distribution or novelty detection for individual images (as e.g. [14, 17]),
but aims to indicate an elevated risk of decreased classification accuracy for
a dataset that is assumed to be similar to the original data. As the metric
does not require annotated data, it can serve as a simple initial test to
evaluate if new data (e.g. from a new scanner) is handled well by an already
trained model, i.e. if the learned feature representation applies to the new
data. Also, as the metric is tightly connected to the specific model used, we
expect that different models and training strategies will results in very
different robustness to changes in statistics (Fig. 4).

Table 1: Mean accuracies (%), training a tumor classifier on Scanner 1, testing it on same-scanner data and other data, using two different model architectures. | Simple CNN | Mini-GoogLeNet
---|---|---
Train Test | Scanner 1 | Scanner 2 | Scanner 3 | Scanner 1 | Scanner 2 | Scanner 3
Orig. data | $88.0\pm 3.4$ | $51.2\pm 0.7$ | $78.4\pm 1.7$ | $87.2\pm 2.8$ | $60.0\pm 14.4$ | $74.3\pm 9.6$
Color aug. | $87.2\pm 4.4$ | $80.2\pm 1.7$ | $\mathbf{82.3\pm 0.8}$ | $87.3\pm 4.5$ | $\mathbf{84.5\pm 1.2}$ | $\mathbf{82.8\pm 5.8}$
Stain norm. | $89.4\pm 3.0$ | $\mathbf{84.1\pm 2.0}$ | $77.6\pm 4.6$ | $88.1\pm 1.5$ | $80.5\pm 4.1$ | $77.9\pm 6.9$
CycleGAN | $\mathbf{94.5\pm 0.2}$ | $83.2\pm 5.3$ | $80.9\pm 3.1$ | $\mathbf{91.0\pm 1.8}$ | $83.4\pm 3.1$ | $78.2\pm 2.2$

##  3 Results

###  3.1 Cross-dataset Generalization

Table 1 shows validation accuracies when training a model on data from one
scanner, and evaluating it on another. Both model architectures suffer from
poor generalization when no augmentation is used, with significant drop in
mean accuracy ($21.7$ percentage points (p.p.) mean drop from same-scanner
data to unseen data). Adding color augmentation results in better
generalization across the datasets ($4.75$ p.p. drop), with similar
performances on both unseen datasets. Stain normalization performed adequate
on one of the unseen datasets (Scanner 2), but had a quite large drop for the
Scanner 3 data (mean of $9.2$ p.p. drop). CycleGAN gave best performance on
same-scanner data, but these high values in accuracy did not transfer to the
other data sets ($11.45$ p.p drop).

Figure 3: Each row shows example images that maximally activates different
filters in models Simple CNN (left) and Mini-GoogLeNet (right), trained with
(from top to bottom): original data, color aug., stain norm., and CycleGAN.
Best viewed digitally.

###  3.2 Feature visualization

Images visualizing features that give maximal activation of filters in the
last convolutional layer are shown in Fig. 3. Note that color reproduction may
not be completely accurate, since no regularization on maximal pixel values is
done during the optimization process. However, the relative color differences
between the models can be compared, showing that the Mini-GoogLeNet model
trained with color augmentation completely ignores the color components in the
images. This is in strong contrast with the model trained on original data,
which seem to react to one or two colors only. Thus, it is not surprising that
the generalization to another domain is poor – if the representation is highly
dependent on color information, we can expect that small differences in color
distribution will affect the performance negatively. Since differences in
color is one of the most evident visual differences between images from
different scanners, this is indeed the case.

Images from Mini-GoogLeNet reveal that the models have learned cell-like
circular structures, and trained with color augmentation or stain
normalization show even stronger resemblances with nucleus structures and cell
membranes. The Simple CNN architecture seems to find more low-level
structures, which might indicate that the model has not learned the expected
features which separate tumor cells from non-tumorous tissue. Both stain
normalization and CycleGAN transformation unifies the color variations in the
dataset, which is visible by a more homogeneous color scale in these images.

Figure 4: Relation between patch level accuracy (see Table 1) and
representation shift (Equation 2) for Simple CNN (left) and Mini-GoogLeNet
(right) with and without data transformation. A regression line shows the
correlation between the variables.

###  3.3 Representation shift

Fig. 4 shows the representation shift in relation to model patch level
accuracy. Models trained with original data show large representation shifts,
which is to be expected as no measures where taken to handle the domain shift.
The representation shift is largely reduced with any of the data
transformation techniques. There is not a one-to-one mapping between the
representation shift and classification accuracy, which is expected as the
final accuracy depends on the final fully connected layers. However there is a
clear negative correlation between representation shift and classification
accuracy.

##  4 Discussion

This study is an initial attempt to analyze how histopathological data shapes
DNNs differently than natural images do. Although the presented techniques can
be straight-forwardly used for detecting model vulnerability, we also believe
that such information can be used to increase our knowledge about deep
learning in histopathology.

Using feature visualization (Fig. 3), we see that different preparations of
training data results in very different latent representations. Even if this
way of visualizing the learned features only give us parts of the truth, it is
an important insight into the “black box” of CNNs in histopathology. We see
that these different representations are more or less sensitive to changes in
input distsribution (Table 1). For example stain normalization or CycleGAN may
unify the data and reduce much of the variance, which help the cross-dataset
generalization, but does not necessarily create robust models.

The results also show how the representation shift metric can give a valuable
contribution for domain shift analysis on a new dataset (Fig. 4), without
requiring annotations on the new data. While it is still a coarse measure, we
argue that it is useful tool for a first investigation of the magnitude of a
domain shift.

In future work, we will further investigate the correlation between
activations and domain shift. In particular, we aim to further explore
representation shift metrics as a tool for precise domain shift analysis, and
how this is affected by different models, datasets, and domain adaptation
techniques.

## References

  * [1] Arvidsson, I., et al.: Generalization of prostate cancer classification for multiple sites using deep learning. In: 2018 IEEE 15th International Symposium on Biomedical Imaging (ISBI 2018). pp. 191–194 (Apr 2018)
  * [2] Bentaieb, A., Hamarneh, G.: Adversarial Stain Transfer for Histopathology Image Analysis. IEEE Transactions on Medical Imaging 37(3), 792–802 (Mar 2018)
  * [3] Ciompi, F., et al.: The importance of stain normalization in colorectal tissue classification with convolutional networks. In: 2017 IEEE 14th International Symposium on Biomedical Imaging (ISBI 2017). pp. 160–163 (Apr 2017)
  * [4] Deng, J., et al.: ImageNet: A large-scale hierarchical image database. In: 2009 IEEE Conference on Computer Vision and Pattern Recognition. pp. 248–255 (Jun 2009). https://doi.org/10.1109/CVPR.2009.5206848
  * [5] Erhan, D., et al.: Visualizing Higher-Layer Features of a Deep Network. University of Montreal (Jan 2009)
  * [6] Frosst, N., et al.: Analyzing and Improving Representations with the Soft Nearest Neighbor Loss. arXiv:1902.01889 [cs, stat] (Feb 2019), arXiv: 1902.01889
  * [7] Geirhos, R., et al.: ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness. arXiv:1811.12231 [cs, q-bio, stat] (Nov 2018), arXiv: 1811.12231
  * [8] He, K., et al.: Deep Residual Learning for Image Recognition. In: 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). pp. 770–778 (Jun 2016)
  * [9] Keskar, N.S., et al.: On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima. arXiv:1609.04836 [cs, math] (Sep 2016), arXiv: 1609.04836
  * [10] Kurakin, A., et al.: Adversarial examples in the physical world. arXiv:1607.02533 [cs, stat] (Jul 2016), arXiv: 1607.02533
  * [11] Lafarge, M.W., et al.: Domain-adversarial neural networks to address the appearance variability of histopathology images. arXiv:1707.06183 [cs] 10553 (2017), arXiv: 1707.06183
  * [12] Litjens, G., et al.: 1399 H&E-stained sentinel lymph node sections of breast cancer patients: the CAMELYON dataset. Gigascience 7(6) (Jun 2018)
  * [13] Olah, C., Mordvintsev, A., Schubert, L.: Feature Visualization. Distill 2(11), e7 (Nov 2017)
  * [14] Papernot, N., McDaniel, P.: Deep k-Nearest Neighbors: Towards Confident, Interpretable and Robust Deep Learning. arXiv:1803.04765 [cs, stat] (Mar 2018), arXiv: 1803.04765
  * [15] Quiñonero-Candela, J. (ed.): Dataset shift in machine learning. Neural information processing series, MIT Press, Cambridge, Mass (2009), oCLC: ocn227205909
  * [16] Raghu, M., et al.: Transfusion: Understanding Transfer Learning with Applications to Medical Imaging. arXiv:1902.07208 [cs, stat] (Feb 2019), arXiv: 1902.07208
  * [17] Schultheiss, A., et al.: Finding the Unknown: Novelty Detection with Extreme Value Signatures of Deep Neural Activations. In: Roth, V., Vetter, T. (eds.) Pattern Recognition. pp. 226–238. Lecture Notes in Computer Science, Springer International Publishing (2017)
  * [18] Szegedy, C., et al.: Going Deeper with Convolutions. arXiv:1409.4842 [cs] (Sep 2014), arXiv: 1409.4842
  * [19] Szegedy, C., et al.: Rethinking the Inception Architecture for Computer Vision. In: 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). pp. 2818–2826. IEEE, Las Vegas, NV, USA (Jun 2016)
  * [20] Tellez, D., et al.: H and E stain augmentation improves generalization of convolutional networks for histopathological mitosis detection. In: Medical Imaging 2018: Digital Pathology. vol. 10581, p. 105810Z. International Society for Optics and Photonics (Mar 2018)
  * [21] Tellez, D., et al.: Quantifying the effects of data augmentation and stain color normalization in convolutional neural networks for computational pathology. arXiv:1902.06543 [cs] (Feb 2019), arXiv: 1902.06543
  * [22] Torralba, A., Efros, A.A.: Unbiased look at dataset bias. In: CVPR 2011. pp. 1521–1528 (Jun 2011)
  * [23] Vahadane, A., et al.: Structure-Preserving Color Normalization and Sparse Stain Separation for Histological Images. IEEE Transactions on Medical Imaging 35(8), 1962–1971 (Aug 2016)
  * [24] Yagi, Y.: Color standardization and optimization in Whole Slide Imaging. Report, BioMed Central (Dec 2011)
  * [25] Zhu, J.Y., et al.: Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks. arXiv:1703.10593 [cs] (Mar 2017), arXiv: 1703.10593