# Deep Reinforcement Learning for Data-Driven Adaptive Scanning in
Ptychography

M. Schloz<EMAIL_ADDRESS>Humboldt Universität zu Berlin, Institut für
Physik & IRIS Adlershof, Berlin, Germany  J. Müller  Humboldt Universität zu
Berlin, Institut für Physik & IRIS Adlershof, Berlin, Germany  T. C. Pekin
Humboldt Universität zu Berlin, Institut für Physik & IRIS Adlershof, Berlin,
Germany  W. Van den Broek  Humboldt Universität zu Berlin, Institut für Physik
& IRIS Adlershof, Berlin, Germany  C. T. Koch  Humboldt Universität zu Berlin,
Institut für Physik & IRIS Adlershof, Berlin, Germany

###### Abstract

We present a method that lowers the dose required for a ptychographic
reconstruction by adaptively scanning the specimen, thereby providing the
required spatial information redundancy in the regions of highest importance.
The proposed method is built upon a deep learning model that is trained by
reinforcement learning (RL), using prior knowledge of the specimen structure
from training data sets. We show that equivalent low-dose experiments using
adaptive scanning outperform conventional ptychography experiments in terms of
reconstruction resolution.

††preprint: APS/123-QED

##  I Introduction

Ptychography is a coherent diffractive imaging (CDI) method that has found use
in light, x-ray and scanning transmission electron microscopies (STEM). The
method combines whole diffraction patterns from spatially overlapping regions
to reconstruct the structure of a specimen for arbitrarily large fields of
view [1], with many advantages over other imaging methods [2, 3, 4, 5]. The
development of new hardware [6, 7] and reconstruction algorithms [8, 9] has
led to ptychography becoming a mature electron microscopy technique [4].
Current research to further improve this technique is driven by the desire to
investigate thick samples [10, 11, 12, 13, 14] as well as to lower the
required electron dose [15, 16, 17, 18].

In order to lower the electron dose used, researchers have tried to vary
various experimental parameters while preserving information redundancy
through overlapping probes. One approach involves a defocused probe rastered
across the specimen, with a less dense scan pattern. This uses therefore a
lower dose than focused probe ptychography, but introduces additional
complications for the reconstruction algorithm due to an increased need to
account for partial spatial coherence in the illuminating probe [18]. Another
approach is simply to scan faster - by lowering the probe dwell time per probe
position, an overall decrease in dose can be realized. However, this comes
with its own limitations, as the physical limits of the electron source,
microscope, and camera all must be considered. Finally, a third approach is
the optimization of the scan pattern, deviating from a raster grid in favour
of a generally more efficient pattern [19]. This approach can, however, only
yield a limited improvement in reconstruction quality as it is not capable of
taking into account the structure of the specimen in the scan pattern.

In this paper we present an approach particularly tailored for electron
ptychography that enables reduction of the electron dose through adaptive
scanning. It is based upon the idea that, at atomic resolution, ptychography
requires an increased information redundancy through overlapping illuminating
beams only at regions that contain atomic structure of the scanned specimen.
We present here an algorithm that scans only the regions with the highest
information content in order to strongly improve the ptychographic
reconstruction quality while keeping the total number of scan positions, and
therefore the total dose, low. The scan positions are predicted sequentially
during the experiment and the only information required for the prediction
process is the diffraction data acquired at previous scan positions.

The scan position prediction model of the algorithm is a mixture of deep
learning models, and the model training is performed with both supervised and
reinforcement learning. The synergy of deep learning and reinforcement
learning has already shown strong performance in various dynamic decision
making problems, such as playing Atari games [20], or Go [21], as well as
tasks in robotics [22, 23] and visual recognition [24, 25, 26]. The success of
this approach, despite the complexity of the problems that they had to
overcome, can be attributed to their algorithms’ ability of learning
independently from data.

Similarly, the proposed algorithm here solves a sequential decision making
problem by learning from a large amount of simulated or, if available,
experimental ptychographic data consisting of hundreds to thousands of
diffraction patterns. Here, the focus of the learning is specifically designed
to maximize the dynamic range in the reconstruction for each individual scan
position. The algorithm then transfers the learned behaviour it developed
offline to a realistic experimental environment.

Our approach is conceptually related to the subfield of computer vision that
focuses on identifying relevant regions of images or video sequences for the
purpose of classification or recognition. However, there are fundamental
differences not only in the purpose, but also in the solution strategy for our
application in contrast to computer vision tasks. Differences include a lack
of direct access to images (updated real space information is only accessible
through a highly optimized reconstruction algorithm), non-optimal parameter
settings of the reconstruction algorithm and experimental uncertainties such
as imprecise scan positioning of the microscope or contamination of the
specimen requiring pre-processing of the reconstructed image, and the
necessity of a much larger number of measurements requiring methods that
improve the performance of the sequential decision making process.

Work in adaptive scanning for x-ray fluorescence imaging [27] and for scanning
probe microscopy [28] has recently been reported. The work in [27] uses RL to
determine the exposure time on a per pixel basis sequentially for multiple
apertures that vary in their respective resolution. It is therefore more
closely related to previous work in scanning electron microscopy that divides
the measurement into a low-dose raster scan and a subsequent high-dose
adaptive scan [29]. The latter work in [28] uses Gaussian processes based
Bayesian optimization to sequentially explore the image space with the
scanning probe. However, it has been reported that this model suffers in
performance as it lacks prior knowledge of the domain structure, which can be
compensated by including a deep learning model with domain specific knowledge.
Our proposed algorithm is the first application of adaptive scanning to
ptychography, and is further unique in that the scan pattern is predicted
using prior knowledge about the sample in the form of a pre-trained deep RL
network, thereby improving performance. Our research forms a basis for a new
avenue of automated and autonomous microscopy [30].

With the ever-increasing data storage capacities, implementations of data
infrastructures and data sharing platforms [31, 32, 33, 34], access to
ptychographic data will be further facilitated and data-driven adaptive
scanning schemes can be applied to a vast number of ptychographic experiments.

We demonstrate the performance of our algorithm using experimentally acquired
data. Our analysis shows that the algorithm can sufficiently learn information
about the structure of a material from data in order to optimize the scan
behaviour of the microscope in a real experiment. For low dose experiments, we
show that adaptive scanning can improve the ptychographic reconstruction
quality by up to $25.75\%$ and the resolution by up to $31.59\%$ compared to a
non-adaptive (random) scan method. Adaptive scanning allows for the retrieval
of the material’s structure in this low dose regime and even improves the
resolution of the reconstruction when compared to the reconstruction obtained
using the conventional high dose raster grid scan approach.

##  II Methods

###  II.1 Image formation in ptychography

Single slice ptychography can be expressed by a multiplicative approximation
that describes the interaction of a wavefunction $\psi^{in}_{p}(\textbf{r})$
of an incoming beam with the transmission function $t(\textbf{r})$ of a
specimen. For each measurement $p$, the beam is shifted by $\textbf{R}_{p}$
and a diffraction pattern is acquired with the intensity $I_{p}$ that is
expressed by:

$\begin{split}I_{p}=|\Psi^{ex}_{p}(\textbf{k})|^{2}=|\mathcal{F}\left[\psi^{in}_{p}(\textbf{r}-\textbf{R}_{p})t(\textbf{r})\right]|^{2},\end{split}$
(1)

where $\mathcal{F}$ is the Fourier propagator, r the real space coordinate, k
the reciprocal space coordinate and $\Psi^{ex}_{p}(\textbf{k})$ the exit
wavefunction at the detector. According to the strong phase object
approximation, the transmission function can be defined as
$t(\textbf{r})=e^{i\sigma V(\textbf{r})}$, with the interaction constant
$\sigma$ and the complex quantity $V(\textbf{r})$, where the real part
$V_{re}(\textbf{r})$ corresponds to the local projected electrostatic
potential and the imaginary part $V_{im}(\textbf{r})$ accounts for absorption
or scattering outside the range of scattering angles and energy losses
recorded by the detector. Throughout the remainder of this paper, the variable
$\sigma$ is absorbed into $V(\textbf{r})$. X-ray and optical ptychography is
mathematically described similarly with the only difference that the
transmission function $t(\textbf{r})$ is related to the complex refractive
index of the specimen. Figure 1 illustrates the experimental configuration of
conventional ptychography.

Figure 1: Experimental setup in ptychography. At the scan position
$\textbf{R}_{p}$ of the scan sequence, the beam illuminates a sample, where
the incident electron wave $\psi^{in}_{p}(\textbf{r}-\textbf{R}_{p})$
interacts with the transmission function $t(\textbf{r})$. The wave exiting the
sample is propagated by a Fourier transform to the detector located in the far
field and the intensity $I_{p}=|\Psi^{ex}_{p}(\textbf{k})|^{2}$ is recorded.

The potential of the specimen is recovered from data of experimentally
acquired diffraction patterns $J_{p}$ using a reconstruction algorithm. Here,
we apply a gradient based algorithm [17] with a gradient decent optimization
and the potential is retrieved by iteratively minimizing the loss function:

$\mathcal{L}(V)=\frac{1}{P}\sum^{P}_{p=1}\|I_{p}(V)-J_{p}\|^{2}_{2}.$ (2)

Although the approach described in this paper is compatible with multisclice
ptychography, in light of the application to a 2D material we constrain
ourselves to single-slice ptychography.

###  II.2 Generation of scan sequences

We consider a recurrent neural network (RNN) [35, 36, 37] for the generation
of scan sequences. Its network architecture is designed to model temporal
sequences with recurring input information. Memory cells combine the current
input information $X_{t}$ with the hidden state $H_{t}$ and map it to the next
hidden state $H_{t+1}$. These hidden states represent the memory gathered from
all the previous time steps. At every time step $t$, an output is generated on
the basis of the current hidden state. In the implementation shown here, the
output corresponds to a sub-sequence of scan positions, given by a vector of
2D coordinates $\vec{\textbf{R}}_{P_{t}}$. In principle, the output can be
reduced to a single scan position $\textbf{R}_{p_{t}}$, but we do not for to
practical reasons that will be discussed later. The sub-sequence is predicted
via a fully connected layer (FC) that is parameterized by the layer weights
$\theta_{H}$:

$\vec{\textbf{R}}_{P_{t}}=\text{FC}_{\theta_{H}}(H_{t}).$ (3)

At the predicted scan positions $\vec{\textbf{R}}_{P_{t}}$, diffraction
patterns $\bm{J}_{P_{t}}$ are acquired by the microscope and from these
diffraction patterns a potential $V_{t}(\textbf{r})$ is reconstructed
minimizing Eq. (2). The intermediate reconstruction $V_{t}(\textbf{r})$
combined with its corresponding sub-sequence of scan positions
$\vec{\textbf{R}}_{P_{t}}$ can then be used for the input information $X_{t}$
of the RNN. However, the bandwidth of the information given in
$V_{t}(\textbf{r})$ and $\vec{\textbf{R}}_{P_{t}}$ differs strongly and thus
pre-processing is required before the two components can be concatenated and
mapped to $X_{t}$. For the processed location information $\textbf{L}_{t}$
based on the sub-sequence $\vec{\textbf{R}}_{P_{t}}$, a FC that is
parameterized by the weights $\theta_{R}$ is used:

$\textbf{L}_{t}=\text{FC}_{\theta_{R}}(\vec{\textbf{R}}_{P_{t}}).$ (4)

For the processed structure information $\textbf{C}_{t}$ based on the
reconstructed potential $V_{t}(\textbf{r})$, a compressed representation
$z_{t}$ is generated by using the encoder part of a convolutional autoencoder
[38]. This processing step is described in more detail in appendix A. The
compressed representation $z_{t}$ is then fed into a FC that is parameterized
by the weights $\theta_{z}$:

$\textbf{C}_{t}=\text{FC}_{\theta_{z}}(z_{t}).$ (5)

The processed location information $\textbf{L}_{t}$ is subsequently
concatenated with the processed structure information $\textbf{C}_{t}$ and
mapped to the input information $X_{t}$ with a FC that is parameterized by the
weights $\theta_{LC}$. The whole process of predicting sub-sequences of scan
positions and acquiring the corresponding diffraction patterns is repeated
until a ptychographic dataset of desired size is reached.

In practice and even after a reduction through adaptive scanning, several
hundreds to thousands of diffraction patterns are required for effective
ptychography. Covering this range of scan positions with a strong prediction
performance requires efficient training of a large RNN. Backpropagation
through time (BPTT) is typically used to generate the required gradients to
update the network weights
$\theta=\\{\theta_{H},\theta_{\text{GRU}},\theta_{LC},\theta_{R},\theta_{z}\\}$
of the RNN. Its foundation on the chain-rule, with terms being multiplied by
themselves as many times as the length of the network, can result in problems
with training efficiency. For even the most basic RNN architectures, BPTT
fails for relatively short sequences due to the so-called vanishing or
exploding gradient problem [39]. To circumvent this issue a more complex RNN
architecture was proposed by Hochreiter et al. [40]. The Long-Short-Term-
Memory (LSTM) network uses a more complex mapping between the input
information and hidden state to the output, which allows a more efficient
training using the BPTT method for larger networks. The gated recurrent units
(GRU) network, which is a computationally faster, simplified version of the
LSTM network, is used in this paper [41]. A very large network would,
nevertheless, be difficult to train using BPTT and also greatly increase
acquisition time in adaptive scanning due to, e.g., a more frequent data
transfer and generation of intermediate reconstructions $V_{t}(\textbf{r})$.
Therefore, a sub-sequence of scan positions $\vec{\textbf{R}}_{P_{t}}$ is
preferred over a single scan position $\textbf{R}_{p_{t}}$ as the RNN output.
Figure 2 shows the prediction process modeled by the RNN in full detail.

Figure 2: Schematic of the forward propagation process of the RNN model. The
RNN consists of GRU units that use the hidden state $H_{t}$ from the previous
time step and the hybrid input information $X_{t}$ to create a new hidden
state $H_{t+1}$. The hybrid input is the concatenation of the pre-processed
information from the sub-sequence of scan positions $\vec{\textbf{R}}_{P_{t}}$
and the corresponding compressed representation of the partial reconstruction
$z_{t}$. The output of the GRU cell is used to predict the positions of the
next sub-sequence $\vec{\textbf{R}}_{P_{t+1}}$ and is also used as the input
for the next GRU cell. The process is repeated until the full length of the
scan sequence, consisting of $T$ sub-sequences, is reached.

###  II.3 Training through reinforcement learning

A RNN, such as the one described in the previous section, can be combined with
RL to provide a formalism for modelling behaviour to solve decision making
problems. In RL, a learning agent interacts with an environment, while trying
to maximize a reward signal. This is generally formalized as a Markov decision
process (MDP) described by a 5-tuple:
$\langle\mathcal{S},\mathcal{A},\rho,r,\gamma\rangle$. At each time-step $t$
the agent has complete knowledge of the environment by observing the state
$s_{t}\in\mathcal{S}$ and makes an optimal decision by selecting an action
$a_{t}\in\mathcal{A}$. Based on $s_{t}$ and $a_{t}$, the next state $s_{t+1}$
is generated according to a transition function
$\rho:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow[0,1]$. The
agent additionally receives a feedback through a scalar reward function
$r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$. This reward $r$
contributes to the total reward computed at the end of the sequence,
$G=\sum^{T}_{t=0}\gamma^{t}r(a_{t},s_{t})$, also known as the return. The
discount factor $\gamma\in[0,1]$ controls the emphasis of long-term rewards
versus short-term rewards.

In the case of adaptive scanning in ptychography, complete knowledge of the
specimen structure is not known and the previous described formalism, where
observations are equivalent to states, is not quite applicable. A partially
observable Markov decision process (POMDP) generalizes the MDP to a 7-tuple:
$\langle\mathcal{S},\mathcal{A},\rho,r,\mathcal{O},\omega,\gamma\rangle$ by
considering the observation $o_{t}\in\mathcal{O}$ to contain only partial or
incomplete information about the state $s_{t}$, and which is generated
according to a observation function
$\omega:\mathcal{A}\times\mathcal{S}\times\mathcal{O}\rightarrow[0,1]$.
Therefore, $o_{t}$ can not sufficiently represent the state $s_{t}$ and
instead the entire history of observations and actions up to the current time
$h_{t}=\\{o_{1},a_{1},...,o_{t-1},a_{t-1},o_{t}\\}$ is used as basis for
optimal or near-optimal decision making. A stochastic policy
$\pi_{\theta}(a_{t}|h_{t})$ maps the history of past interactions $h_{t}$ to
action probabilities. Given a continuous action space, the policy can be
represented by a two-dimensional Gaussian probability distribution:

$\pi_{\theta}(a_{t}|h_{t})=\mathcal{N}(\bm{\mu}_{\theta}(h_{t}),\Sigma),$ (6)

with its mean vector $\bm{\mu}_{\theta}(h_{t})$ corresponding to
$\textbf{R}_{p_{t}}$, where the history $h_{t}$ is summarized in the hidden
state $H_{t}$ of the RNN and the covariance matrix $\Sigma$ with fixed
variances $\sigma_{x}^{2}\in[0,1]$ and $\sigma_{y}^{2}\in[0,1]$.

In this POMDP formalism, however, a single action $a_{t}$ is drawn from the
probability distribution $\pi_{\theta}(a_{t}|h_{t})$, which corresponds to a
single agent interacting with the environment. This is incompatible with scan
control in ptychography where we seek to predict multiple scan positions at
each time step. A partially observable stochastic game (POSG) extends the
POMDP formalism to a 8-tuple, $\langle
M,\mathcal{S},\\{\mathcal{A}^{m}\\}_{m\in M},\rho,\\{r^{m}\\}_{m\in
M},\\{\mathcal{O}^{m}\\}_{m\in M},\omega,\gamma\rangle$, with multiple agents
$M$, each selecting an action $a^{m}_{t}$ and making an observation
$o^{m}_{t}$ given the state $s_{t}$. Thus, joint actions $\bm{a}_{t}=\langle
a^{1}_{t},...,a^{m}_{t}\rangle$ from the joint action space
$\bm{\mathcal{A}}=\mathcal{A}_{1}\times...\times\mathcal{A}_{M}$ are executed
and joint observations $\bm{o}_{t}=\langle o^{1}_{t},...,o^{m}_{t}\rangle$
from the joint observation space
$\bm{\mathcal{O}}=\mathcal{O}_{1}\times...\times\mathcal{O}_{M}$ are received
from the environment at each time step. In this case, the transition function
is given by
$\rho:\mathcal{S}\times\bm{\mathcal{A}}\times\mathcal{S}\rightarrow[0,1]$, the
observation function is given by
$\omega:\bm{\mathcal{A}}\times\mathcal{S}\times\bm{\mathcal{O}}\rightarrow[0,1]$
and each agent receives its immediate reward defined by the reward function
$r^{m}:\mathcal{S}\times\bm{\mathcal{A}}\rightarrow\mathbb{R}$. Here, we
consider the individual agent to have access to the actions and observations
of all other agents, which allows the optimization of its individual policy
$\pi_{\theta^{m}}(a^{m}_{t}|\bm{h}_{t})$ using the joint history of
observations and actions
$\bm{h}_{t}=\\{\bm{o}_{1},\bm{a}_{1},...,\bm{o}_{t-1},\bm{a}_{t-1},\bm{o}_{t}\\}$.
The joint policy of all agents is then defined as
$\bm{\pi}_{\bm{\theta}}(\bm{a}_{t}|\bm{h}_{t})=\prod_{m=1}^{M}\pi_{\theta^{m}}(a^{m}_{t}|\bm{h}_{t})$,
with $\bm{\theta}=\\{\theta^{m}\\}_{m\in M}$. The goal of RL is now to learn a
joint policy that maximizes the expected total reward for each agent $m$ with
respect to its parameters $\theta^{m}$:

$\begin{split}\mathcal{J}^{m}(\bm{\theta})=\mathbb{E}_{\bm{\pi}_{\bm{\theta}}(\bm{\tau})}\left[G^{m}\right]\approx\frac{1}{N}\sum^{N}_{n=1}G^{m}_{n},\end{split}$
(7)

where the expected total reward can be approximated by Monte Carlo sampling
with $N$ samples. In this paper, improvement of the policy is achieved by
updating the policy parameters
$\theta^{m}=\\{\theta^{m}_{H},\theta_{\text{GRU}},\theta_{LC},\theta_{R},\theta_{z}\\}$
with ’REINFORCE’ [42], a policy gradient method:

$\begin{split}\nabla_{\theta^{m}}\mathcal{J}^{m}(\bm{\theta})=\mathbb{E}_{\bm{\pi}_{\bm{\theta}}(\bm{\tau})}\left[\nabla_{\theta^{m}}\text{log}\bm{\pi}_{\bm{\theta}}(\bm{\tau})G^{m}\right]\\\
\approx\frac{1}{N}\sum^{N}_{n=1}\sum^{T}_{t=0}\nabla_{\theta^{m}}\text{log}\pi_{\theta^{m}}(a^{m}_{n,t}|\bm{h}_{n,t})&G^{m}_{n}.\end{split}$
(8)

The derivation of $\nabla_{\theta^{m}}\mathcal{J}^{m}(\bm{\theta})$ is given
in the appendix B.

###  II.4 Learning to adaptively scan in ptychography

While policy gradient methods are the preferred choice to solve reinforcement
learning problems in which the action spaces are continuous [43], they come
with significant problems. Like any gradient based method, policy gradient
solutions mainly converge to local, not global, optima [44]. In this paper, we
reduce the effect of this problem during training by splitting the training of
the RNN into supervised learning and RL. The first training step initializes
the policy parameters such that the scan pattern follows a conventional grid
pattern, thereby avoiding relatively poor local optima during subsequent
policy gradient steps. This training step is explained in more detail in the
appendix C. In the second training step, the pre-trained policy is fine tuned
through RL, resulting in a scan pattern that has been adapted to the structure
of the material.

A high variance of gradient estimates is another problem that particularly
strongly affects the Monte Carlo policy gradient method [45, 43, 46]. Due to
this, the sampling efficiency is relatively low, which causes a slow
convergence to a solution. This makes deep RL applied to ptychography
challenging as the image reconstruction itself requires iterative processing
(see section II.1).

The high variance can be in part attributed to the difficulty of assigning
credit from the overall performance to an individual agent’s action. This
credit assignment problem is limited to a temporal problem in the single agent
RL case [47]. In this case, methods for variance reduction to better assign
credit to individual actions are for instance reward-to-go [48] or using a
baseline [44, 45].

In RL involving multiple cooperating agents with a shared reward function
$r^{1}(\bm{a}_{t}|s_{t})=r^{2}(\bm{a}_{t}|s_{t})=...=r^{m}(\bm{a}_{t}|s_{t})$,
the challenge to overcome the credit assignment problem further increases due
to the necessity of now identifying the contribution of each agent’s action to
the total reward. This challenge can be tackled with difference reward [49,
50, 51, 52], which replaces the shared reward with a shaped reward that is
formed by comparing the global reward with a reward that an agent would
receive when performing a default action.

Following the idea of a difference reward in spirit, we introduce a way to
estimate the reward function in order to tackle the credit assignment problem
for adaptive scanning in ptychography. The reward function should naturally
correspond to the quality of the ptychographic reconstruction. We have found
empirically that a high reconstruction quality correlates positively with a
high dynamic range in the phase. Therefore, the reward function could
intuitively be formalized by
$r^{m}(\bm{a}_{t}|s_{t})=P^{-1}\sum_{\textbf{r}\in\text{FOV}}V(\textbf{r})$,
where $P$ is the total number of scan positions. This formulation, however,
does not solve the credit assignment problem and results in an insufficient
training performance, as shown in Figure 3a). To estimate the reward for the
actions of each individual agent, we use a tessellation method that partitions
the atomic potential into small segments. A Voronoi diagram [53], where each
position corresponds to a seed for one Voronoi cell, enables assignment of
only a part of the total phase to each position. More precisely, the Voronoi
diagram formed by the predicted scan positions is overlaid with the
corresponding ptychographic reconstruction at the end of the prediction
process and the summed phase within each Voronoi cell is the reward for that
cell’s seed position. The reward function can be expressed by
$r^{m}(\bm{a}_{t}|s_{t})=P^{-1}\sum_{\textbf{r}\in\text{Cell}^{m}}V(\textbf{r})$.
Figure 3b) shows a Voronoi diagram generated by predicted scan positions.

Figure 3: a) Learning curves of RL with multiple agents that use a shared
reward or shaped reward, illustrated in orange and blue, respectively. b) A
Voronoi diagram is used to assign a unique reward to each scan position of the
predicted sequence. The scan positions are shown as red dots, where the first
50 positions are distributed on the right side within the dark blue area. For
visualization purpose, the ground truth reconstruction is included in the
diagram.

###  II.5 Experiment and model design

For the experimental investigation, we acquired multiple ptychographic
datasets from a monolayer molybdenum disulfide (MoS2) specimen with a NION
HERMES microscope. The microscope was operated with a $60$ kV acceleration
voltage, a convergence angle of $33$ mrad and diffraction patterns with a
pixel size of $0.84$ mrad were acquired using a Dectris ELA direct electron
detector mounted at the electron energy loss spectroscopy (EELS) camera port.
Distortions induced by the EEL spectrometer were corrected with in-house
developed software. For the ptychographic dataset acquisition, a conventional
grid scan with a scanning step size of $0.02$ nm was used.

From the experimentally acquired datasets we created $175$ smaller datasets,
each with 10,000 diffraction patterns. The diffraction patterns were binned by
a factor of $2$ to $64\times 64$ pixels. The adaptive scanning algorithm was
then trained on the smaller datasets with the goal of predicting optimal scan
sequences of $250$ to $500$ probe positions, out of the possible 10,000, which
corresponds to a dose reduction by a factor of $40$ to $20$. Each sub-sequence
contains 50 to 100 positions, where the first sub-sequence follows a quasi-
random Halton sequence.

The ptychographic reconstructions were performed with an optimized version of
ROP [17] that allows simultaneous reconstruction from a batch of different
datasets, which was required for efficient model training. A gradient descent
step size $\alpha_{\text{ROP}}$ of $5.25\text{{e}}2$ was chosen and the
potential was retrieved at iteration 5. The reconstructed potential was
$200\times 200$ pixels with a pixel size of $0.0154$ nm, for a field of view
of $2\times 2$ nm. For the generation of the reward function, Voronoi diagrams
were generated with the Jump Flooding Algorithm [54] and for the
implementation of the network models, PyTorch [55] was used. For the
compression of structure information, we used a convolutional autoencoder
consisting of $6$ convolutional layers with kernels of dimension $3$, a stride
of $1$ and channels that ranged from $16$ to $512$ for the encoder and decoder
part, respectively. The input of the autoencoder had a dimension of $512$ with
a pixel size of $0.0064$ nm and thus a scaling and an interpolation was
required before the potential generated by ROP could be compressed. In
addition, the value of the potential $V_{i}$ at each pixel $i$ was transformed
to zero mean and unit variance. For the prediction of the scan sequences, pre-
training and fine-tuning was performed with a RNN model composed of $2$
stacked GRU layers with hidden states $H_{t}$ of size $2048$, the Adam
optimizer [56] with a learning rate $\alpha_{\text{RNN}}$ of $1\text{{e}-}5$
and a batch size of $24$. For the fine-tuning, a policy with variances of
$\sigma_{x}^{2}=\sigma_{y}^{2}=0.0125^{2}$ was chosen and a myopic behavior
was enforced by setting the discount factor for the return, $G$, to
$\gamma=0$. All settings used for training the adaptive scanning algorithm are
summarized in Table 2.

##  III Results

###  III.1 Adaptive scanning on experimental MoS2 data

Figure 4 shows the result of adaptive scanning on experimentally acquired MoS2
data and compares it to the result of a random scanning and the conventional
grid scanning procedure. The data used for the comparison was not part of the
training data for the adaptive scanning model. While the full data set
consisting of 10,000 diffraction patterns has been used to obtain a ground
truth reconstruction, only $250$ diffraction patterns have been used for the
adaptive scanning as well as the random scanning reconstruction. Figure 4a)
shows the ptychographic reconstruction when using a random scanning procedure.
The structure of the material is not clearly resolved and large parts of the
field of view are not covered by the scanning procedure. Figure 4c) shows the
reconstruction when the scan positions are predicted by the adaptive scanning
algorithm. The structure of the MoS2 material is now much better resolved and
is closer to the ground truth reconstruction of the full data grid scanning
procedure, shown in Figure 4e). Figure 4b), d) and f) show the diffractograms
of the corresponding reconstructions and the circled diffraction spots show
that the highest resolution of $1.08$ Å is achieved by the adaptive scanning
procedure, while the lowest resolution of $3.25$ Å is obtained by the random
scanning procedure. Further examples of reconstructions and their
corresponding scan sequences are shown in Figure 8.

Figure 4: Ptychographic reconstructions of a MoS2 data set with different
scanning procedures. a) Reconstruction from $250$ diffraction patterns of the
data set that correspond to scan positions which follow a random sequence and
c) an adaptively predicted sequence. e) Reconstruction of the full data set
with 10,000 diffraction patterns acquired with the conventional grid scan. b),
d) and f) Corresponding diffractograms of the reconstructions. The real-space
distance of the circled diffraction spots is labeled.

The results suggest that probe delocalization due to scattering plays an
important role as to why an improved ptychographic reconstruction can be
achieved by distributing the scan positions predominantly on the atoms of the
specimen. When the beam is positioned on an atom, it scatters to higher angles
and thus experiences spatial delocalization. It therefore also probes a larger
range in real space, i.e. the scattering includes the local environment of the
atom the beam hits. This implies that similar results could be achieved by
using RL with a reward function that specifically emphasizes the scattered
electrons in the recorded diffraction patterns, which is an interesting area
for future research.

The final point of our investigation into adaptive scanning in ptychography
evaluates the performance of the method for various prediction settings. We
compare the structural similarity index measure (SSIM) [57] between the
reconstruction obtained from the reduced data and the ground truth
reconstruction obtained from the full data to quantify the improvement when
using adaptive scanning. Here, SSIMa and SSIMr are the SSIM using a
reconstruction of reduced data that is obtained with the adaptive scanning and
the random scanning procedure, respectively. Table 1 shows the relative
reconstruction quality improvement
$Q_{\text{SSIM}}=(\text{SSIM}_{a}-\text{SSIM}_{r})/(\text{SSIM}_{r})$ for
different experimental settings averaged over $25$ data sets. Additionally,
the relative resolution improvement $Q_{\text{res}}$ averaged over the same
datasets is given. In the case of $250$ scan positions, which corresponds to a
dose reduction by a factor of $40$ with respect to the original data, tests
were performed for multiple sub-sequences, i.e. predictions. The quality
improvement ranges from $16.71\%$ to $25.75\%$ and the resolution improvement
ranges from $9.74\%$ to $27.57\%$ for a number of $2$ to $5$ sub-sequences
(corresponding to 1 to 4 predictions), respectively. Further tests were
performed using a larger number of total scan positions and 5 sub-sequences.
However, while the relative resolution slightly improves for an increasing
amount of scan positions, the difference in quality between the reconstruction
generated with the positions of the adaptive scan and the random scan
decreases with the total number of positions used, as can be expected, since
the random sampling covers the sampled area in an increasingly complete
manner. It should be noted that for all tests which used adaptive scanning
with $5$ sub-sequences, a higher resolution of the reconstruction compared to
that of the reconstruction of the full data set was achieved. These results
indicate that the reconstruction quality and resolution improves with the
frequency by which the positions are predicted, and that low dose experiments
benefit the most from the adaptive scanning scheme.

Table 1: Performance of adaptive scanning for various experimental settings that differ in the number of scan positions and the total number of sub-sequences. For each setting, the oversampling ratio $N_{\text{k}}/N_{\text{u}}$, which is calculated following [17], and the electron dose is given. # Pos. | $N_{\text{k}}/N_{\text{u}}$ | Dose ($\mathrm{e^{-}}$/Å-2) | # Sub-seq. | Q${}_{\text{SSIM}}$ | Q${}_{\text{res}}$
---|---|---|---|---|---
$250$ | $8.21$ | $1.34\text{$\text{{e}}$5 }$ | $2$ | 16.71% | $9.74\%$
|  |  | $3$ | 22.89% | $15.06\%$
|  |  | $4$ | 24.96% | $19.79\%$
|  |  | $5$ | $25.75\%$ | $27.57\%$
$335$ | $10.83$ | $1.79\text{$\text{{e}}$5 }$ | $5$ | $12.53\%$ | $28.71\%$
$420$ | $13.43$ | $2.25\text{$\text{{e}}$5 }$ | $5$ | $12.29\%$ | $29.23\%$
$500$ | $15.86$ | $2.68\text{$\text{{e}}$5 }$ | $5$ | $8.16\%$ | $31.59\%$

##  IV Conclusion

In this paper we present a method for electron ptychography that reduces the
electron dose through adaptive scanning. It is based on the idea that
ptychography requires only an increased information redundancy through
overlapping illuminating beams at regions of the sample that contain atomic
structure. The prediction algorithm is a mixture of deep learning models being
trained using supervised and reinforcement learning.

We show an improved reconstruction quality and resolution when using an
adaptive scanning approach on experimentally acquired monolayer MoS2 datasets
in comparison to another dose reduction scanning approach. In a low dose
experiment the adaptive scanning procedure could improve on average the
reconstruction quality by up to $27.75\%$ and the resolution by up to
$31.59\%$. The resolution achieved by adaptive scanning was also higher than
that of the reconstruction from the full data set, but without the same
homogeneous reconstruction quality throughout the entire field of view.

These improvements show that adaptive scanning for ptychography is a useful
technique to lower the dose needed for the analysis of sensitive samples. In
addition to that, the proposed algorithm can be taken as a blueprint for a
broad range of scanning based microscopy methods and thus paves the way for
future research in machine learning supported, automated and autonomous
microscopy.

###### Acknowledgements.

M.S., T.C.P. and W.V.d.B. acknowledge financial support from the Deutsche
Forschungsgemeinschaft (DFG, German Research Foundation) Grant No. BR
5095/2-1. M.S., J.M. and C.T.K acknowledge support from the DFG - Project-ID
182087777 - SFB 951. M.S. and C.T.K. acknowledge financial support from the
DFG - Project-ID 414984028 - SFB 1404. We thank Prof. Sang Ho Oh and Dr.
Jinsol Seo (Korea Institute of Energy Technology, Naju, Korea), as well as Dr.
Bumsu Park (CEMES, Toulouse, France) for providing the MoS2 sample. M.S. would
like to thank Jayesh K. Gupta (Microsoft, Stanford, USA) for his helpful
suggestion about the description of the POSG formalism.

##  Appendix A Structure content compression through convolutional
autoencoders

For the processing of the intermediate reconstruction $V_{t}(\textbf{r})$ that
is the basis for the compressed representation $z_{t}$ of the RNN, we make use
of a convolutional autoencoder [38], which is a deep learning model based on
convolutional neural networks [58]. This model is composed of two parts,
namely an encoder and a decoder. The encoder transforms potentially degraded
input data into a compressed representation and the decoder recovers the
original input from this representation [59]. By reducing the dimensionality,
the model ensures that only the most important information for the
reconstruction is extracted. In the convolutional autoencoder applied in this
paper, an image $V\in\mathbb{R}^{H\times W\times 2}$ reconstructed from
diffraction patterns is mapped to the latent representation
$z\in\mathbb{R}^{\frac{H}{2^{b}}\times\frac{W}{2^{b}}\times f}$ by using the
encoder network $E_{\phi_{e}}$:

$z=E_{\phi_{e}}(V),$ (9)

where $\phi_{e}$ corresponds to the encoder network weights. The encoder
network consists of a concatenation of $b$ convolution layers that increase
the feature space $f$ and halve the dimension size $H\times W$. The latent
representation $z$ is then mapped back to the original dimensions by the
decoder network:

$\hat{V}=D_{\phi_{d}}(z),$ (10)

where $\phi_{d}$ are the decoder network weights and $\hat{V}$ is the
predicted image. Transposed convolutional layers form the decoder network,
which has the inverse effect to the layers used in the encoder network. Figure
5 illustrates the encoder-decoder architecture in full detail.

Figure 5: Schematic of the convolutional autoencoder model. A reconstruction
$V$ generated from diffraction patterns is mapped to the compressed
representation $z$ by using the encoder network $E_{\phi_{e}}(V)$. The
compressed representation $z$ is then basis for a reverse mapping by the
decoder network $D_{\phi_{d}}(z)$ to generate a prediction of the potential
$\hat{V}$.

After estimating the network weights $\phi_{e}$ and $\phi_{d}$ by minimizing
the loss function:

$\mathcal{M}(\phi_{e},\phi_{d})=\|D_{\phi_{d}}(E_{\phi_{e}}(V))-V\|^{2}_{2},$
(11)

we can utilize the encoder network $E_{\phi_{e}}$ for the compression of
$V_{t}(\textbf{r})$. Figure 6 shows a compression of a partial reconstruction
$V_{t}(\textbf{r})$ and the decompression of its corresponding compressed
representation $z_{t}$. This pre-processing helps the algorithm form the
hybrid input information $X_{t}$ by reducing the structure input information
size, but also handle reconstructions from experimentally acquired data that
may suffer from noise, contamination and/or incorrect scan positions.

Figure 6: Convolutional autoencoder applied on partial structure information,
given by the reconstruction of data from a sub-sequence of scan positions. a)
The reconstruction $V_{t}(\textbf{r})$ from a sub-sequence of scan positions
that is used as input for the convolutional autoencoder. b)-e) $4$ channels of
the compressed representation $z_{t}$ of the structure information. f) Decoded
structure information $\hat{V}_{t}$ from $z_{t}$.

##  Appendix B The ”REINFORCE” algorithm

In the case of the multi agent RL problem, where we use the POSG formalism,
the objective of an agent $m$ given by Eq. (7) can be expressed by:

$\mathcal{J}^{m}(\bm{\theta})=\mathbb{E}_{\bm{\pi}_{\bm{\theta}}(\bm{\tau})}\left[G^{m}\right]=\int\bm{\pi}_{\bm{\theta}}(\bm{\tau})G^{m}d\bm{\tau},$
(12)

with the trajectory
$\bm{\tau}=\\{s_{0},\bm{o}_{0},\bm{a}_{0},s_{1},...,s_{T},\bm{o}_{T},\bm{a}_{T}\\}$
and the policy induced trajectory distribution
$\bm{\pi}_{\bm{\theta}}(\bm{\tau})=q(s_{0})\prod^{T}_{t=0}\rho(s_{t+1}|s_{t},\bm{a}_{t})\bm{\pi}_{\bm{\theta}}(\bm{a}_{t}|\bm{h}_{t})\omega(\bm{o}_{t}|s_{t})$
and where $q(s_{0})$ is the distribution of initial states. Applying the
gradient $\nabla_{\theta^{m}}$ to the objective and using the identity
$\nabla_{\bm{\theta}}\bm{\pi}_{\bm{\theta}}(\bm{\tau})=\bm{\pi}_{\bm{\theta}}(\bm{\tau})\nabla_{\bm{\theta}}\text{log}\bm{\pi}_{\bm{\theta}}(\bm{\tau})$,
we obtain:

$\begin{split}&\nabla_{\theta^{m}}\mathcal{J}^{m}(\bm{\theta})=\int\bm{\pi}_{\bm{\theta}}(\bm{\tau})\nabla_{\theta^{m}}\text{log}\bm{\pi}_{\bm{\theta}}(\bm{\tau})G^{m}d\bm{\tau}\cr&=\int\bm{\pi}_{\bm{\theta}}(\bm{\tau})\nabla_{\theta^{m}}\text{log}[q(s_{0})\prod^{T}_{t=0}\rho(s_{t+1}|s_{t},\bm{a}_{t})\cr&\times\prod^{M}_{m=1}\pi_{\theta^{m}}(a^{m}_{t}|\bm{h}_{t})\omega(\bm{o}_{t}|s_{t})]G^{m}d\bm{\tau}\cr&=\mathbb{E}_{\bm{\pi}_{\theta}(\bm{\tau})}\Biggl{[}\sum^{T}_{t=0}\nabla_{\theta^{m}}\text{log}\pi_{\theta^{m}}(a^{m}_{t}|\bm{h}_{t})\cr&\times\left(\sum^{T}_{t^{{}^{\prime}}=t}\gamma^{t^{{}^{\prime}}-t}r^{m}(\bm{a}_{t^{{}^{\prime}}},s_{t^{{}^{\prime}}})\right)\Biggr{]}\cr\end{split}$
(13)

##  Appendix C Policy initialization through supervised learning

While training in RL can be performed with a policy whose parameters are
arbitrarily initialized, this is not ideal. Having an adequate initial guess
of the policy and using RL subsequently to only fine tune the policy is a much
easier problem to solve. A quasi-random Halton sequence [60] with equally
spaced probe positions is a reasonable initialization. Pre-training of the
parameterized policy for the RL model can then be performed by supervised
learning applied on the RNN such that the discrepancy between the predicted
scan positions
$\vec{\textbf{R}}_{P_{t}}=\vec{\bm{\mu}}_{\bm{\theta}}(\bm{h}_{t})$ and the
scan positions of the initialization sequence
$\vec{\textbf{R}}^{\text{init}}_{P_{t}}$ is minimized:

$\mathcal{K}(\bm{\theta})=\sum^{T}_{t}\|\vec{\textbf{R}}_{P_{t}}-\vec{\textbf{R}}^{\text{init}}_{P_{t}}\|^{2}_{2}.$
(14)

Figure 7 illustrates the scan positions during the fine tuning of the policy
through RL for the first 10,000 iterations when either a) a policy that has
not been initialized via supervised learning or b) an initialized policy is
used. While the scan positions in both cases converge to the atomic structure,
the positions predicted by the non-initialized policy are distributed only
within a small region of the field of view during the entire training.

Figure 7: Fine tuning of a policy with RL that a) has not been initialized and
b) has been initialized via supervised learning. Positions A indicate the scan
positions of the first sub-sequence $\vec{\textbf{R}}_{P_{0}}$ that is
provided to the RNN as part of the initial input. Positions B and C are the
scan positions of all predicted sub-sequences at iteration $0$ and 10,000,
respectively. The trajectories they form during the optimization process are
indicated by a dashed blue lines.

##  Appendix D Settings

The various settings used for ROP, the convolutional autoencoder, and the RNN
can be found in Table 2.

Table 2: Settings for the potential reconstruction with ROP, structure content compression with the convolutional autoencoder and prediction of the scan sequences with the RNN. ROP |
---|---
acceleration voltage (kV) | $60$
convergence angle (mrad) | $33$
object dimension (px) | $200$
real space pixel size (nm) | $0.0154$
diffraction pattern dimension (px) | $64$
reciprocal space pixel size $\left(\text{nm}^{-1}\right)$ | $0.4759$
scanning step size (nm) | $0.02$
iterations | $5$
step size $\alpha_{\text{ROP}}$ | $5.25\text{{e}}2$
batch size | $24$
Conv. autoencoder |
input dimension (px) | $512$
pixel size (nm) | $0.0064$
encoder/decoder kernel sizes (px) | $[3,3,3,3,3,3]$
encoder/decoder kernel strides (px) | $[1,1,1,1,1,1]$
encoder output channels | $[16,32,64,128,256,512]$
decoder output channels | $[256,128,64,32,16,2]$
iterations | $100000$
step size $\alpha_{\text{CAE}}$ | $1\text{{e}-}5$
batch size | $24$
RNN |
sequence length | $250$
sub-sequence length | $50$
hidden state $H_{t}$ size | $2048$
stacked GRU layers | $2$
iterations | $20000$
step size $\alpha_{\text{RNN}}$ | $1\text{{e}-}6$
batch size | $24$

##  Appendix E Examples

Several different exemplary reconstructions using adaptive scanning can be
seen in Figure 8.

Figure 8: Ptychographic reconstructions of different MoS2 data sets and with
different scanning procedures. Reconstruction from $250$ diffraction patterns
of a data set that correspond to scan positions which follow a)-d) a random
sequence and e)-h) an adaptively predicted sequence. e)-h) Ground truth
reconstruction of the full data set with 10,000 diffraction patterns shown
with the scan positions used for the corresponding reconstructions a)-d) in
green and e)-h) in red.

## References

  * Rodenburg [2008] J. M. Rodenburg, Ptychography and related diffractive imaging methods, Advances in imaging and electron physics 150, 87 (2008).
  * Humphry  _et al._ [2012] M. Humphry, B. Kraus, A. Hurst, A. Maiden, and J. Rodenburg, Ptychographic electron microscopy using high-angle dark-field scattering for sub-nanometre resolution imaging, Nature communications 3, 1 (2012).
  * Lozano  _et al._ [2018] J. G. Lozano, G. T. Martinez, L. Jin, P. D. Nellist, and P. G. Bruce, Low-dose aberration-free imaging of li-rich cathode materials at various states of charge using electron ptychography, Nano letters 18, 6850 (2018).
  * Jiang  _et al._ [2018a] Y. Jiang, Z. Chen, Y. Han, P. Deb, H. Gao, S. Xie, P. Purohit, M. W. Tate, J. Park, S. M. Gruner, _et al._ , Electron ptychography of 2d materials to deep sub-ångström resolution, Nature 559, 343 (2018a).
  * Zhou  _et al._ [2020] L. Zhou, J. Song, J. S. Kim, X. Pei, C. Huang, M. Boyce, L. Mendonça, D. Clare, A. Siebert, C. S. Allen, _et al._ , Low-dose phase retrieval of biological specimens using cryo-electron ptychography, Nature communications 11, 1 (2020).
  * McMullan  _et al._ [2016] G. McMullan, A. Faruqi, and R. Henderson, Direct electron detectors, Methods in enzymology 579, 1 (2016).
  * Tate  _et al._ [2016] M. W. Tate, P. Purohit, D. Chamberlain, K. X. Nguyen, R. Hovden, C. S. Chang, P. Deb, E. Turgut, J. T. Heron, D. G. Schlom, _et al._ , High dynamic range pixel array detector for scanning transmission electron microscopy, Microscopy and Microanalysis 22, 237 (2016).
  * Thibault  _et al._ [2008] P. Thibault, M. Dierolf, A. Menzel, O. Bunk, C. David, and F. Pfeiffer, High-resolution scanning x-ray diffraction microscopy, Science 321, 379 (2008).
  * Maiden and Rodenburg [2009] A. M. Maiden and J. M. Rodenburg, An improved ptychographical phase retrieval algorithm for diffractive imaging, Ultramicroscopy 109, 1256 (2009).
  * Van den Broek and Koch [2012] W. Van den Broek and C. T. Koch, Method for retrieval of the three-dimensional object potential by inversion of dynamical electron scattering, Physical review letters 109, 245502 (2012).
  * Maiden  _et al._ [2012] A. M. Maiden, M. J. Humphry, and J. Rodenburg, Ptychographic transmission microscopy in three dimensions using a multi-slice approach, JOSA A 29, 1606 (2012).
  * Van den Broek and Koch [2013] W. Van den Broek and C. T. Koch, General framework for quantitative three-dimensional reconstruction from arbitrary detection geometries in tem, Physical Review B 87, 184108 (2013).
  * Tsai  _et al._ [2016] E. H. Tsai, I. Usov, A. Diaz, A. Menzel, and M. Guizar-Sicairos, X-ray ptychography with extended depth of field, Optics express 24, 29089 (2016).
  * Jiang  _et al._ [2018b] Y. Jiang, Z. Chen, I. El Baggari, L. F. Kourkoutis, V. Elser, and D. A. Muller, Breaking the rayleigh limit in thick samples with multi-slice ptychography, Microscopy and Microanalysis 24, 192 (2018b).
  * Pelz  _et al._ [2017] P. M. Pelz, W. X. Qiu, R. Bücker, G. Kassier, and D. R. J. Miller, Low-dose cryo electron ptychography via non-convex bayesian optimization, Scientific Reports 7, 9883 (2017).
  * Song  _et al._ [2019] J. Song, C. S. Allen, S. Gao, C. Huang, H. Sawada, X. Pan, J. Warner, P. Wang, and A. I. Kirkland, Atomic resolution defocused electron ptychography at low dose with a fast, direct electron detector, Scientific reports 9, 1 (2019).
  * Schloz  _et al._ [2020] M. Schloz, T. C. Pekin, Z. Chen, W. Van den Broek, D. A. Muller, and C. T. Koch, Overcoming information reduced data and experimentally uncertain parameters in ptychography with regularized optimization, Optics Express 28, 28306 (2020).
  * Chen  _et al._ [2020] Z. Chen, M. Odstrcil, Y. Jiang, Y. Han, M.-H. Chiu, L.-J. Li, and D. A. Muller, Mixed-state electron ptychography enables sub-angstrom resolution imaging with picometer precision at low dose, Nature communications 11, 1 (2020).
  * Huang  _et al._ [2014] X. Huang, H. Yan, R. Harder, Y. Hwu, I. K. Robinson, and Y. S. Chu, Optimization of overlap uniformness for ptychography, Optics Express 22, 12634 (2014).
  * Mnih  _et al._ [2013] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller, Playing atari with deep reinforcement learning, arXiv preprint arXiv:1312.5602 (2013).
  * Silver  _et al._ [2016] D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, _et al._ , Mastering the game of go with deep neural networks and tree search, nature 529, 484 (2016).
  * Levine  _et al._ [2016] S. Levine, C. Finn, T. Darrell, and P. Abbeel, End-to-end training of deep visuomotor policies, The Journal of Machine Learning Research 17, 1334 (2016).
  * Andrychowicz  _et al._ [2020] O. M. Andrychowicz, B. Baker, M. Chociej, R. Jozefowicz, B. McGrew, J. Pachocki, A. Petron, M. Plappert, G. Powell, A. Ray, _et al._ , Learning dexterous in-hand manipulation, The International Journal of Robotics Research 39, 3 (2020).
  * Sharma  _et al._ [2015] S. Sharma, R. Kiros, and R. Salakhutdinov, Action recognition using visual attention, arXiv preprint arXiv:1511.04119 (2015).
  * Ba  _et al._ [2014] J. Ba, V. Mnih, and K. Kavukcuoglu, Multiple object recognition with visual attention, arXiv preprint arXiv:1412.7755 (2014).
  * Mnih  _et al._ [2014] V. Mnih, N. Heess, A. Graves, _et al._ , Recurrent models of visual attention, in  _Advances in neural information processing systems_ (2014) pp. 2204–2212.
  * Betterton  _et al._ [2020] J.-R. Betterton, D. Ratner, S. Webb, and M. Kochenderfer, Reinforcement learning for adaptive illumination with x-rays, in  _2020 IEEE International Conference on Robotics and Automation (ICRA)_ (IEEE, 2020) pp. 328–334.
  * Vasudevan  _et al._ [2020] R. K. Vasudevan, K. Kelley, H. Funakubo, S. Jesse, S. V. Kalinin, and M. Ziatdinov, Autonomous experiments in scanning probe microscopy and spectroscopy: Choosing where to explore polarization dynamics in ferroelectrics, arXiv preprint arXiv:2011.13050 (2020).
  * Dahmen  _et al._ [2016] T. Dahmen, M. Engstler, C. Pauly, P. Trampert, N. De Jonge, F. Mücklich, and P. Slusallek, Feature adaptive sampling for scanning electron microscopy, Scientific reports 6, 1 (2016).
  * Kalinin  _et al._ [2021] S. V. Kalinin, M. A. Ziatdinov, J. Hinkle, S. Jesse, A. Ghosh, K. P. Kelley, A. R. Lupini, B. G. Sumpter, and R. K. Vasudevan, Automated and autonomous experiment in electron and scanning probe microscopy, arXiv preprint arXiv:2103.12165 (2021).
  * O’Mara  _et al._ [2016] J. O’Mara, B. Meredig, and K. Michel, Materials data infrastructure: a case study of the citrination platform to examine data import, storage, and access, Jom 68, 2031 (2016).
  * Draxl and Scheffler [2018] C. Draxl and M. Scheffler, Nomad: The fair concept for big data-driven materials science, Mrs Bulletin 43, 676 (2018).
  * Draxl and Scheffler [2019] C. Draxl and M. Scheffler, The nomad laboratory: from data sharing to artificial intelligence, Journal of Physics: Materials 2, 036001 (2019).
  * Dillen  _et al._ [2019] M. Dillen, Q. Groom, D. Agosti, and L. H. Nielsen, Zenodo, an archive and publishing repository: A tale of two herbarium specimen pilot projects., Biodiversity Information Science and Standards (2019).
  * Rumelhart  _et al._ [1986] D. E. Rumelhart, G. E. Hinton, and R. J. Williams, Learning representations by back-propagating errors, nature 323, 533 (1986).
  * Elman [1990] J. L. Elman, Finding structure in time, Cognitive science 14, 179 (1990).
  * Werbos [1988] P. J. Werbos, Generalization of backpropagation with application to a recurrent gas market model, Neural networks 1, 339 (1988).
  * Masci  _et al._ [2011] J. Masci, U. Meier, D. Cireşan, and J. Schmidhuber, Stacked convolutional auto-encoders for hierarchical feature extraction, in  _International conference on artificial neural networks_ (Springer, 2011) pp. 52–59.
  * Bengio  _et al._ [1994] Y. Bengio, P. Simard, and P. Frasconi, Learning long-term dependencies with gradient descent is difficult, IEEE transactions on neural networks 5, 157 (1994).
  * Hochreiter and Schmidhuber [1997] S. Hochreiter and J. Schmidhuber, Long short-term memory, Neural computation 9, 1735 (1997).
  * Chung  _et al._ [2014] J. Chung, C. Gulcehre, K. Cho, and Y. Bengio, Empirical evaluation of gated recurrent neural networks on sequence modeling, arXiv preprint arXiv:1412.3555 (2014).
  * Williams [1992] R. J. Williams, Simple statistical gradient-following algorithms for connectionist reinforcement learning, Machine learning 8, 229 (1992).
  * Liu  _et al._ [2020] Y. Liu, K. Zhang, T. Basar, and W. Yin, An improved analysis of (variance-reduced) policy gradient and natural policy gradient methods., in  _NeurIPS_ (2020).
  * Sutton  _et al._ [2000] R. S. Sutton, D. A. McAllester, S. P. Singh, and Y. Mansour, Policy gradient methods for reinforcement learning with function approximation, in  _Advances in neural information processing systems_ (2000) pp. 1057–1063.
  * Wu  _et al._ [2018] C. Wu, A. Rajeswaran, Y. Duan, V. Kumar, A. M. Bayen, S. Kakade, I. Mordatch, and P. Abbeel, Variance reduction for policy gradient with action-dependent factorized baselines, arXiv preprint arXiv:1803.07246 (2018).
  * Feriani and Hossain [2021] A. Feriani and E. Hossain, Single and multi-agent deep reinforcement learning for ai-enabled wireless networks: A tutorial, IEEE Communications Surveys & Tutorials (2021).
  * Agogino and Tumer [2004] A. K. Agogino and K. Tumer, Unifying temporal and structural credit assignment problems, in  _Autonomous Agents and Multi-Agent Systems Conference_ (2004).
  * Di Castro  _et al._ [2012] D. Di Castro, A. Tamar, and S. Mannor, Policy gradients with variance related risk criteria, arXiv preprint arXiv:1206.6404 (2012).
  * Wolpert and Tumer [2002] D. H. Wolpert and K. Tumer, Optimal payoff functions for members of collectives, in  _Modeling complexity in economic and social systems_ (World Scientific, 2002) pp. 355–369.
  * Tumer and Agogino [2007] K. Tumer and A. Agogino, Distributed agent-based air traffic flow management, in  _Proceedings of the 6th international joint conference on Autonomous agents and multiagent systems_ (2007) pp. 1–8.
  * Colby  _et al._ [2014] M. K. Colby, W. Curran, C. Rebhuhn, and K. Tumer, Approximating difference evaluations with local knowledge, in  _Proceedings of the 2014 international conference on Autonomous agents and multi-agent systems_ (2014) pp. 1577–1578.
  * Castellini  _et al._ [2020] J. Castellini, S. Devlin, F. A. Oliehoek, and R. Savani, Difference rewards policy gradients, arXiv preprint arXiv:2012.11258 (2020).
  * Okabe [2016] A. Okabe, Spatial tessellations, International Encyclopedia of Geography: People, the Earth, Environment and Technology: People, the Earth, Environment and Technology , 1 (2016).
  * Rong and Tan [2006] G. Rong and T.-S. Tan, Jump flooding in gpu with applications to voronoi diagram and distance transform, in  _Proceedings of the 2006 symposium on Interactive 3D graphics and games_ (2006) pp. 109–116.
  * Paszke  _et al._ [2019] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, Pytorch: An imperative style, high-performance deep learning library, in  _Advances in Neural Information Processing Systems 32_, edited by H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (Curran Associates, Inc., 2019) pp. 8024–8035.
  * Kingma and Ba [2014] D. P. Kingma and J. Ba, Adam: A method for stochastic optimization, arXiv preprint arXiv:1412.6980 (2014).
  * Wang  _et al._ [2004] Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli, Image quality assessment: from error visibility to structural similarity, IEEE transactions on image processing 13, 600 (2004).
  * LeCun  _et al._ [1989] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel, Backpropagation applied to handwritten zip code recognition, Neural computation 1, 541 (1989).
  * Hinton and Salakhutdinov [2006] G. E. Hinton and R. R. Salakhutdinov, Reducing the dimensionality of data with neural networks, science 313, 504 (2006).
  * Halton [1964] J. H. Halton, Algorithm 247: Radical-inverse quasi-random point sequence, Communications of the ACM 7, 701 (1964).


# Variational auto-encoders with Student’s t-prior

Najmeh Abiri
<EMAIL_ADDRESS>Department of Astronomy and Theoretical Physics, Lund
University, Lund, Sweden Mattias Ohlsson 11footnotemark: 1
<EMAIL_ADDRESS>Center for Applied Intelligent Systems Research, Halmstad
University, Halmstad, Sweden

###### Abstract

We propose a new structure for the variational auto-encoders (VAEs) prior,
with the weakly informative multivariate Student’s t-distribution. In the
proposed model all distribution parameters are trained, thereby allowing for a
more robust approximation of the underlying data distribution. We used
Fashion-MNIST data in two experiments to compare the proposed VAEs with the
standard Gaussian priors. Both experiments showed a better reconstruction of
the images with VAEs using Student’s t-prior distribution.

##  1 Introduction

Variational auto-encoders (VAEs) [1] are complex generative models that use
variational inference with auto-encoder [2] like architecture to learn the
underlying representation. Unlike auto-encoders which only learn the
functional representation of data, VAEs as Bayesian approaches, learn the data
probability distribution, and have shown great outcome in modeling complex
distributions on increasingly larger data.

One way to improve any Bayesian model is to change the prior distribution
depend on the data. Although Kingma and Welling [1] suggested any prior from
location-scale family distribution for VAEs models, most of the studies have
been done only on Gaussian priors. We show, with reformulating Student’s
t-distribution [3] and by using implicit differentiation [4], it is possible
to use the Student’s t-prior. In the case of the presence of outliers in the
data, the adaptive tail thickness can give enormous flexibility to the latent
model.

##  2 Proposed approach

A suitable prior is essential in Bayesian analysis. In lack of knowledge of
data, weakly or non-informative priors are recommended and have shown better
results in many studies [5], e.g., Student’s t-distribution. With knowledge,
an informative distribution like the normal distribution with specified mean
and covariance can be used. Although this prior has shown excellent results in
many studies, O’Hagan [6] showed that normal priors could be outlier-resistant
and never reject outliers in data modeling. Student’s t-distributions allow
for more heavy tails and can reject outliers when used as priors in the
modeling. This motivates an implementation of the VAE with a Student’s
t-prior.

A location shifted p-variant Student’s t-distribution has the probability
density function

$\displaystyle
f(z)=\frac{\Gamma(\frac{\nu+p}{2})}{\Gamma(\frac{\nu}{2})\sqrt{(\pi\nu)^{p}}{\sigma}}\bigg{(}1+\frac{(z-\mu)^{T}\sigma^{-2}(z-\mu)}{\nu}\bigg{)}^{-\frac{\nu+p}{2}},$

with $\nu$ the degree of freedom, $\mu$ the location parameter and $\sigma$
scale parameter of the distribution.

###  2.1 Architecture of VAE with Student’s t-distribution

As explained in Kingma and Welling [1], the objective of VAE is

$\displaystyle\max_{\theta,\phi}[\mathop{{}\mathbb{E}}_{q_{\phi}(\mathbf{z}|\mathbf{x})}\left[\log(p_{\theta}(\mathbf{x}|\mathbf{z}))\right]-\mathbb{D_{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z}))],$
(1)

with the expected likelihood term as a reconstruction loss and the Kullback-
Leibler divergence term as a regularizer factor.

The proposed structure of the VAE is illustrated in the figure 1.

Figure 1: VAE architecture with Student’s t-prior. The outputs of the encoder
part are the three parameters of the Student’s t-distribution.

To be able to consider the Student’s t-distribution as a prior for the VAE we
need to adjust the objective of the network (the equation 1), to have a
differentiable transformation of the probability
$q_{\phi}(\mathbf{z}|\mathbf{x})$ w.r.t. $\phi$.

Student’s t-distribution belongs to location-scale family distributions and it
is possible to use the reparameterization trick,

$\displaystyle X\sim\operatorname{St}(\mu,\sigma,\nu),\;X=\mu+\sigma
T,\;T\sim\operatorname{St}(0,I,\nu).$

Still, the replaced probability distribution is not a differentiable
transformation. Another reformulation of this distribution [3] is,
$x\sim\mathcal{N}(0,I)$ and $z\sim\text{Gamma}(\frac{\nu}{2},0.5)$ then
$t=x/\sqrt{\frac{z}{\nu}}$ is equivalent to $t\sim\operatorname{St}(0,I,\nu)$.
Figurnov and Mohamed [4] found an implicit differentiation of the cumulative
function of the Gamma distribution. With the above formulation of the
Student’s t-distribution it is now differentiable w.r.t all three parameters.

###  2.2 Training the Student’s t-VAE

In this study the prior over the latent variables is a multivariate Student’s
t-distribution, $p_{\theta}(\mathbf{z})=\operatorname{St}(0,I,\nu)$ and the
variational approximate posterior is
$q_{\phi}(\mathbf{z}|\mathbf{x})=\operatorname{St}(\mu,\sigma,\nu)$. The
parameters; $\mu$, $\sigma$ and $\nu$ are neural networks, as in figure 1,
with the weights $\phi$. The KL divergence of the objective takes the form

$\displaystyle\mathbb{D_{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z}))=\int
q_{\phi}(\mathbf{z}|\mathbf{x})\log\left(\frac{q_{\phi}(\mathbf{z}|\mathbf{x})}{p_{\theta}(\mathbf{z})}\right)d\mathbf{z}$
(2) $\displaystyle=$
$\displaystyle-\log\lvert\sigma\rvert-\left(\frac{\nu+p}{2}\right)\mathop{{}\mathbb{E}}_{q_{\phi}}\left[\log\left(1+\frac{(\mathbf{z}-\mu)^{T}\sigma^{-2}(\mathbf{z}-\mu)}{\nu}\right)\right]$
$\displaystyle\hskip
36.98866pt+\left(\frac{\nu+p}{2}\right)\mathop{{}\mathbb{E}}_{q_{\phi}}\bigg{[}\log\left(1+\frac{z^{2}}{\nu}\right)\bigg{]}.$

Zografos [7] shows that the first expectation in equation 2 is equivalent to

$\displaystyle\psi\left(\frac{\nu+p}{2}\right)-\psi\left(\frac{\nu}{2}\right),$

where $\psi(x)=\frac{{\Gamma^{\prime}(x)}}{\Gamma(x)}$ is the digamma
function. The second expectation in equation 2 is

$\displaystyle\frac{\Gamma(\frac{\nu+p}{2})}{\Gamma(\frac{\nu}{2})\sqrt{(\pi\nu)^{p}}\lvert\sigma\rvert}\int\bigg{[}1+\frac{(\mathbf{z}-\mu)^{T}\sigma^{-2}(\mathbf{z}-\mu)}{\nu}\bigg{]}^{\frac{-\nu+p}{2}}\log\left(1+\frac{\mathbf{z}^{2}}{\nu}\right)dz,$

that requires numerical integration. With the expression of the KL-divergence
and a suitable reconstruction error, taking the type of the data into account,
stochastic gradient methods can now be used to train the models.

##  3 Experiments

We trained VAE with two different priors, a multivariate Gaussian (VAE-G) and
our proposed multivariate Student’s t-prior (VAE-St) on the Fashion-MNIST [8]
dataset. Although Fashion-MNIST has the same shape and grayscale format as
MNIST, the images are more diverse which force the models to learn more
advanced features. In each experiment, the models are trained on $60000$
images and tested on $10000$. During training, a fixed validation dataset (20%
of training data) was used for model selection. To estimate the maximum
likelihood
$\mathop{{}\mathbb{E}}\left[\log(p_{\theta}(\mathbf{x}|\mathbf{z}))\right]$,
we used cross-entropy. In all experiments, computations were performed on a
single GPU (GeForce GTX 1080 Ti).

###  3.1 VAEs with complete data

Although the data are images, both VAE-G and VAE-St used multi-layer
perceptrons (MLP) for both the encoder and the decoder. The choice of MLPs
instead of convolutional neural networks was to simplify the hyperparameter
optimization process. To evaluate the performance of VAEs, the generated
images were compared with the original ones using the structural similarity
[9] (SSIM) index. Unlike mean squared difference that resembles the images on
individual pixels, SSIM treats images more holistically by comparing local
regions of the images. The best models, after extensive random search [10] on
hyperparameters, are shown in table 1.

Model | encoder | decoder | act | lr | batch | epochs | SSIM
---|---|---|---|---|---|---|---
VAE-G | $[1453,44]$ | $[44,702]$ | $\tanh$ | $10^{-3}$ | $500$ | $200$ | $0.829$
VAE-St | $[1419,42]$ | $[42,759]$ | relu | $10^{-4}$ | $500$ | $150$ | $\mathbf{0.949}$
Table 1: The best sets of hyperparameters for each model, together with the
average SSIM index over the test data. The numbers for the encoder and the
decoder are the size of the hidden layer and the dimension of the latent space
(last number in the encoder list and first number in the decoder list).

Examples of images generated by the two VAEs are shown in figure 2. In the
first row, there are $11$ randomly selected images, at least one from each
class, from Fashion-MNIST test set. The second and third rows are the
corresponding outputs of VAE-G and VAE-St respectively.

Figure 2: Generated images from the two VAEs. The first row shows input images
from the test set, row two and three are the corresponding conditionally
generated images from VAE-G and VAE-St, respectively. Note, only one sampled
image for each input is shown. The last column shows the average SSIM index
over the test data.

###  3.2 VAEs with incomplete data

The importance of using Student’s t-distribution as a prior in VAE is even
more evident in the presence of missing data. The Student’s t-prior, with its
adjustable tails, lets the VAEs find more profound features which improve the
estimation of missing data. To evaluate how well VAEs can impute missing data,
we manually corrupted Fashion-MNIST which enables a comparison with original
images. Image data have high local correlations; therefore a random pixel-wise
corruption will not be difficult enough. Instead, we used Square-lattice
Ising-like [11], $28\times 28$ images of masks. In such a mask approximately
$78\%$ of the pixels are corrupted, i.e., replaced by NaN. The second row of
figure 3 shows examples of such masks.

To train the VAE with missing data, we replaced each NaN pixel with the
average over the known samples for that pixel. The loss function is calculated
only on non-missing pixels. Similar to the previous experiment, the VAEs are
evaluated with the SSIM index. The best models, after hyperparameter search,
are shown in table 2.

Model | encoder | decoder | act | lr | batch | epochs | SSIM
---|---|---|---|---|---|---|---
VAE-G | $[905,24]$ | $[24,902]$ | relu | $10^{-3}$ | $500$ | $300$ | $0.678$
VAE-St | $[1139,463,37]$ | $[37,1308]$ | $\tanh$ | $10^{-3}$ | $4000$ | $300$ | $\mathbf{0.831}$
Table 2: The best sets of hyperparameters for each model, together with the
average SSIM index over the test data. The numbers for the encoder and the
decoder are the size of the hidden layers and the dimension of the latent
space (last number in the encoder list and first number in the decoder list).

Figure 3 shows example images from the test set. The first and the third row
show the original (uncorrupted) images and the images after applying the Ising
masks, respectively. The fourth row shows the actual input images to the VAEs
where missing pixels are replaced by mean values over known data (mean
imputation). Mean imputation does not work for this degree and structure of
missing data. Both VAE-G and VAE-St improves over mean imputation as indicated
by the SSIM index. From a visual inspection and as shown by the SSIM index,
the use of a Student’s t-prior is improving the reconstructed images.

Figure 3: The first row shows examples of original test images. Applying the
masks, shown in row two, results in the corrupted images in row three. The
input of the trained VAE-G and VAE-St is the fourth row (average over known
pixels). The fifth and sixth rows show one sampled image for each input of
VAE-G and VAE-St, respectively. The last column shows the average SSIM index
over the test data.

##  4 Conclusion

In this paper, we have implemented a variational auto-encoder with a
multivariate Student’s t-prior distribution which was accomplished using the
implicit reparameterization gradient of the Gamma function [4].Two different
experiments were performed on Fashion-MNIST data with both VAE-G and VAE-St.
In the first experiment, we trained both models with complete Fashion-MNIST
images and evaluated the models on test data. For the second experiment, we
used VAE-G and VAE-St as imputation methods to recover missing pixels on
corrupted Fashion-MNIST data. Both models used cross-entropy as loss function
and their corresponding KL divergence as regularization. With the change of
prior from Gaussian to Student’s t-distribution in the VAE, we were able to
get less blurry and more detailed images for both experiments that were
quantified by the SSIM index being larger for VAE-St in both trials. Although
the computations for the VAE with Student’s t-prior distribution are more
extensive, the difference in execution time on a GPU was not significant. We
believe that using the Student’s t-prior distribution for the VAE allows
modeling of more complex data, such as medical data with a mixture of
different input variable types. A typical challenge in working with medical
data is to impute missing values. A flexible generative model such as the
proposed VAE-St can have a significant impact on imputation, which will be our
focus for future developments.

## References

  * [1] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114, 2013.
  * [2] Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierre-Antoine Manzagol.  Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion.  Journal of Machine Learning Research, 11(Dec):3371–3408, 2010.
  * [3] Catherine Forbes, Merran Evans, Nicholas Hastings, and Brian Peacock.  Statistical distributions.  John Wiley & Sons, 2011.
  * [4] Michael Figurnov, Shakir Mohamed, and Andriy Mnih.  Implicit reparameterization gradients.  arXiv preprint arXiv:1805.08498, 2018.
  * [5] Andrew Gelman, Aleks Jakulin, Maria Grazia Pittau, Yu-Sung Su, et al.  A weakly informative default prior distribution for logistic and other regression models.  The Annals of Applied Statistics, 2(4):1360–1383, 2008.
  * [6] Anthony O’Hagan.  On outlier rejection phenomena in bayes inference.  Journal of the Royal Statistical Society. Series B (Methodological), pages 358–367, 1979.
  * [7] Konstantinos Zografos et al.  On maximum entropy characterization of pearson’s type ii and vii multivariate distributions.  Journal of Multivariate Analysis, 71(1):67–75, 1999.
  * [8] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  arXiv preprint arXiv:1708.07747, 2017.
  * [9] Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli.  Image quality assessment: from error visibility to structural similarity.  IEEE transactions on image processing, 13(4):600–612, 2004.
  * [10] James Bergstra and Yoshua Bengio.  Random search for hyper-parameter optimization.  Journal of Machine Learning Research, 13(Feb):281–305, 2012.
  * [11] Rodney J Baxter.  Exactly solved models in statistical mechanics.  Elsevier, 2016.