# Improving Compositionality of Neural Networks by Decoding Representations to
Inputs

Mike Wu, Noah Goodman, Stefano Ermon
Department of Computer Science
Stanford University
Stanford, CA 94303
<EMAIL_ADDRESS>

###### Abstract

In traditional software programs, we take for granted how easy it is to debug
code by tracing program logic from variables back to input, apply unit tests
and assertion statements to block erroneous behavior, and compose programs
together. But as the programs we write grow more complex, it becomes hard to
apply traditional software to applications like computer vision or natural
language. Although deep learning programs have demonstrated strong performance
on these applications, they sacrifice many of the functionalities of
traditional software programs. In this paper, we work towards bridging the
benefits of traditional and deep learning programs by jointly training a
generative model to constrain neural network activations to “decode” back to
inputs. Doing so enables practitioners to probe and track information encoded
in activation(s), apply assertion-like constraints on what information is
encoded in an activation, and compose separate neural networks together in a
plug-and-play fashion. In our experiments, we demonstrate applications of
decodable representations to out-of-distribution detection, adversarial
examples, calibration, and fairness — while matching standard neural networks
in accuracy.

##  1 Introduction

Figure 1: Decoding activations to images for misclassified examples.

Traditional hand-written computer programs are comprised of a computational
graph of typed variables with associated semantic meaning. This structure
enables practitioners to interact with programs in powerful ways (even if they
are not the author) — such as debug code by tracing variables back to inputs,
apply unit tests and assertions to block against erroneous behavior, and
compose separate programs together for more complex functionality. This
compositionality and flexibility has enabled software to permeate every
industry and application. Even still, traditional software has its
limitations: it is difficult to hand-write programs to classify images or
extract sentiment from natural language. For these functionalities, deep
learning and neural networks [53] have become the dominant approach [29, 2,
60].

This distinction has manifested in a budding shift from the classical stack of
Python, C++, etc. to “software 2.0” [26, 63, 48], which is built on neural
network weights. While neural networks have made impressive progress on
complex tasks, switching to software 2.0 comes at a sacrifice of many of the
desirable properties of “software 1.0” programs. The closest approximation to
a variable in a neural network is an activation. It is difficult to understand
a neural network’s computation from an activation value, and there is little
to associate an activation with semantic meaning. A practitioner cannot impose
unit tests on neural networks, let alone write assertion statements to
constrain valid logic: checking the values that an activation takes is usually
not enough to gauge correctness nor meaning. Moreover, given multiple neural
networks (e.g. one for classification and one for segmentation), composing
them together requires retraining from scratch. Consequently, neural networks
deployed in practice are brittle and require maintenance [54].

In this paper, we work towards a “software 1.5”, where the overarching goal is
to bridge the expressivity of deep learning with the engineering practicality
of traditional software. Our desiderata is to construct neural networks whose
activations can be decoded to semantic meanings — taking a step towards
enabling deep learning practitioners to probe the model by tracing activations
back to inputs, apply assertion-like constraints on what information is stored
in an activation, and also compose neural network modules together in new and
useful ways.

To start, we propose to train a neural network classifier jointly with a
generative model whose role is to map the classifier’s activations back to the
input, approximating invertibility. This is only an approximation as
activations deeper into the network drop information, meaning decoded inputs —
while semantically related to the original input — are themselves new. This
joint training optimizes the classifier not only to be highly performant, but
constrains its activations to a manifold that the generative model is able to
decode. Consequently, the neural network’s entire computation is constrained
to a “decodable” vector space as all activations can be mapped back to the
input space. Continuing the software analogy, we are enforcing our neural
network activations to have the same “type” as training examples. It is this
insight that enables us to treat our neural networks more like traditional
software, where practitioners can analyze and repurpose decoded activations.

In our experiments, we make progress towards each of the aforementioned
features of software 1.0 that we could not do with software 2.0: probing our
“variables” (i.e. activations) for meaning, constraining model logic with
“assertions”, and composing neural network modules together. In doing so, we
will explore applications of decodable neural networks to out-of-distribution
detection, adversarial examples, calibration, and fairness — a representative
collection of engineering features necessary for real world applications of
deep learning. Throughout this, we show that decodability comes at low cost as
we find equivalent accuracy as a standard neural network.

(a) Decodable NN

(b) Recursively Decodable NN

Figure 2: Decodable Neural Networks: (left) A generative model $g_{\phi}$ is
jointly trained with a classifier $f_{\theta}$ to map activations to inputs.
The images shown are actual reconstructions from a trained network that
misclassified whether a celebrity is wearing a hat. (right) A recursively
decodable neural network.

##  2 Background

Neural Networks  We will focus on supervised neural networks for
classification, although it is simple to extend to regression tasks. We denote
a neural network by $f_{\theta}$ where $\theta$ represents its trainable
parameters. A neural network $f_{\theta}$ maps an example $x$ to predicted
probabilities over $K$ classes. Typically, a neural network is composed of $L$
blocks. For example, if $f_{\theta}$ is a multi-layer perception, each block
is a linear layer followed by an non-linearity. If $f_{\theta}$ is a ResNet
[21], each layer is a residual block. In computing its prediction, a neural
network $f_{\theta}$ produces activations $\\{h_{1},\ldots,h_{L}\\}$ where
$h_{l}$ is the output of the $l$-th block.

To train a neural network, we solve an optimization function of the form

$\mathcal{L}_{\textup{nn}}(x;\theta)=\log
p(y|f_{\theta}(x))+\beta\Omega(x;\theta)$ (1)

using stochastic gradient descent. In Equation 1, the notation $x$ denotes an
input example and $y$ its label. The function $\Omega$ represents an auxiliary
objective such as a regularizer. The hyperparameter $\beta>0$ is used to scale
the auxiliary loss. For classification, $\log p(y|\cdot)$ is cross entropy.

Generative Models  Broadly, we are given a latent variable $H$ and an observed
variable $X$. A generative model $g_{\phi}$ captures the distribution
$p(x|h)$, mapping samples of the latent code $h$ to samples of the observed
variable $x$. For a deep generative model, $g_{\phi}$ is parameterized by a
neural network. In the simplest case where $g_{\phi}$ is an autoencoder, we
are also given an encoder $q$ that maps observed samples $x$ to latent codes
$h$. Then, we optimize the objective

$\mathcal{L}_{\textup{ae}}(x;\phi)=\log p(x|g_{\phi}(q(x)))$ (2)

where for continuous $x$, the expression $\log p(x|h)$ is the log Gaussian
probability with mean $h$ and fixed variance (or equivalently, mean squared
error). More complex generative models might impose a prior distribution on
$h$ [27], an adversarial discriminator [35, 18], or invertibility of
$g_{\phi}$ [44].

##  3 Neural Networks with Decodable Activations

In Section 1, we motivated our desiderata for activations that can be mapped
back to the input space. We now propose to do this by “inverting” a neural
network at a particular activation. Concretely, given an input $x$, a neural
network $f_{\theta}$, we can compute an activation $h$ by computing the
forward pass $f_{\theta}(x)$. Then, suppose we could construct the set:

$S(h)=\\{\hat{x}:\hat{x}\sim p_{\mathcal{D}},f_{\theta}(\hat{x})=h\\}$ (3)

where $p_{\mathcal{D}}$ represents the empirical distribution that $x$ is
sampled from. Now, any input $\hat{x}\in S(h)$ is a valid mapping of $h$ back
to the input space. (Trivially, $x\in S(h)$.) We call this $\hat{x}$ a
decoding of the activation $h$. Unfortunately, in practice we do not know how
to construct such a set $S(h)$.

But, we do know how to (approximately) sample from $S(h)$. We can train a
generative model $g_{\phi}$ to reconstruct the original input $x$ from neural
network activations $h$ (recall $x\in S(h)$). Then, we can sample $x\sim
g_{\phi}(h)$, which we know how to do, to sample from $S(h)$. If we jointly
train this generative model $g_{\phi}$ with the neural network $f_{\theta}$,
we can optimize for activations that both solve the classification task and
decode back to inputs at once.

So far, we have only considered a single activation despite our neural network
having $L$ such activations. Although one could train a separate generative
model per layer, we found it efficient and performant to share weights $\phi$
across all layers. Our joint objective is:

$\mathcal{L}(x;\theta,\phi)=\log
p(y|f_{\theta}(x))+\beta\frac{1}{L}\sum_{l=1}^{L}\log
p(x|g_{\phi}(h_{l})),\qquad h_{1},...,h_{L}\textup{ come from }f_{\theta}(x)$
(4)

where the first term is cross entropy as in Equation 1 and the second term is
a reconstruction error (Equation 2) averaged over layers. Comparing Equation 4
to Equation 1, we can interpret the generative loss as a regularization term
$\Omega$. Intuitively, in a standard neural network, the activations
$h_{1},\ldots,h_{L}$ are unconstrained, chosen only to minimize the loss.
However, in Equation 4, the network $f_{\theta}$ is constrained to produce
activations that the generative model $g_{\phi}$ can invert. We call this
setup a Decodable Neural Network, or DecNN. See Figure 2a for an illustration.

After training, given a new example $x^{*}$, we first compute hidden
activations $h^{*}_{1},\ldots h^{*}_{L}$ through $f_{\theta}(x^{*})$. For each
activation $h^{*}_{l}$, we can decode it to an input using
$g_{\phi}(h^{*}_{l})$. In the remaining sections, we show how DecNN enables us
to probe what information is stored in an activation (Section 4), compose
neural networks modules (Section 5), and constrain what a network learns
(Section 6).

##  4 Probing Decodable Activations

For a classifier, the Data Processing Inequality [7, 4] tells us that
information about the input $x$ embedded in activation $h_{l}$ decreases as
$l$ increases. More formally, the mutual information follows a sequence of
inequalities
$\mathcal{I}(X;H_{1})\geq\cdots\geq\mathcal{I}(X;H_{L})\geq\mathcal{I}(X;Y)$
where $X$, $H_{l}$, and $Y$ are random variables representing the inputs,
activations, and outputs, respectively. In practice, this theory leads us to
expect “useless” information to be dropped deeper into the network, while
“useful” information is kept to solve the classification task.

In a traditional software program, debugging code amounts to creating
breakpoints and printing values to understand what information a variable
contains. For a neural network, we similarly wish to probe what information an
activation captures. While many algorithms have been developed for specific
architectures [55, 1, 34], they require a nontrivial amount of additional
computation. With DecNN, probing activations is more straightforward. Since
decoded activations have the same type as inputs (which we assume are human
interpretable), we can directly visualize information. If we examine all $L$
activations at once, we can track what information is dropped through the
network.

Table 1(a) shows test accuracies comparing DecNN to an standard network,
finding no performance loss. Given this, Figure 3 shows randomly chosen
examples taken from DecNNs trained on FashionMNIST and CelebA (see Appendix
for full figure). In each subfigure, the leftmost image is from the dataset
while the remaining eight images are decoded activations ($f_{\theta}$ is an
MLP with $L=8$ layers).

(a)

(b)

(c)

(d)

Figure 3: Visualizing activations: the left two images shows decoded
activations when the model correctly classified examples; the right two shows
misclassifications.

The left two images show decoded activations for correctly-classified examples
whereas the right two images show mis-classified examples. For correctly-
classified examples, we observe that the decoded activations tends towards
class prototypes. Decoded images from the 7th or 8th activation lose details
(e.g. patterns, logos, or shoe straps disappear from clothing). The top row of
Figure 3a presents a good example: the tank-top (a rare form of the t-shirt)
is projected to a more prototypical t-shirt. We also find that mis-classified
examples morph into images of the incorrect class. The seventh row of Figure
3c morphs the two articles of clothing that compose the dress to build a
shirt. The seventh row of Figure 3d show the celebrity’s hat transform into a
background color.

As a practitioner, it is useful to understand what features a neural network
is attending to in making a prediction, or to understand what features make it
susceptible to making an error. Decodable activations take a step towards
“printing activations” for debugging.

##  5 Composing Decodable Neural Networks

In traditional software, given two functions $f_{1}$ and $f_{2}$, we can
easily compose them together by $f_{1}(f_{2}(x))$ for some input $x$. With
this as inspiration, we wish to similarly compose neural networks together.
But given two image classifiers, it is not clear how to do so. How do we tie
outputs from one neural network to inputs for the other? Both networks expect
images as input.

Since decoded activations have the same type as inputs, we are free to treat
decoded activations as inputs for another model. In other words, we can easily
compose neural networks together using decoded activations as a conduit: given
two models $f_{1}$ and $f_{2}$, and an input $x$, we can compute
$f_{2}(g_{\phi}(h))$ where $h$ is derived from $f_{1}(x)$ and use it to
optimize a joint objective. In the next two subsections, we explore two kinds
of neural network compositions: one by composing a model with itself and the
other by composing a model with a pretrained one.

###  5.1 Recursive Self-Composition

Our first composition will be to compose a DecNN model with itself. But if we
can do that, we can also recursively compose a DecNN model with itself
infinitely:

Given a decoded activation $g_{\phi}(h^{(0)}_{l})$ for any layer $l$, we can
feed it as input back into the classifier, $f_{\theta}(g_{\phi}(h^{(0)}_{l}))$
and maximize agreement with the true label $y$. Intuitively, if
$g_{\phi}(h_{l})$ is a good reconstruction of the input $x$ (i.e.
$g_{\phi}(h_{l})$ and $x$ belong to the same set $S(h_{l})$), then a good
classifier $f_{\theta}$ should predict the correct label with it as input.
Now, we also observe that the computation $f_{\theta}(g_{\phi}(h^{(0)}_{l}))$
produces $L$ additional activations, $h^{(1)}_{1},\ldots,h^{(1)}_{L}$ where we
use the superscript to represent the number of compositions. We can repeat
this exact process with the same intuition, decoding $h^{(1)}_{l}$ and passing
it to $f_{\theta}$ to produce $h^{(2)}_{l}$. Computational limits aside, this
can be repeated ad infinitum.

One interpretation of this procedure is as data augmentation: we saw in
Section 4 that decoded activations share features relevant to the prediction
task (e.g. generated celebrities wear hats) but vary other irrelevant features
(e.g. face shape, hair color, etc.). Notably, these generated images are not
in the original dataset. Thus, training a model using decoded activations as
inputs is similar to augmenting the initial dataset. The only difference is
that gradients are propagated back through recursion to the inputs, which is
also a form of regularization.

Admittedly, what we have described so far is intractable. Besides recursing
infinitely, each activation produces $L$ more activations, creating a tree
(with root node $x$) with a large branching factor $L$. In practice, we make
two adjustments — First, we limit the recursion to a max depth $D$. Second,
instead of constructing the full graph of activations, which would have
$O(L^{D})$ nodes, we randomly sample “paths” in the graph. Specifically,
starting at the root, we compute $\\{h^{(0)}_{1},\ldots,h^{(0)}_{L}\\}$
through $f_{\theta}(x)$ and randomly choose an integer $l_{0}\in[1,L]$. Then,
we compute $\\{h^{(1)}_{1},\ldots,h^{(1)}_{L}\\}$ through
$f_{\theta}(g_{\phi}(h^{(0)}_{l_{0}}))$ i.e. we classify the $l_{0}$-th
decoded activation from depth 0. Again, we randomly pick $l_{1}$ to compute
$\\{h^{(2)}_{1},\ldots,h^{(2)}_{L}\\}$ through
$f_{\theta}(g_{\phi}(h^{(1)}_{l_{1}}))$ and repeat until depth $D$.

Building on Equation 4, we can write this “recursive” objective as:

$\hat{\mathcal{L}}(x;\theta,\phi)=\log
p(y|f_{\theta}(x))+\sum_{d=1}^{D}\alpha^{d}\log
p(y|f_{\theta}(g_{\phi}(h^{(d-1)}_{l_{d-1}})))+\beta\sum_{d=1}^{D}\alpha^{d}\log
p(x|g_{\phi}(h^{(d-1)}_{l_{d-1}}))$ (5)

where the sequence of integers $l_{0},\ldots l_{D-1}$ are each sampled
uniformly from $[1,L]$. We introduce a new hyperparameter $\alpha\in[0,1]$
that geometrically downweights later depths (if $\alpha=0$, Equation 5 reduces
to maximum likelihood). The first term in Equation 5 is the usual
classification objective. The second term is a sum of cross entropy terms that
encourage the decoded activation at depth $d$ to predict the correct label
through $f_{\theta}$. The third and final term is a sum of reconstruction
losses that encourage the decoded activation at depth $d$ to approximate the
original input $x$.

We call this model a “Recursively Decodable Neural Network”, abbreviated
ReDecNN. Note that ReDecNN is a special case of an ensemble network. That is,
every path sampled in the activation tree by picking $l_{0},\ldots l_{D-1}$
builds a legitimate classifier. There are $O(L^{D})$ such classifiers implicit
in the ReDecNN. We can interpret optimizing Equation 5 as training individual
classifiers that are Monte Carlo sampled from the tree every gradient step.
See Figure 2b for an illustration.

Our goal in the next few paragraphs is to utilize this ensemble to measure
uncertainty.

Method | MNIST | Fashion | CelebA
---|---|---|---
Standard |  $97.1$ (0.11) |  $87.7$ (0.19) |  $90.8$ (<0.1)
DecNN |  $\mathbf{97.8}$ (0.17) |  $\mathbf{88.8}$ (0.19) |  $90.8$ (<0.1)
ReDecNN |  $97.5$ (0.20) |  $88.1$ (0.16) |  $90.8$ (0.16)
Dropout |  $95.9$ (0.13) |  $80.7$ (0.24) |  $88.5$ (0.12)
MC-Dropout |  $91.3$ (0.14) |  $78.6$ (0.23) |  $87.8$ (0.19)
BayesNN |  $96.3$ (0.51) |  $87.0$ (0.36) |  $88.2$ (0.22)
Ensemble |  $97.6$ (0.11) |  $88.4$ (0.12) |  $\mathbf{90.9}$ (0.11)
(a) Classification performance on test set.

Method | MNIST | Fashion | CelebA
---|---|---|---
ReDecNN | $0.869$ | $0.795$ | $\mathbf{0.692}$
MC-Dropout | $\mathbf{0.878}$ | $\mathbf{0.818}$ | $0.657$
BayesNN | $0.537$ | $0.574$ | $0.587$
Ensemble | $0.531$ | $0.559$ | $0.605$
(b) ROC-AUC of separating correctly classified and misclassified examples
using Eq. 6 as a score. The stdev. for all entries over 3 runs are <0.01.

Measuring Uncertainty  As a first step, we compare the performance of ReDecNN
to several baselines: a standard neural network, a neural network with
dropout, two kinds of Bayesian neural networks — Monte Carlo dropout [15],
abbreviated MC-Dropout, and weight uncertainty [6], abbreviated BayesNN — and
a naive ensemble network where we train $D$ copies of the standard neural
network with different initializations. See supplement for training details.
Table 1(a) shows accuracies over a held-out test set, averaged over three
runs. While we observe a drop in accuracy in Bayesian NNs, we observe equal
performance of ReDecNN to a standard neural network.

The next step is to compare the quality of model uncertainty. For the ReDecNN,
for a given input $x$, we measure uncertainty as follows: sample $N$ different
classifiers from the activation tree, and make predictions with each
classifier on $x$. The uncertainty is the entropy of the predictions:

$\textup{Uncertainty}(x)=-\sum_{c=1}^{K}\hat{p}(y=c)\log\hat{p}(y=c)$ (6)

where $\hat{p}(y=c)$ is the empirical probability (e.g. normalized count) of
predicting class $c$ out of $N$ classifiers. A higher uncertainty metric would
represent more disagreement between classifiers. We can compute the same
metric for a naive ensemble, as well as for MC-Dropout (where repeated calls
drop different nodes) and BayesNN (by sampling weights from the learned
posterior).

We want to test if the model’s uncertainty is useful. One way is to correlate
the uncertainty with when the model makes prediction errors. It would be
useful if the model was less uncertain on correct predictions and more
uncertain on incorrect ones. A practitioner could then use uncertainty to
anticipate when a model might make a mistake. Table 1(b) reports the area
under the ROC curve, or ROC-AUC of separating correctly and incorrectly
classified examples in the test set (using uncertainty as a score). A higher
ROC-AUC, closer to 1, represents a better separation. We find that ReDecNN is
competitive with MC-Dropout, even out performing it on CelebA. In Section A.3,
we include examples of images with low and high uncertainty.

Out-of-Distribution Detection  A second application of uncertainty is
detecting out-of-distribution (OOD) inputs. Given an input that is far from
the training distribution, a model is likely to make a mistake. Instead, the
model should recognize its uncertainty, and refuse to make a prediction. There
is a rich literature on detecting OOD inputs using downstream computation on a
trained model [22, 31, 30, 68, 23]. Unlike those works, we study OOD detection
using uncertainty alone.

| MNIST | FashionMNIST | CelebA
---|---|---|---
| ReDecNN | MC | Ensemble | ReDecNN | MC | Ensemble | ReDecNN | MC | Ensemble
Adversarial (FGSM) | $0.787$ | $\mathbf{0.817}$ | $0.540$ | $0.711$ | $\mathbf{0.760}$ | $0.589$ | - | - | -
OOD (MNIST) | - | - | - | $0.793$ | $\mathbf{0.864}$ | $0.501$ | $\mathbf{0.647}$ | $0.548$ | $0.591$
OOD (FashionMNIST) | $0.812$ | $\mathbf{0.888}$ | $0.534$ | - | - | - | $\mathbf{0.641}$ | $0.572$ | $0.599$
OOD (CelebA) | $0.893$ | $\mathbf{0.943}$ | $0.675$ | $0.753$ | $\mathbf{0.793}$ | $0.704$ | - | - | -
Corrupt (Mean) | $0.727$ | $\mathbf{0.785}$ | $0.566$ | $0.676$ | $\mathbf{0.691}$ | $0.606$ | $\mathbf{0.686}$ | $0.569$ | $0.632$
Corrupt (Stdev) | $0.075$ | $0.113$ | $0.087$ | $0.072$ | $0.108$ | $0.106$ | $0.038$ | $0.017$ | $0.016$
Table 2: ROC-AUC of predicting which examples are out-of-distribution (OOD).
We vary OOD examples to be adversarial, corrupted, or from a different
dataset. Standard deviation for all entries are < 0.01 over three runs.

We explore three types of OOD examples: (1) adversarial examples crafted for a
single classifier using FGSM [19], (2) examples from a different dataset (e.g.
train on MNIST and consider FashionMNIST as OOD) as done in [31], and (3) and
corrupted examples by 14 image transformations (e.g. adding pixel noise or
rotating an image), borrowed from [39].

Table 2 reports the ROC-AUC of separating inlier examples taken from the test
split of the dataset the models were trained on, and outlier examples. From
Table 2, we observe ReDecNN is just under MC-Dropout, but while MC-Dropout
achieves these results at the cost of classification accuracy, ReDecNN does
not (Table 1(a)). Furthermore, we find ReDecNN generalizes better to CelebA, a
more complex image dataset, where it outperforms MC-Dropout (and other
baselines).

OOD: Wearing a Hat
---
Method | ROC-AUC
ReDecDNN |  $\mathbf{0.604}$ (<0.01)
MC-Dropout |  $0.519$ (<0.01)
BayesNN |  $0.510$ (0.01)
Ensemble |  $0.526$ (<0.01)

OOD: Blond Hair
---
Method | ROC-AUC
ReDecDNN |  $\mathbf{0.593}$ (<0.01)
MC-Dropout |  $0.502$ (<0.01)
BayesNN |  $0.508$ (0.01)
Ensemble |  $0.509$ (<0.01)

OOD: Bald
---
Method | ROC-AUC
ReDecDNN |  $\mathbf{0.615}$ (<0.01)
MC-Dropout |  $0.502$ (<0.01)
BayesNN |  $0.508$ (<0.01)
Ensemble |  $0.503$ (<0.01)

Table 3: We hold out a group of CelebA attributes, such as those wearing a
hat, when training. We compute the ROC-AUC of labeling the held-out group as
OOD (mean and stdev. over 3 runs).

Focusing on CelebA, we study a domain-specific OOD challenge by holding out
all images with positive annotations for a single attribute as out-of-
distribution. Here, inlier and outlier examples are very similar in
distribution. We report ROC-AUCs in Table 3 for three different held-out
attributes: is the celebrity in the image “wearing a hat”, has “blonde hair”,
or “is bald”. As this is a more challenging problem, the performance is lower
than in Table 2. But while MC-Dropout, BayesNN, and a naive ensemble all
perform near chance, ReDecNN is consistently near 0.6.

Method | MNIST | Fashion | CelebA
---|---|---|---
Standard |  $1.18$ (0.10) |  $0.56$ (0.14) |  $2.75$ (0.13)
ReDecNN |  $\mathbf{0.33}$ (0.05) |  $0.21$ (0.03) |  $\mathbf{1.92}$ (0.13)
Dropout |  $0.63$ (0.09) |  $0.48$ (0.09) |  $2.51$ (0.13)
MC-Dropout |  $0.41$ (0.06) |  $\mathbf{0.13}$ (0.03) |  $2.99$ (0.15)
BayesNN |  $0.91$ (0.20) |  $0.18$ (0.03) |  $2.41$ (0.19)
Ensemble |  $1.80$ (0.12) |  $0.89$ (0.07) |  $3.12$ (0.12)
Table 4: We compare the expected calibration error (ECE) of ReDecNN to a
standard neural network with various regularization and ensembling.

Calibration  A third application of uncertainty we explore is calibration.
Standard neural networks are notoriously over-confident and are incentived in
training to predict with extreme probabilities (e.g. 0.99). A model with
useful uncertainty would make predictions with proper probabilities. Although
many calibration algorithms exist [41, 20], we want to measure how calibrated
each model is out-of-the-box to compare the model uncertainties.

Table 4 reports the expected calibration error, or ECE [40], which
approximates the difference in expectation between confidence accuracy using
discrete bins. A lower number (closer to 0) is a more calibrated model. We
observe that while a standard neural network has a relatively high ECE,
dropout, MC-Dropout, BayesNN, and ReDecNN all reduce ECE. In two of three
datasets, ReDecNN achieves the lowest ECE whereas MC-Dropout achieves the
lowest in the third. These results should be viewed in parallel to Table 1(a),
which measures “sharpness”. Otherwise, we can trivially reduce ECE by ignoring
the input.

###  5.2 Composing with Pretrained Models

Apart from self-composition, we can also compose our neural networks with off-
the-shelf pretrained models as a form of regularization or distillation.
Suppose we have a pretrained model $m$ that maps an input $x$ to a class
prediction. We can modify the ReDecNN objective as follows:

$\tilde{\mathcal{L}}(x;\theta,\phi)=\hat{\mathcal{L}}(x;\theta,\phi)+\sum_{d=1}^{D}\alpha^{d}\log
p(f_{\theta}(g_{\phi}(h^{(d-1)}_{l_{d-1}}))|m(g_{\phi}(h^{(d-1)}_{l_{d-1}})))$
(7)

where $\hat{\mathcal{L}}$ is as defined in Equation 5. The additional term
acts as a divergence, bringing the classifer’s predictions close to those of
the pretrained model. A similar edit can be made to the DecNN. We revisit
Section 5 experiments now using two pretrained models — a linear classifier or
a ResNet-18 classifier. For each, we optimize Equation 7 and compute
performance on OOD detection.

| MNIST | CelebA
---|---|---
Method | None | Linear | ResNet | None | Linear | ResNet
Accuracy |  $97.50$ (0.20) |  $97.17$ (0.12) |  $\mathbf{98.30}$ (0.24) |  $90.80$ (<0.1) |  $89.76$ (0.11) |  $\mathbf{91.28}$ (0.13)
Misclassfication |  $0.869$ (<0.01) |  $\mathbf{0.904}$ (<0.01) |  $0.742$ (<0.01) |  $0.692$ (<0.01) |  $\mathbf{0.732}$ (<0.01) |  $0.615$ (<0.01)
Adversarial |  $0.787$ (<0.01) |  $\mathbf{0.805}$ (0.01) |  $0.729$ (<0.01) | - | - | -
OOD (MNIST) | - | - | - |  $0.647$ (<0.01) |  $\mathbf{0.693}$ (<0.01) |  $0.609$ (<0.01)
OOD (Fashion) |  $0.812$ (<0.01) |  $\mathbf{0.845}$ (<0.01) |  $0.650$ (<0.01) |  $0.641$ (<0.01) |  $\mathbf{0.673}$ (<0.01) |  $0.585$ (<0.01)
OOD (CelebA) |  $0.893$ (<0.01) |  $\mathbf{0.928}$ (<0.01) |  $0.862$ (<0.01) | - | - | -
Corruption (Mean) |  $0.727$ (<0.01) |  $\mathbf{0.790}$ (<0.01) |  $0.694$ (<0.01) |  $0.686$ (<0.01) |  $\mathbf{0.708}$ (<0.01) |  $0.616$ (<0.01)
Calibration (ECE) |  $0.334$ (0.05) |  $\mathbf{0.296}$ (0.01) |  $0.385$ (0.02) |  $\mathbf{1.927}$ (0.13) |  $1.937$ (0.04) |  $3.002$ (0.10)
Table 5: OOD detection using uncertainty for ReDecNN composed with pretrained
models.

Table 5 compares the results of incorporating a linear model to a residual
network to no pretrained model (None). Using a linear model, we observe a
slight drop in test accuracy but an increase across all OOD experiments and
calibration. In fact, these new results rival or surpass MC-Dropout from Table
2. Conversely, with ResNet-18, we observe a 1 point increase in test accuracy
but notable drops in OOD and calibration results. These polarizing outcomes
show two use cases of pretrained models: using simpler models allow for
regularization that trades off accuracy for robustness while more complex
models encourage distillation, making the opposing tradeoff.

##  6 Constraining Decodable Computation

In traditional software, assertion statements are a popular way to make sure
the program is executing as expected. Ideally, we would have the same
functionality in neural networks: the ability to specify what information an
activation should not capture. This has a close relationship with fairness in
machine learning [36], where we have knowledge of sensitive or protected
attributes that should not be abused in solving a task. Our desired “assertion
statement” for neural networks should proactively block the representations
from utilizing the protected attributes in prediction.

Suppose we have access to a pretrained classifier $m$ that predicts the
assignment for a $K$-way protected feature. Then, building on Section 5.2, we
can optimize our classifier to ignore information about the protected feature
through $m$. To do this, we define the objective:

$\tilde{\mathcal{L}}(x;\theta,\phi)=\hat{\mathcal{L}}(x;\theta,\phi)+\sum_{d=1}^{D}\alpha^{d}\log
p(m(g_{\phi}(h^{(d-1)}_{l_{d-1}}))|\frac{1}{K}\mathbf{1}_{K})$ (8)

where $\mathbf{1}_{K}$ is a $K$-dimensional vector of ones. That is, Equation
8 encourages the decoded activation $g_{\phi}(h^{(d-1)}_{l_{d-1}})$, when
passed through the protected classifier $m$, to predict chance probabilities
i.e. have no information about the protected attribute. A high performing
classifier trained in this manner must have solved the prediction task without
abusing the protected attribute(s).

To test this, we extract two attributes from CelebA: “attractiveness” and
“pale skin”. As the vast majority of CelebA images contain celebrities with
pale skin, naively training a binary classifier for “attractiveness” could
lead to a discrepancy in performance between images of celebrities with and
without pale skin. Indeed, Table 6 shows a 8 point difference in F1 and a 9
point difference in average precision (AP) in classifying attractiveness
between the two groups.

Model | F1 (Pale / Not Pale) | AP (Pale / Not Pale) | ECE (Pale / Not Pale)
---|---|---|---
Standard |  $0.843/0.776$ (<0.01 / <0.01) |  $78.55/69.36$ (0.05 / 0.14) |  $2.22/2.94$ (0.01 / 0.12)
ReDecNN (Eq. 4) |  $0.845/0.780$ (<0.01 / 0.01) |  $78.28/70.34$ (0.07 / 0.14) |  $2.05/2.69$ (0.04 / 0.08)
ReDecNN (Eq. 8) |  $\mathbf{0.833/0.806}$ (<0.01 / 0.01) |  $\mathbf{77.36/76.27}$ (0.09 / 0.14) |  $\mathbf{2.06/2.12}$ (0.07 / 0.09)
Table 6: We use a pretrained model on CelebA that predicts pale skin to
optimize the activations of a second classifier to ignore skin color when
predicting attractiveness.

The second row of Table 6 provides a baseline ReDecNN optimized without
knowledge of the protected attribute. We find a similar discrepany between
groups as to a standard neural network. In the third row, we evaluate a
ReDecNN that was optimized to ignore the “pale skin” attribute using a
pretrained classifier for pale skin. Critically, we find much more balanced
F1, AP, and ECE across the two groups.

##  7 Generalizing to other Modalities

While we have focused on image classification, the proposed approach is more
general and can be extended to other modalities. We apply the same decodable
representations to speech, in particular utterance and action recognition. We
compare DecNN and ReDecNN to the same baselines, measuring accuracy and
uncertainty through similar experiments as we did for images.

Method | Acc | Misclass. | OOD | ECE
---|---|---|---|---
Standard |  $94.5$ (0.1) | - | - | -
DecNN |  $93.8$ (0.2) | - | - | -
ReDecNN |  $93.4$ (0.2) | $0.766$ | $0.705$ |  $\mathbf{0.225}$ (0.1)
MC-Dropout |  $82.1$ (0.4) | $\mathbf{0.788}$ | $0.745$ |  $0.429$ (0.2)
BayesNN |  $91.6$ (0.7) | $0.545$ | $0.518$ |  $1.039$ (0.4)
Ensemble |  $\mathbf{96.3}$ (0.1) | $0.529$ | $0.506$ |  $1.103$ (0.2)
(a) AudioMNIST

Method | Acc | Misclass. | OOD | ECE
---|---|---|---|---
Standard |  $42.4$ (0.4) | - | - | -
DecNN |  $41.4$ (0.7) | - | - | -
ReDecNN |  $41.2$ (0.5) | $\mathbf{0.642}$ | $\mathbf{0.605}$ |  $0.523$ (0.2)
MC-Dropout |  $34.5$ (0.7) | $0.629$ | $0.562$ |  $\mathbf{0.515}$ (0.2)
BayesNN |  $40.3$ (1.2) | $0.523$ | $0.500$ |  $0.918$ (0.1)
Ensemble |  $\mathbf{44.1}$ (0.1) | $0.541$ | $0.522$ |  $1.189$ (0.5)
(b) Fluent Speech Commands

Table 7: Performance on speech classification. If not specified, stdev. is
<0.01 over three test runs. Figure 4: Decoding activations to images for
Fluent spectrograms.

We utilize the AudioMNIST [5] and Fluent speech commands [33] datasets, the
former being a corpus of 60k recorded utterances of the digits 0 through 9 and
the latter a corpus of 100k recordings of 97 speakers interacting with smart-
home appliances. Each utterance demarcates one of six actions. Audio waveforms
are preprocessed to log Mel spectrograms, outputting a 32 by 32 dimensional
matrix, as done in [61]. Figure 4 shows the input on the left-most column
along with the 2nd, 4th, 6th, and 8th decoded activations for 3 random test
examples. We can observe information being dropped through the network,
although the primary structure is preserved. Table 7 reports the findings,
where like the image experiments, we find ReDecNN has strong performance on
OOD detection and calibration in return for only a small drop in performance
compared to a standard neural network.

##  8 Related Work

#### Probing Neural Networks

A number of works seek to probe neural network representations post-training
to understand the underlying computation. Some propose to do this through
distillation to simpler models [49], visualizing gradients [55, 42, 10],
estimating feature importance [1] and node importance [3], while others
leverage influence functions [28] or natural language [38]. Most relevant to
our work is an approach to understand CNN representations by inverting them
through reconstruction [34]. Unlike these approaches, we are not proposing
computation post-hoc. Rather, we are interested in optimizing neural networks
such that their computation is more easily probed by mapping representations
back to the input domain. There are some works that optimize for
interpretability [64, 51], such as regularizing a network’s predictions to
agree with a decision tree. However, these methods have difficulty
generalizing to richer modalities, such as images.

#### Model Robustness and Fairness

Out-of-distribution detection [32, 30, 68, 52, 24, 23], selective
classification [16, 17], adversarial perturbation detection [65, 14, 37, 43],
and neural network calibration [20, 67, 66, 40, 47] each have a rich subfield
of algorithms. While task-specific methods likely outperform our results, we
are excited about representations that enable uncertainty in a way that is
useful — requiring no task information nor additional algorithms. For model
fairness, our approach shares similarities to [56] where the authors maximize
the mutual information of representations with the input while minimizing the
mutual information with the protected attribute. Our approach is an
approximation of this where we treat the pretrained classifier for the
protected attribute as a proxy for mutual information with the sensitive
attribute.

#### Invertible and Ensemble Networks

The DecNN is akin to invertible generative models [9, 45] as DecNN is training
a generative model to “invert” the classifier up to a hidden layer. However,
unlike invertible flows, DecNN does not impose an explicit prior on the
activations (i.e. latent variables), and further, DecNN has a supervised
objective. Our proposed recursive network, ReDecNN, has similarities to
Bayesian neural networks [15, 6] and ensemble networks [62] — two baselines we
compare against for evaluating model uncertainty. However, we find ReDecNN to
be easier to train than doing inference over weights, and cheaper in parameter
count than naive ensembles.

##  9 Limitations and Future Work

We explored building neural networks with decodable representations, inspired
by both the difficulty of maintaining, composing, and debugging modern neural
networks, as well as the gap between traditional software and deep learning in
usability. By enforcing decoded activations to have the same “type” as the
input, we were able to visualize information embedded in an activation,
compose neural networks together in useful ways for out-of-distribution
detection and calibration, and impose real-world constraints on what
information a network uses to learn. We are optimistic that this is only the
start and look to future research to explore more applications of decodable
representations.

We discuss a few limitations. First, our approach is bottlenecked by the
quality of the generative model. Without a good reconstruction, optimization
will be intractable. However, in light recent work [59, 58, 8, 25, 57], this
is becoming less of a problem as new generative models surface. Second, we
only explored feedforward classifiers in our experiments for simplicitly. Our
approach extends naturally to residual and transformer blocks, and future
research could explore this direction. Third, optimizing recursive networks,
although parameter efficient, costs more compute as backpropagation is more
expensive. Future work could explore weight sharing across activations (not
just depth).

On a broader scale, one of the challenges of practical deep learning is
understanding model computation. High dimensional vectors with nonlinearities
are not easily interpreted by humans and as a result, deploying and
maintaining neural networks is a constant challenge. While a significant
amount research is dedicated to “white-boxing” neural networks, we view our
work as an alternative approach: rather than try to understand what a model
has learned after training, we seek to build models that bake in a way to
understand and utilize network activations in a human medium. From a different
perspective, by its very nature, decodable representations expose data through
every activation. For sensitive, personal, or secure data, using DecNN
essentially reveals that information to every actor that interacts with the
model, facing privacy concerns.

## References

  * [1] Sebastian Bach, Alexander Binder, Grégoire Montavon, Frederick Klauschen, Klaus-Robert Müller, and Wojciech Samek.  On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation.  PloS one, 10(7):e0130140, 2015.
  * [2] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.  Neural machine translation by jointly learning to align and translate.  arXiv preprint arXiv:1409.0473, 2014.
  * [3] David Bau, Bolei Zhou, Aditya Khosla, Aude Oliva, and Antonio Torralba.  Network dissection: Quantifying interpretability of deep visual representations.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6541–6549, 2017.
  * [4] Normand J Beaudry and Renato Renner.  An intuitive proof of the data processing inequality.  arXiv preprint arXiv:1107.0740, 2011.
  * [5] Sören Becker, Marcel Ackermann, Sebastian Lapuschkin, Klaus-Robert Müller, and Wojciech Samek.  Interpreting and explaining deep neural networks for classification of audio signals.  CoRR, abs/1807.03418, 2018.
  * [6] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.  Weight uncertainty in neural network.  In International Conference on Machine Learning, pages 1613–1622. PMLR, 2015.
  * [7] Thomas M Cover.  Elements of information theory.  John Wiley & Sons, 1999.
  * [8] Prafulla Dhariwal and Alex Nichol.  Diffusion models beat gans on image synthesis.  arXiv preprint arXiv:2105.05233, 2021.
  * [9] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp.  arXiv preprint arXiv:1605.08803, 2016.
  * [10] Dumitru Erhan, Yoshua Bengio, Aaron Courville, and Pascal Vincent.  Visualizing higher-layer features of a deep network.  University of Montreal, 1341(3):1, 2009.
  * [11] Piero Esposito.  Blitz - bayesian layers in torch zoo (a bayesian deep learing library for torch).  https://github.com/piEsposito/blitz-bayesian-deep-learning/, 2020\.
  * [12] et al. Falcon, WA.  Pytorch lightning.  GitHub. Note: https://github.com/PyTorchLightning/pytorch-lightning, 3, 2019.
  * [13] William Falcon and Kyunghyun Cho.  A framework for contrastive self-supervised learning and designing a new approach.  arXiv preprint arXiv:2009.00104, 2020.
  * [14] Reuben Feinman, Ryan R Curtin, Saurabh Shintre, and Andrew B Gardner.  Detecting adversarial samples from artifacts.  arXiv preprint arXiv:1703.00410, 2017.
  * [15] Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In international conference on machine learning, pages 1050–1059. PMLR, 2016.
  * [16] Yonatan Geifman and Ran El-Yaniv.  Selective classification for deep neural networks.  arXiv preprint arXiv:1705.08500, 2017.
  * [17] Yonatan Geifman and Ran El-Yaniv.  Selectivenet: A deep neural network with an integrated reject option.  In International Conference on Machine Learning, pages 2151–2159. PMLR, 2019.
  * [18] Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial networks.  arXiv preprint arXiv:1406.2661, 2014.
  * [19] Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy.  Explaining and harnessing adversarial examples.  arXiv preprint arXiv:1412.6572, 2014.
  * [20] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger.  On calibration of modern neural networks.  In International Conference on Machine Learning, pages 1321–1330. PMLR, 2017.
  * [21] 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.
  * [22] Dan Hendrycks and Kevin Gimpel.  A baseline for detecting misclassified and out-of-distribution examples in neural networks.  arXiv preprint arXiv:1610.02136, 2016.
  * [23] Dan Hendrycks, Mantas Mazeika, and Thomas Dietterich.  Deep anomaly detection with outlier exposure.  arXiv preprint arXiv:1812.04606, 2018.
  * [24] Dan Hendrycks, Mantas Mazeika, Saurav Kadavath, and Dawn Song.  Using self-supervised learning can improve model robustness and uncertainty.  In Advances in Neural Information Processing Systems, pages 15663–15674, 2019.
  * [25] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  arXiv preprint arXiv:2006.11239, 2020.
  * [26] Andrej Karpathy, Nov 2017.
  * [27] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114, 2013.
  * [28] Pang Wei Koh and Percy Liang.  Understanding black-box predictions via influence functions.  In International Conference on Machine Learning, pages 1885–1894. PMLR, 2017.
  * [29] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  Advances in neural information processing systems, 25:1097–1105, 2012.
  * [30] Kimin Lee, Kibok Lee, Honglak Lee, and Jinwoo Shin.  A simple unified framework for detecting out-of-distribution samples and adversarial attacks.  In Advances in Neural Information Processing Systems, pages 7167–7177, 2018.
  * [31] Shiyu Liang, Yixuan Li, and R Srikant.  Enhancing the reliability of out-of-distribution image detection in neural networks.  In International Conference on Learning Representations, 2018.
  * [32] Shiyu Liang, Yixuan Li, and Rayadurgam Srikant.  Enhancing the reliability of out-of-distribution image detection in neural networks.  arXiv preprint arXiv:1706.02690, 2017.
  * [33] Loren Lugosch, Mirco Ravanelli, Patrick Ignoto, Vikrant Singh Tomar, and Yoshua Bengio.  Speech model pre-training for end-to-end spoken language understanding.  arXiv preprint arXiv:1904.03670, 2019.
  * [34] Aravindh Mahendran and Andrea Vedaldi.  Understanding deep image representations by inverting them.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5188–5196, 2015.
  * [35] Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey.  Adversarial autoencoders.  arXiv preprint arXiv:1511.05644, 2015.
  * [36] Ninareh Mehrabi, Fred Morstatter, Nripsuta Saxena, Kristina Lerman, and Aram Galstyan.  A survey on bias and fairness in machine learning.  arXiv preprint arXiv:1908.09635, 2019.
  * [37] Jan Hendrik Metzen, Tim Genewein, Volker Fischer, and Bastian Bischoff.  On detecting adversarial perturbations.  arXiv preprint arXiv:1702.04267, 2017.
  * [38] Jesse Mu and Jacob Andreas.  Compositional explanations of neurons.  arXiv preprint arXiv:2006.14032, 2020.
  * [39] Norman Mu and Justin Gilmer.  Mnist-c: A robustness benchmark for computer vision.  arXiv preprint arXiv:1906.02337, 2019.
  * [40] Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht.  Obtaining well calibrated probabilities using bayesian binning.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015.
  * [41] Alexandru Niculescu-Mizil and Rich Caruana.  Predicting good probabilities with supervised learning.  In Proceedings of the 22nd international conference on Machine learning, pages 625–632, 2005.
  * [42] Chris Olah, Alexander Mordvintsev, and Ludwig Schubert.  Feature visualization.  Distill, 2017.  https://distill.pub/2017/feature-visualization.
  * [43] Tianyu Pang, Chao Du, Yinpeng Dong, and Jun Zhu.  Towards robust detection of adversarial examples.  In NeurIPS, 2018.
  * [44] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference.  arXiv preprint arXiv:1912.02762, 2019.
  * [45] George Papamakarios, Theo Pavlakou, and Iain Murray.  Masked autoregressive flow for density estimation.  arXiv preprint arXiv:1705.07057, 2017.
  * [46] Guim Perarnau, Joost Van De Weijer, Bogdan Raducanu, and Jose M Álvarez.  Invertible conditional gans for image editing.  arXiv preprint arXiv:1611.06355, 2016.
  * [47] John Platt et al.  Probabilistic outputs for support vector machines and comparisons to regularized likelihood methods.  Advances in large margin classifiers, 10(3):61–74, 1999.
  * [48] Christopher Ré.  Software 2.0 and snorkel: beyond hand-labeled data.  In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 2876–2876, 2018.
  * [49] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin.  " why should i trust you?" explaining the predictions of any classifier.  In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pages 1135–1144, 2016.
  * [50] Robin Rombach, Patrick Esser, and Bjorn Ommer.  Network-to-network translation with conditional invertible neural networks.  Advances in Neural Information Processing Systems, 33, 2020.
  * [51] Andrew Ross and Finale Doshi-Velez.  Improving the adversarial robustness and interpretability of deep neural networks by regularizing their input gradients.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
  * [52] Chandramouli Shama Sastry and Sageev Oore.  Detecting out-of-distribution examples with in-distribution examples and gram matrices.  arXiv preprint arXiv:1912.12510, 2019.
  * [53] Jürgen Schmidhuber.  Deep learning in neural networks: An overview.  Neural networks, 61:85–117, 2015.
  * [54] David Sculley, Gary Holt, Daniel Golovin, Eugene Davydov, Todd Phillips, Dietmar Ebner, Vinay Chaudhary, Michael Young, Jean-Francois Crespo, and Dan Dennison.  Hidden technical debt in machine learning systems.  Advances in neural information processing systems, 28:2503–2511, 2015.
  * [55] 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.
  * [56] Jiaming Song, Pratyusha Kalluri, Aditya Grover, Shengjia Zhao, and Stefano Ermon.  Learning controllable fair representations.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2164–2173. PMLR, 2019.
  * [57] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  arXiv preprint arXiv:2010.02502, 2020.
  * [58] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  arXiv preprint arXiv:1907.05600, 2019.
  * [59] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  arXiv preprint arXiv:2011.13456, 2020.
  * [60] Ilya Sutskever, Oriol Vinyals, and Quoc V Le.  Sequence to sequence learning with neural networks.  arXiv preprint arXiv:1409.3215, 2014.
  * [61] Alex Tamkin, Mike Wu, and Noah Goodman.  Viewmaker networks: Learning views for unsupervised representation learning.  arXiv preprint arXiv:2010.07432, 2020.
  * [62] Sean Tao.  Deep neural network ensembles.  In International Conference on Machine Learning, Optimization, and Data Science, pages 1–12. Springer, 2019.
  * [63] Pete Warden, Nov 2017.
  * [64] Mike Wu, Michael Hughes, Sonali Parbhoo, Maurizio Zazzi, Volker Roth, and Finale Doshi-Velez.  Beyond sparsity: Tree regularization of deep models for interpretability.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
  * [65] Weilin Xu, David Evans, and Yanjun Qi.  Feature squeezing: Detecting adversarial examples in deep neural networks.  arXiv preprint arXiv:1704.01155, 2017.
  * [66] Bianca Zadrozny and Charles Elkan.  Obtaining calibrated probability estimates from decision trees and naive bayesian classifiers.  In Icml, volume 1, pages 609–616. Citeseer, 2001.
  * [67] Bianca Zadrozny and Charles Elkan.  Transforming classifier scores into accurate multiclass probability estimates.  In Proceedings of the eighth ACM SIGKDD international conference on Knowledge discovery and data mining, pages 694–699, 2002.
  * [68] Ev Zisselman and Aviv Tamar.  Deep residual flow for out of distribution detection.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 13994–14003, 2020.

##  Appendix A Appendix

###  A.1 Training Details

In our experiments, the classifier $f_{\theta}$ is a 8-layer MLP with 128
hidden dimensions per layer. We assume the same dimensionality per layer for
simplicity but the approach easily supports MLPs with variable hidden
dimensionalities: one can train seperate generative models or use the
technique from [50] where smaller dimensionalities are padded. The generative
model $g_{\theta}$ is a ResNet18 decoder that maps a 128 dimensional vector to
a 3 by 32 by 32 pixel image for RGB images and 1 by 32 by 32 for grayscale
images and speech spectrograms. We repurpose the implementation from the
PyTorch Lightning Bolts [13] repository:
https://github.com/PyTorchLightning/lightning-
bolts/blob/master/pl_bolts/models/autoencoders/components.py (Apache 2.0
License). (All input images are reshaped to 32 by 32 pixels. No additional
image transformations were used in training nor evaluation.) All models are
trained for 50 epochs with batch size 128, Adam optimizer with learning rate
1e-4. We set $\beta=10$ in all cases. For all ReDecNN models, we use a max
depth $D=8$ and set $\alpha=0.5$ in all cases. To make for a fair comparison,
for all naive ensemble networks, we train 8 copies of a neural network with
different weight initializations. For baseline models using dropout and MC-
Dropout, we use 0.5 dropout probability. For weight uncertainty (i.e.
BayesNN), we use the Blitz library [11], https://github.com/piEsposito/blitz-
bayesian-deep-learning (GNU V3 License). All pretrained models used were
ResNet-18 classifiers. Our ResNet classifier implementation was adapted from
torchvision. For each of the architectures, to support 32x32 images (which are
smaller than the standard), we replace the first 7x7 convolutional layer with
a 3x3 convolutional layer and reduce the first max pooling layer. We did not
use their pretrained weights from torchvision and instead trained them
ourselves on MNIST, FashionMNIST, and CelebA from scratch (again 50 epochs).
All models were trained on a single Titan Xp GPU with 4 CPUs for data loading.
An average model takes 1 hours to train for MNIST and FashionMNIST and 4 hours
for CelebA, and speech experiments. In our experiments, we compute uncertainty
using 30 random samples (e.g. for ReDecNN, this is 30 paths; for MC-Dropout,
this is 30 different dropout configurations, etc.). ROC-AUC computation is
done through scikit-learn (BSD License). In CelebA, only a subset of 18
attributes are used, chosen for visual distinctiveness as done in [46]. In
experiments, we utilize the PyTorch Lightning framework [12] (Apache 2.0
License) and Weights and Biases (MIT License) for tracking. For speech
experiments, the AudioMNIST dataset is found at
https://github.com/soerenab/AudioMNIST (MIT License) and the Fluent dataset
can be downloaded at https://fluent.ai/fluent-speech-commands-a-dataset-for-
spoken-language-understanding-research/ (academic license). All speech
spectrograms are normalized using a mean and standard deviation computed from
the training dataset. We use torchaudio and librosa to efficiently compute Mel
spectrograms.

###  A.2 Extended OOD Results

We provide a more thorough breakdown of performance for the 14 different image
corruptions. In the main paper, we only show the average performance over all
corruptions.

| MNIST | FashionMNIST | CelebA
---|---|---|---
| ReDecNN | MC | Ensemble | ReDecNN | MC | Ensemble | ReDecNN | MC | Ensemble
Adversarial (FGSM) | $0.787$ | $\mathbf{0.817}$ | $0.540$ | $0.711$ | $\mathbf{0.760}$ | $0.589$ | todo | todo | todo
OOD (MNIST) | - | - | - | $0.793$ | $\mathbf{0.864}$ | $0.501$ | $\mathbf{0.647}$ | $0.548$ | $0.591$
OOD (FashionMNIST) | $0.812$ | $\mathbf{0.888}$ | $0.534$ | - | - | - | $\mathbf{0.641}$ | $0.572$ | $0.599$
OOD (CelebA) | $0.893$ | $\mathbf{0.943}$ | $0.675$ | $0.753$ | $\mathbf{0.793}$ | $0.704$ | - | - | -
Corrupt (Brightness) | $0.851$ | $0.955$ | $0.702$ | $0.732$ | $0.667$ | $0.816$ | $0.653$ | $0.558$ | $0.623$
Corrupt (Dotted Line) | $0.672$ | $0.684$ | $0.546$ | $0.602$ | $0.589$ | $0.535$ | $0.647$ | $0.559$ | $0.626$
Corrupt (Glass Blur) | $0.664$ | $0.718$ | $0.515$ | $0.617$ | $0.584$ | $0.525$ | $0.641$ | $0.596$ | $0.652$
Corrupt (Impulse Noise) | $0.702$ | $0.848$ | $0.551$ | $0.702$ | $0.746$ | $0.642$ | $0.683$ | $0.555$ | $0.640$
Corrupt (Rotate) | $0.689$ | $0.714$ | $0.523$ | $0.754$ | $0.839$ | $0.600$ | $0.645$ | $0.585$ | $0.652$
Corrupt (Shear) | $0.685$ | $0.661$ | $0.502$ | $0.611$ | $0.809$ | $0.543$ | $0.750$ | $0.581$ | $0.636$
Corrupt (Spatter) | $0.647$ | $0.683$ | $0.502$ | $0.617$ | $0.615$ | $0.534$ | $0.715$ | $0.569$ | $0.641$
Corrupt (Translate) | $0.809$ | $0.889$ | $0.511$ | $0.725$ | $0.826$ | $0.555$ | $0.673$ | $0.577$ | $0.610$
Corrupt (Canny Edges) | $0.696$ | $0.824$ | $0.521$ | $0.703$ | $0.807$ | $0.571$ | $0.649$ | $0.569$ | $0.650$
Corrupt (Fog) | $0.868$ | $0.940$ | $0.732$ | $0.748$ | $0.670$ | $0.810$ | $0.729$ | $0.608$ | $0.643$
Corrupt (Scale) | $0.768$ | $0.846$ | $0.535$ | $0.808$ | $0.805$ | $0.521$ | $0.750$ | $0.562$ | $0.632$
Corrupt (Shot Noise) | $0.617$ | $0.559$ | $0.526$ | $0.541$ | $0.537$ | $0.528$ | $0.685$ | $0.542$ | $0.592$
Corrupt (Stripe) | $0.804$ | $0.892$ | $0.755$ | $0.684$ | $0.522$ | $0.772$ | $0.674$ | $0.563$ | $0.628$
Corrupt (Zigzag) | $0.709$ | $0.785$ | $0.514$ | $0.626$ | $0.658$ | $0.532$ | $0.714$ | $0.555$ | $0.627$
Corrupt (Mean) | $0.727$ | $\mathbf{0.785}$ | $0.566$ | $0.676$ | $\mathbf{0.691}$ | $0.606$ | $\mathbf{0.686}$ | $0.569$ | $0.632$
Corrupt (Stdev) | $0.075$ | $0.113$ | $0.087$ | $0.072$ | $0.108$ | $0.106$ | $0.038$ | $0.017$ | $0.016$
Table 8: We report the ROC-AUC of predicting which examples are out-of-
distribution (OOD) using uncertainty. We vary OOD examples to be adversarial,
corrupted, or taken from a different dataset.

###  A.3 Visualizing OOD Examples

Figure 5 shows one image from each class with low and high uncertainty. We
observe that high uncertainty images are less prototypical. For example, the
high uncertainty digits in Figure 5d have accented curvature, and the high
uncertainty images of celebrities in Figure 5f, although annotated as not
bald, are wearing hats or have thin hairlines and exposed foreheads.

(a) MNIST (certain)

(b) FashionMNIST (certain)

(c) Not Bald (certain)

(d) MNIST (uncertain)

(e) FashionMNIST (uncertain)

(f) Not Bald (uncertain)

Figure 5: Randomly sampled images with low and high uncertainty. Images with
high uncertainty appear less “prototypical”. Figure 5(c,h) shows images from
CelebA annotated as “not bald”. We observe images that the model is uncertain
about depict indviduals with a thin hair line, exposed foreheads, or wearing
helmets.

###  A.4 Probing Decoded Activations

We include the full figure (with MNIST results) of Figure 3 in the main text.
The top row shows decoded activations for correctly-classified examples
whereas the bottom row shows mis-classified examples. For correctly-classified
examples, we observe that the decoded activations tends towards class
prototypes. Decoded images from the 7th or 8th activation lose details (e.g.
patterns, logos, or shoe straps disappear from clothing). The top row of
Figure 6b presents a good example: the tank-top (a rare form of the t-shirt)
is projected to a more prototypical t-shirt. We also find that mis-classified
examples morph into images of the incorrect class. The seventh row of Figure
6e morphs the two articles of clothing that compose the dress to build a
shirt. The seventh row of Figure 6f show the celebrity’s hat transform into a
background color.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 6: Visualizing activations: the top row shows decoded activations when
the model correctly classified examples; the bottom row shows
misclassifications.


# Computational ghost imaging for transmission electron microscopy

Akhil Kallepalli  1,* Lorenzo Viani  2,3 Daan Stellinga  4 Enzo Rotunno  3,**
Ming-Jie Sun  5 Richard Bowman  6 Paolo Rosi  2 Stefano Frabboni  2 Roberto
Balboni  7 Andrea Migliori  7 Vincenzo Grillo  3 and Miles Padgett1  1School
of Physics and Astronomy, University of Glasgow, Glasgow, G12 8QQ, UK
2 Department of Physics, IT and Mathematics, University of Modena and Reggio
Emilia
3 CNR-NANO Via G. Campi 213/a, I-41125 Modena, Italy
4 Faculty of Science and Technology, University of Twente, The Netherlands
5 School of Instrumentation and Optoelectronic Engineering, Beihang
University, Beijing 100191, China
6 Department of Physics, University of Bath, BA2 7AY, UK
7 CNR-IMM Via Gobetti 101, 40129 Bologna, Italy
*<EMAIL_ADDRESS>**<EMAIL_ADDRESS>
###### Abstract

While transmission electron microscopes (TEM) can achieve a much higher
resolution than optical microscopes, they face challenges of damage to samples
during the high energy processes involved. Here, we explore using
computational ghost imaging techniques in electron microscopy to reduce the
total required intensity. The technological lack of the equivalent high-
resolution, optical spatial light modulator for electrons means that a
different approach needs to be pursued. To this end, we show a beam shaping
technique based on the use of a distribution of electrically charged metal
needles to structure the beam, alongside a novel reconstruction method to
handle the resulting highly non-orthogonal patterns. Second, we illustrate the
application of this ghost imaging approach in electron microscopy. To test the
full extent of the capabilities of this technique, we realised an analogous
optical setup method. In both regimes, the ability to reduce the amount of
total illumination intensity is evident in comparison to raster scanning.

††journal: osajournal††articletype: Research Article

##  1 Introduction

Scanning transmission electron microscopes (STEM) achieve high resolution
imaging by raster scanning a focused beam of electrons over the sample and
measuring the transmission to form an image. Raster scanning is a simple
approach to creating an image in the absence of complicated detector arrays.
An alternative approach that has been pioneered in the optical regime has been
measuring the transmission of the sample when illuminated with more
complicated spatial patterns. This approach, often referred to as classical
computational ghost imaging, inverts knowledge of the known patterns and their
measured transmission to reveal the image. Furthermore, the inversion can be
constrained using known image properties (e.g. sparsity of spatial
frequencies) such that the image can be estimated from fewer measurements than
would have been required by a raster scan; a form of compressed sensing [1,
2].

The spatial structuring of optical beams has been extended to matter waves,
primarily electrons. This structuring of electron beams paved the way for the
production of vortex electron beams [3, 4, 5], self-accelerating beams [6],
non-diffracting beams [7, 8, 9] and orbital angular momentum analysers [10,
11]. In the case of electron beams, an arbitrary way of flexibly spatially
programming a beam is problematic and existing methods for the dynamic shaping
of electron beams are limited. However, recent progress could inspire to build
a tuneable phase plate based on a set of needle-shaped electrodes producing a
continuous phase landscape. In this work, we use electron beam shaping
technologies to demonstrate that the technique can be applied in electron
microscopy and show it is possible to create complicated far-field patterns
that when combined with a novel inversion algorithm produce images comparable
to those produced by raster scanning but using much lower peak intensities.
The technique is firstly developed with a light-optics analogue toy experiment
and then deployed in the electron microscope real measurement.

The potential advantages of the computational ghost imaging approach are two-
fold. First, that single element detectors usually have better
characteristics, e.g. for light an extended range of operating wavelengths,
for electrons a larger quantum efficiency. Second, the flexibility of the
illuminating patterns make them potentially more suitable for imaging
optimisation by compressed sensing algorithms, reducing the number of required
measurements [12, 13, 14]. This ability to image with fewer measurements
potentially increases the frame rate and/or reduces the exposure of
potentially delicate samples to the illumination. Using these computational
ghost imaging techniques, novel imaging systems have been developed that
include imaging at terahertz wavelengths [15], time of flight imaging [16,
17], and imaging sensitive to fluorescent lifetime [18].

In addition to optical wavelengths, electron beams can also be used for
precision imaging. Since the time-independent wave equation for the light and
electrons are perfectly analogous [19], components like lenses for electron
beams are defined with a similar nomenclature as for light. However, optical
elements for light are inherently easier to build and spatial light modulators
(SLM) that shape the optical beams in a programmable, and near arbitrary way,
do not explicitly exist for electrons. Although single-element detector-based
imaging has previously been introduced for specialised electron imaging
systems [20], the barrier to more general use is absence of SLM-like
technologies.

In the context of electron beams, holographic plates [19, 21] laid the
foundation for applications related to the shaping of electron beams. However,
these are passive elements and cannot be dynamically addressed like optical
SLM. A most welcomed recent development has been the introduction of active
optical elements or programmable phase plates [22]. These are based on simple
electrostatic elements to be positioned along the beam path. Two technological
solutions are currently available: the research group led by Verbeeck has been
developing devices based on Einzel lenses. These lenses work like the
individual pixels of an optical spatial light modulator but the number of
addressable points are few in comparison the the optical devices [22]. An
additional challenge for these lenses arises from the difficulty to bias and
drive many connections from outside the microscope column, which necessarily
limits the useful fill fraction of such a pixel-based approach. As an
alternative, other groups are using separated electrodes to define specific
phase landscapes across a wider area. This electrode-based approach has the
benefit that the beam itself only interacts with the electric field created
rather than the modulator components themselves, avoiding the main problems of
the other approach. However, whereas the solution is defined everywhere the
phases $\phi$ are limited by the harmonic condition $\nabla^{2}\phi=0$
everywhere except on the electrodes [23]. In other words, the variations in
possible phase modulation is far more restricted.

The most trivial electrode used thus far is a simple metallic needle [24]. The
definition of the electrostatic field around a charged microtip has been
studied in detail [25, 26], wherein charge applied to these needles induces a
phase shift around their tips to the electron beam. Interestingly, the phase
shift generated by a metallic needle was used as key to implement an orbital
angular momentum (OAM) sorter [27]. More importantly, multiple metallic
needles have been demonstrate to be capable of producing complex tunable
caustic phenomena in the far field[28]. This current work uses an arrangement
of six such needles.

##  2 Computational ghost imaging in TEM

We propose a design of the experimental setup for TEM implementation of
Computational Ghost Imaging. Figure 1 shows a schematic representing the key
steps of the experiment.

Figure 1: A TEM sample holder carrying our needle-based beam-shaping device is
inserted in the column at the condenser system level. Dedicated electronics
are used to randomly select six random biases and apply them to the needles.
The electron beam crosses the potential distribution and is imaged on the
sample. Then, the total transmitted intensity is acquired by a single-pixel
detector, whose output value is acquired by the electronics and recorded
together with the six biases. After that, the patterns used for the
measurement are calculated starting from the six biases. Finally, the
combination of the single-pixel measurements and the calculated patterns
through the CGI algorithm yields the ghost image of the sample.

In complete analogy with the optical realisation of ghost imaging, the
technique requires a way to reliably produce a form of structured
illumination. As mentioned before, in optical systems this is possible thanks
to SLM technology, but there is not any equivalent form of beam-shaping
technology available for electrons. Thus, we developed a device that uses six
electrically biased needles to generate a potential distribution, as shown in
figure 2. When the device is inserted into the TEM column the potential
distribution interacts with the electron beam and applies a position-dependent
phase shift to it (see figure 2(c)).

Figure 2: SEM images (a) and (b) illustrate the six-needle device and its
active region. The active region is limited by a 10 $\mu$m diameter aperture
and a resulting simulation of a typical phase profile produced by these
needles is shown in (c). The white bars in (c) show the positions of the
needles. (d) shows three TEM experimental images of the intensity distribution
in the far-field produced by the device.

The beam shaping device produces complex patterns in far field observation,
with a shape dependent on the six biases applied to the needles. Thus, our
device allows us to control the illumination, albeit with far fewer degrees of
freedom than those available with a SLM. Some examples are shown in figure
2(d). Notice how there can be a single intensity peak or more than one with
variable separation, from superimposed to very far apart.

A critical limitation of the patterns produced by the beam-shaping device is
that they are highly non-orthogonal. In stark contrast, SLMs can produce
completely orthogonal pattern sets, like the Hadamard basis, which are
inherently better for CGI. To this end, we also developed a modified CGI
algorithm that adds orthogonalisation to the image reconstruction process.

##  3 Novel algorithm for ghost imaging with non-orthogonal pattern basis

The choice of mask patterns is key to the performance of the system as it
defines the resolving capability of the system and the applicability of the
reconstruction algorithm. The field of ghost imaging has seen an evolution in
both experimental setups and algorithms [29, 30]. Although the vast majority
of computational ghost imaging and the related work on single pixel cameras
use a spatial modulator positioned in the image plane of the object, the
original work by Shapiro [31] recognised that this need not be the case.
Rather, given any specific modulation, the resulting form of the light field
in the plane of the object could be calculated using numerical beam
propagation techniques enabling both lens free and far field implementations
[32].

Another important aspect of ghost imaging are the algorithms used to recover
an image from the measurements and knowledge of the patterns used. This
problem can be cast into a simple matrix form:

$A\vec{x}=\vec{b},$ (1)

where $A$ is a matrix containing the patterns $\vec{p_{i}}^{\top}$ as rows,
$\vec{x}$ represents the (discretised) scene, and $\vec{b}$ is the set of
measurements $b_{i}$. The methods used to retrieve $\vec{x}$ and thereby an
image fall roughly into two categories, offline or online [33]. Offline
approaches infer the image from the final complete set of patterns and
corresponding measurements. The majority of well known methods for solving
linear systems fall in this category, e.g. the matrix (pseudo)inverse,
Newton’s method, etc. Conversely, online approaches update a best guess of the
image with every new measurement and light pattern. The most common online
algorithm involves a straightforward weighted sum of the patterns, where the
weights are adjusted based on a statistical interpretation of the
measurements:

$\vec{g}_{i}=\vec{g}_{i-1}+(b_{i}-\langle b\rangle_{i})\vec{p}_{i}.$ (2)

Here $\vec{g}_{i}$ is the best estimate after $i$ patterns and $\langle
b\rangle_{i}$ indicates the mean detector value up to the current pattern, and
$\vec{g}_{0}=\vec{0}$. This algorithm, which we will refer to as traditional
ghost imaging (TGI), has been extensively explored with different setups [34,
35].

The primary advantage of online reconstruction algorithms lies in the low
computational overhead and robustness to under-sampling. In the absence of
additional constraints, to achieve a reasonable spatial resolution, the number
of illumination patterns required is typically equal to the number of pixels
in the image. This quickly makes the matrix $A$ from equation 1 very large,
leading to high memory requirements and long computation times for offline
methods. Online methods only need to store the current best guess of $\vec{x}$
and the current pattern, can run in parallel with the measurement, and only
ever operate on vectors of limited size. Moreover, TGI is by its very nature
impervious to problems with under-sampling and can show approximations to the
correct image with comparatively only a handful of measurements.

Both classes of reconstruction method work significantly better when used with
orthogonal pattern sets, for example Hadamard matrices [36, 37, 38], where
each mask probes a subset of spatial frequencies. For offline methods this is
equivalent to ensuring full row rank, which increases the stability of the
matrix inversion results. For TGI, it ensures that each pattern’s spatial
frequencies are counted democratically in the mean correction term. Either
way, in the absence of noise, the result approaches the true image when the
number of orthogonal patterns equals the number of pixels.

Figure 3: Light from a HeNe at 633 nm is polarised by a polarising beam
splitter (PBS), then spatially filtered and expanded to fill the aperture of
an SLM using an aperture (P), a 50 mm focal length lens (L1), precision
pinhole (P50μ) and 400 mm focal length lens (L2). The SLM displays a
superposition of the six-needles phase masks and a carrier diffraction
grating. The SLM is optimised to project the majority of the modulated light
into the first diffraction order, which is selected using a 250 mm focal
length lens (L3) and aperture (P) placed in the Fourier plane of the SLM. This
light is propagated into the far-field through a combination of lenses (250 mm
focal length lens (L4) and 400 mm focal length lens (L5)). The beam is
propagated into two separate arms using beamsplitter (BS). One arm uses a
power meter to measure the laser power and the second arm consists of the
target (T). The photomultiplier tube (Thorlabs, PMM02) measures the
transmitted light after interaction with the target. Author Figure adopted
from Kallepalli et al. [39]

If the mask patterns are implemented using a pixelated spatial light
modulator, specifying their design to be orthogonal is straightforward (e.g.
Hadamard, or similar). If, however, the patterns are produced in other ways,
e.g. using natural scattering or with the highly limited modulators available
in a TEM, then orthogonality between the patterns can no longer be enforced
and the standard reconstruction algorithms no longer converge to an accurate
image. The TGI algorithm in particular will overemphasise the overlapping
parts of the patterns. To resolve this issue, we report an alternative
approach that is computationally efficient and makes optimal use of patterns
regardless of their orthogonality. This alternative method, which we will call
orthogonalised ghost imaging (OGI), takes inspiration from the Kaczmarz
algorithm for solving linear systems [40, 41, 42, 43, 44].

OGI essentially takes the projection of each new pattern onto the current best
pattern estimate as an extra correction to the mean detector value used in
TGI:

$\vec{g}_{i}=\vec{g}_{i-1}+\left[\left(b_{i}-\langle
b\rangle_{i}\right)-\left(c_{i}-\langle
c\rangle_{i}\right)\right]\vec{p}_{i},$ (3)

where $c_{i}=\vec{p}_{i}\cdot\vec{g}_{i-1}$ is the predicted signal under the
assumption that the previous best estimate is the correct reconstruction. This
has the effect of mitigating double counting stemming from non-orthogonality
and emphasising actual new information present in the measurement. Note that
care needs to be taken such that both the measured and predicted values are
correctly scaled to each other such that the reconstructed image approaches
the true image. This is equivalent to ensuring that the quantity
$[(b_{i}-\langle b\rangle_{i})-(c_{i}-\langle c\rangle_{i})]$ tends to zero.
The approach we adopt here to ensure this scaling is to iteratively normalise
the signals such that $\langle c\rangle=\langle b\rangle$ and
$\sigma_{c}=\sigma_{b}$.

##  4 Optical ghost imaging

The modulation device investigated in this article is inspired from the
charged microtips approach, designed by Matteucci et al. [25]. The needles, to
which random voltages are applied, create a spatially structured phase change
in the transmitted electron beam, leading to a complicated, spatially
structured, point spread function in the far-field that illuminates the
sample. An example is shown in figure 2. In our experiment, the equivalent
phase change to that induced by the charged needles is, in the optical regime,
created using a liquid crystal spatial light modulator (LC-SLM).

Our analogous optical experimental set up [39] for demonstrating the six-
needles approach to computational ghost imaging is shown in Figure 3. The
expanded output beam from a HeNe laser is shaped using randomly switched six-
needles that generate phase masks. These phase masks are displayed on the
spatial light modulator (SLM), structuring the laser beam and propagating
complex intensity patterns into the far-field.

Reflective SLMs have the advantage of faster response rates and a higher fill
factor than those devices based upon transmission. However, these reflective
devices are susceptible to manufacturing limitations that result in deviations
from flatness of the addressed surface, resulting in aberrations such as
astigmatism and defocus. Correction of these distortions is key to accurately
manipulating the phase of the incoming light beam. [45, 46]. Our approach to
aberration correction, adapted from Turtaev et al. [47], uses a diagonal
blazed phase grating with a predetermined periodicity. The aberration-
corrected SLM operates in a diffractive mode, and therefore, the masks are
combined with a phase-grating to diffract the desired pattern to the first-
order which is selected by a pinhole in the Fourier-plane. The chosen order is
then propagated through the lens system and into two optical paths (Figure 3):
(a) laser power meter to monitor the laser power for any variations and (b)
ghost imaging system (photomultiplier tube, PMT). The PMT is positioned to
collect the light after transmission by the object. Based on the known voltage
applied to each of the needles and the phase change this creates, the
calculation of the far-field intensity distribution and the measured signal
are used to reconstruct the image, using the various algorithms described
above.

We benchmark the performance of our six-needle approach against a traditional
raster scan. Raster scanning is achieved by displaying the appropriate series
of linear phase gratings on the SLM to produce a scanning spot in the far-
field. The experiment results are presented in figure 4, which gives a visual
representation of the TGI and OGI algorithms performance compared to that of a
raster scan. Clearly, the non-orthogonal nature of the six-needle pattern sets
has significantly degraded the reconstructed image when using the TGI
algorithm. By comparison we see that the OGI algorithm allows reconstruction
an image with the same resolution as a traditional raster scan. As a more
quantitative comparison, the spatial frequency and resolution is assessed
using a radial contrast function (Figure 5). With this number of measurements
TGI has essentially converged to its final resolution, with any further
improvements primarily consisting of suppressing the noise. OGI on the other
hand has not yet converged, but already significantly outperforms TGI for the
non-orthogonal pattern set created by the six-needle modulator.

Figure 4: Our approach illustrates the viability of using TEM-inspired phase
masks for image reconstruction using a new approach to ghost imaging known as
orthogonalised ghost imaging (OGI). This approach is ideal for cases such as
six-needles phase mask sets, where the pattern set is not orthogonal. The
comparison shows (a) randomised raster scan using traditional ghost imaging
(TGI), (b) six-needle phase masks using TGI and (c) six-needle phase masks
using the proposed OGI technique. In each case, the construction was done
using a set of 50000 patterns and the reconstructed images are of 512 $\times$
512 pixel resolution. The binary amplitude target used is a laser cut 8 spoke
shape. Figure 5: Radial contrast functions associated with the images shown in
figure 4, quantifying the resulting reconstructions using OGI-six needles,
TGI-six needles and TGI-raster scans, each using 50000 patterns. The six-
needles patterns are non-orthogonal by nature.

Although the overall optical energy used in the six-needle OGI and the
traditional raster scan is the same we note that the peak intensity in six-
needle approach is orders of magnitude lower. This reduction would be
extremely important in situations where the sample can be damaged by the
illumination, especially if that damage mechanism was non-linear with respect
to intensity.

##  5 TEM computational ghost imaging

Having proven the applicability of computational ghost imaging with the
optical analogy, we attempted an experimental demonstration with electrons.
The experiments were carried out by mounting the modulation device in a biased
I-V TEM sample holder. The sample to be imaged is a Siemens star target with 5
spokes, located in the Selected Area plane of the microscope. It was
fabricated on a commercially available 200 $nm$ thick Si3N4 membrane, covered
with $\sim$150 $nm$ of gold, so that it would act as an amplitude mask. The
total transmitted intensity is collected by the annular dark-field (ADF)
detector, set off-axis to acquire the whole transmission. The output value of
the detector is read by the control software and synchronised with the biases
from the modulation device.

Figure 6: Comparison of (a) a TEM image of the target and the images
reconstructed with TGI (b) and OGI (c) algorithms. The target is a 3 $\mu$m
10-spoke Siemens star, acquired on a Gatan MSC794 CCD camera.

The acquisition comprised 11500 different illumination patterns. The patterns
produced by the beam shaping device were not directly recorded, but instead
they were recovered using a simulation algorithm after the measurement. The
simulation algorithm predicts the illumination patterns used during the
experiment starting from the six biases applied to the needles of the beam-
shaping device. Figure 6(a) displays the target and the ghost images obtained
with both the traditional (b) and orthogonalised (c) algorithms, starting from
the same data. TGI is clearly outperformed by our OGI formulation (Eq. 3), as
pointed out by the differences between the two reconstructions. The OGI image
exhibits higher resolution and better contrast, resulting in an overall
sharper image.

##  6 Conclusions

We have shown a successful proof-of-principle ghost imaging scheme tested in
both optical and electron microscopy regimes. In the optical regime, we
combine non-orthogonal intensity distribution patterns with a novel ghost
imaging algorithm for image reconstruction. Due to the non-orthogonal nature
of the far-field intensity pattern produced by the six-needle mask, the
reconstruction requires a modified algorithm if image resolution is to be
maintained. The results obtained using six-needle mask in the far-field show a
comparable reconstruction resolution between the proposed method and a
traditional raster scanning with a focused spot. This approach was
transitioned to transmission electron microscopy, wherein the modulation
method approach was used to reconstruct the object, with preliminary results
showing promise for image reconstruction with reduced overall illumination
intensity.

Funding We wish to acknowledge the support from the European Union’s Horizon
2020 Research and Innovation Programme (Grant Agreement No. 766970 Project
"Q-Sort", No. 964591 "SMART-electron"), the Royal Society, EPSRC Research
Council funding to QuantIC [EP/M01326X/1] and National Natural Science
Foundation of China (Grant No. 61922011 and U21B2034).

## References

  * [1] P. Sen, B. Chen, G. Garg, S. R. Marschner, M. Horowitz, M. Levoy, and H. P. A. Lensch, “Dual photography,” ACM Trans. Graph. 24, 745–755 (2005).
  * [2] M. F. Duarte, M. A. Davenport, D. Takhar, J. N. Laska, T. Sun, K. F. Kelly, and R. G. Baraniuk, “Single-pixel imaging via compressive sampling,” IEEE Signal Processing Magazine 25, 83–91 (2008).
  * [3] M. Uchida and A. Tonomura, “Generation of electron beams carrying orbital angular momentum,” Nature 464, 737–739 (2010).
  * [4] J. Verbeeck, H. Tian, and P. Schattschneider, “Production and application of electron vortex beams,” Nature 467, 301–304 (2010).
  * [5] B. J. McMorran, A. Agrawal, I. M. Anderson, A. A. Herzing, H. J. Lezec, J. J. McClelland, and J. Unguris, “Electron Vortex Beams with High Quanta of Orbital Angular Momentum,” Science 331, 192–195 (2011).
  * [6] M. Goutsoulas and N. K. Efremidis, “Dynamics of self-accelerating electron beams in a homogeneous magnetic field,” Phys. Rev. A 103, 013519 (2021).
  * [7] N. Voloch-Bloch, Y. Lereah, Y. Lilach, A. Gover, and A. Arie, “Generation of electron Airy beams,” Nature 494, 331–335 (2013).
  * [8] V. Grillo, E. Karimi, G. C. Gazzadi, S. Frabboni, M. R. Dennis, and R. W. Boyd, “Generation of nondiffracting electron bessel beams,” Phys. Rev. X 4, 011013 (2014).
  * [9] V. Grillo, J. Harris, G. C. Gazzadi, R. Balboni, E. Mafakheri, M. R. Dennis, S. Frabboni, R. W. Boyd, and E. Karimi, “Generation and application of bessel beams in electron microscopy,” Ultramicroscopy 166, 48–60 (2016).
  * [10] V. Grillo, A. H. Tavabi, F. Venturi, H. Larocque, R. Balboni, G. C. Gazzadi, S. Frabboni, P.-H. Lu, E. Mafakheri, F. Bouchard, R. E. Dunin-Borkowski, R. W. Boyd, M. P. J. Lavery, M. J. Padgett, and E. Karimi, “Measuring the orbital angular momentum spectrum of an electron beam,” Nature Communications 8, 1–6 (2017).
  * [11] A. H. Tavabi, P. Rosi, E. Rotunno, A. Roncaglia, L. Belsito, S. Frabboni, G. Pozzi, G. C. Gazzadi, P.-H. Lu, R. Nijland, M. Ghosh, P. Tiemeijer, E. Karimi, R. E. Dunin-Borkowski, and V. Grillo, “Experimental Demonstration of an Electrostatic Orbital Angular Momentum Sorter for Electron Beams,” Physical Review Letters 126, 094802 (2021).
  * [12] M. B. Wakin, J. N. Laska, M. F. Duarte, D. Baron, S. Sarvotham, D. Takhar, K. F. Kelly, and R. G. Baraniuk, “An architecture for compressive imaging,” in _2006 International Conference on Image Processing,_ (2006), pp. 1273–1276.
  * [13] E. J. Candes and M. B. Wakin, “An introduction to compressive sampling,” IEEE Signal Processing Magazine 25, 21–30 (2008).
  * [14] J. Romberg, “Imaging via compressive sampling,” IEEE Signal Processing Magazine 25, 14–20 (2008).
  * [15] R. I. Stantchev, B. Sun, S. M. Hornett, P. A. Hobson, G. M. Gibson, M. J. Padgett, and E. Hendry, “Noninvasive, near-field terahertz imaging of hidden objects using a single-pixel detector,” Science Advances 2, 1–6 (2016).
  * [16] M.-J. Sun and J.-M. Zhang, “Single-pixel imaging and its application in three-dimensional reconstruction: A brief review,” Sensors 19 (2019).
  * [17] G. A. Howland, P. B. Dixon, and J. C. Howell, “Photon-counting compressive sensing laser radar for 3D imaging,” Applied Optics 50, 5917–5920 (2011).
  * [18] G. Futia, P. Schlup, D. G. Winters, and R. A. Bartels, “Spatially-chirped modulation imaging of absorbtion and fluorescent objects on single-element optical detector,” Optics Express 19, 1626–1640 (2011).
  * [19] R. Shiloh, Y. Lereah, Y. Lilach, and A. Arie, “Sculpturing the electron wave function using nanoscale phase masks,” Ultramicroscopy 144, 26–31 (2014).
  * [20] S. Shwartz, “Single-pixel imaging with high-energy electromagnetic radiation and particles,” Science Bulletin 66, 857–859 (2021).
  * [21] P. Rosi, F. Venturi, G. Medici, C. Menozzi, G. C. Gazzadi, E. Rotunno, S. Frabboni, R. Balboni, M. Rezaee, A. H. Tavabi, R. E. Dunin-Borkowski, E. Karimi, and V. Grillo, “Theoretical and practical aspects of the design and production of synthetic holograms for transmission electron microscopy,” Journal of Applied Physics 131, 031101 (2022).
  * [22] J. Verbeeck, A. Béché, K. Müller-Caspary, G. Guzzinati, M. A. Luong, and M. Den Hertog, “Demonstration of a 2x2 programmable phase plate for electrons,” Ultramicroscopy 190, 58–65 (2018).
  * [23] G. Ruffato, E. Rotunno, L. Giberti, and V. Grillo, “Arbitrary conformal transformations of wave functions,” Phys. Rev. Applied 15, 054028 (2021).
  * [24] D. J. Griffiths and Y. Li, “Charge density on a conducting needle,” American Journal of Physics 64, 706–714 (1996).
  * [25] G. Matteucci, G. Missiroli, M. Muccini, and G. Pozzi, “Electron holography in the study of the electrostatic fields: the case of charged microtips,” Ultramicroscopy 45, 77–83 (1992).
  * [26] M. Beleggia, T. Kasama, D. J. Larson, T. F. Kelly, R. E. Dunin-Borkowski, and G. Pozzi, “Towards quantitative off-axis electron holographic mapping of the electric field around the tip of a sharp biased metallic needle,” Journal of Applied Physics 116, 024305 (2014).
  * [27] B. J. McMorran, T. R. Harvey, and M. P. J. Lavery, “Efficient sorting of free electron orbital angular momentum,” New Journal of Physics 19, 023053 (2017).
  * [28] A. H. Tavabi, V. Migunov, C. Dwyer, R. E. Dunin-Borkowski, and G. Pozzi, “Tunable caustic phenomena in electron wavefields,” Ultramicroscopy 157, 57–64 (2015).
  * [29] F. Ferri, D. Magatti, L. A. Lugiato, and A. Gatti, “Differential ghost imaging,” Physical Review Letters 104, 1–4 (2010).
  * [30] G. M. Gibson, S. D. Johnson, and M. J. Padgett, “Single-pixel imaging 12 years on: a review,” Optics Express 28, 28190–28208 (2020).
  * [31] J. H. Shapiro, “Computational ghost imaging,” Phys. Rev. A 78, 061802 (2008).
  * [32] A. Kallepalli, J. Innes, and M. J. Padgett, “Compressed sensing in the far-field of the spatial light modulator in high noise conditions,” Scientific Reports 11, 1–8 (2021).
  * [33] R. M. Karp, “On-line algorithms versus off-line algorithms: How much is it worth to know the future?” in _Proceedings of the IFIP 12th World Computer Congress on Algorithms, Software, Architecture - Information Processing ’92, Volume 1 - Volume I,_ (1992), p. 416–429.
  * [34] A. Gatti, E. Brambilla, M. Bache, and L. A. Lugiato, “Ghost imaging with thermal light: Comparing entanglement and classical correlation,” Phys. Rev. Lett. 93, 093602 (2004).
  * [35] B. Sun, S. S. Welsh, M. P. Edgar, J. H. Shapiro, and M. J. Padgett, “Normalized ghost imaging,” Optics Express 20, 16892 (2012).
  * [36] W. Pratt, J. Kane, and H. Andrews, “Hadamard transform image coding,” Proceedings of the IEEE 57, 58–68 (1969).
  * [37] N. J. A. Sloane and M. Harwit, “Masks for hadamard transform optics, and weighing designs,” Applied Optics 15, 107–114 (1976).
  * [38] Z. Zhang, X. Wang, G. Zheng, and J. Zhong, “Hadamard single-pixel imaging versus fourier single-pixel imaging,” Optics Express 25, 19619–19639 (2017).
  * [39] A. Kallepalli, D. Stellinga, M.-J. Sun, R. Bowman, E. Rotunno, V. Grillo, and M. Padgett, “Ghost imaging with electron microscopy inspired, non-orthogonal phase masks,” Research Square Preprint (2021).
  * [40] C. Popa and R. Zdunek, “Kaczmarz extended algorithm for tomographic image reconstruction from limited-data,” Mathematics and Computers in Simulation 65, 579–598 (2004).
  * [41] A. Bautu, E. Bautu, and C. Popa, “A weighted kaczmarz algorithm in image reconstruction,” in _Proceedings of the Fifth Workshop on Mathematical Modelling of Environmental and Life Sciences Problems,_ (2006), pp. 43–50.
  * [42] T. Strohmer and R. Vershynin, “A randomized kaczmarz algorithm with exponential convergence,” Journal of Fourier Analysis and Applications 15, 262–278 (2008).
  * [43] R. G. Pires, D. R. Pereira, L. A. Pereira, A. F. Mansano, and J. P. Papa, “Projections onto convex sets parameter estimation through harmony search and its application for image restoration,” Natural Computing 15, 493–502 (2016).
  * [44] M.-L. Sun, C.-Q. Gu, and P.-F. Tang, “On randomized sampling kaczmarz method with application in compressed sensing,” Mathematical Problems in Engineering 2020, 1–11 (2020).
  * [45] A. Jesacher, A. Schwaighofer, S. Fürhapter, C. Maurer, S. Bernet, and M. Ritsch-Marte, “Wavefront correction of spatial light modulators using an optical vortex image,” Optics Express 15, 5801–5808 (2007).
  * [46] T. Čižmár, M. Mazilu, and K. Dholakia, “In situ wavefront correction and its application to micromanipulation,” Nature Photonics 4, 388–394 (2010).
  * [47] S. Turtaev, I. T. Leite, K. J. Mitchell, M. J. Padgett, D. B. Phillips, and T. Čižmár, “Comparison of nematic liquid-crystal and dmd based spatial light modulation in complex photonics,” Optics Express 25, 29874–29884 (2017).