# Improved Three-Dimensional Reconstructions in Electron Ptychography through
Defocus Series Measurements

Marcel Schloz<EMAIL_ADDRESS>Institute of Physics and Center for the
Science of Materials Berlin, Humboldt-Universität zu Berlin, Newtonstraße 15,
12489 Berlin, Germany  Thomas C. Pekin  Institute of Physics and Center for
the Science of Materials Berlin, Humboldt-Universität zu Berlin, Newtonstraße
15, 12489 Berlin, Germany  Hamish G. Brown  Ian Holmes Imaging Centre, Bio21
Molecular Science and Biotechnology Institute, University of Melbourne,
Melbourne 3010, Australia  Dana O. Byrne  Department of Chemistry, University
of California, Berkeley, Berkeley, CA 94720, USA  National Center for Electron
Microscopy, Molecular Foundry, Lawrence Berkeley National Laboratory,
Berkeley, CA 94720, USA  Bryan D. Esser  Monash Centre for Electron
Microscopy, Monash University, Melbourne 3800, Australia  Emmanuel Terzoudis-
Lumsden  School of Physics and Astronomy, Monash University, Clayton,
Melbourne 3800, Australia  Takashi Taniguchi  International Center for
Materials Nanoarchitectonics, National Institute for Materials Science, 1-1
Namiki, Tsukuba, 305-0044, Japan  Kenji Watanabe  Research for Functional
Materials, National Institute for Materials Science, 1-1 Namiki, Tsukuba,
305-0044, Japan  Scott D. Findlay  School of Physics and Astronomy, Monash
University, Clayton, Melbourne 3800, Australia  Benedikt Haas  Institute of
Physics and Center for the Science of Materials Berlin, Humboldt-Universität
zu Berlin, Newtonstraße 15, 12489 Berlin, Germany  Jim Ciston  National Center
for Electron Microscopy, Molecular Foundry, Lawrence Berkeley National
Laboratory, Berkeley, CA 94720, USA  Christoph T. Koch  Institute of Physics
and Center for the Science of Materials Berlin, Humboldt-Universität zu
Berlin, Newtonstraße 15, 12489 Berlin, Germany

###### Abstract

A detailed analysis of ptychography for 3D phase reconstructions of thick
specimens is performed. We introduce multi-focus ptychography, which
incorporates a 4D-STEM defocus series to enhance the quality of 3D
reconstructions along the beam direction through a higher overdetermination
ratio. This method is compared with established multi-slice ptychography
techniques, such as conventional ptychography, regularized ptychography, and
multi-mode ptychography. Additionally, we contrast multi-focus ptychography
with an alternative method that uses virtual optical sectioning through a
reconstructed scattering matrix ($\mathcal{S}$-matrix), which offers more
precise 3D structure information compared to conventional ptychography. Our
findings from multiple 3D reconstructions based on simulated and experimental
data demonstrate that multi-focus ptychography surpasses other techniques,
particularly in accurately reconstructing the surfaces and interface regions
of thick specimens.

## Introduction

Electron ptychography is a technique whereby the optical phase (itself a
function of the electrostatic potential) of the sample is directly
reconstructed from a set a diffraction patterns produced by a convergent,
coherent electron probe. This technique has achieved dose-efficient
quantitative phase imaging at sub-Ångstrom resolution, enabling the clear
visualization of atomic structures and nanoscale features ([11]). This
capability has significantly contributed to the understanding of various
material properties, including crystal defects ([8]), surface reconstructions
([14]), and zeolite structures ([31]). However, this remarkable resolution is
limited to the lateral dimension, constraining reliable structural
characterization to 2D-like systems thinner than a few nanometers. Recent
research in integrating ptychography into a tilt-series tomography experiment,
known as ptychographic atomic electron tomography (PAET), has demonstrated,
through simulations ([3]) and subsequent experimental validation ([7, 21]),
the ability to achieve atomic-resolution phase reconstructions across all
three spatial dimensions. However, this gain in depth-resolution comes at the
expense of a considerably more challenging experiment and its subsequent
analysis. Additionally, the tomographic reconstruction algorithm necessitates
specimen projections at a range of tilt angles, provided as single-slice
ptychography reconstructions. This constraint again limits the thickness of
the investigated specimen to a few nanometers. Extending the PAET technique
with multi-slice ptychography holds promise for relaxing the stringent
criteria for tomographic sampling ([10]). This enhancement could potentially
simplify experimental acquisition and broaden the technique’s applicability to
thicker specimens. However, as of now, this capability has not been
demonstrated in electron microscopy.

The thickness limitation of single-slice ptychography is roughly equivalent to
the depth of field (DOF) of the beam ([27]). Reconstructions from specimens
with thicknesses surpassing the beam’s DOF become unreliable due to a
breakdown of the phase object approximation, which assumes that the probe and
specimen interact in a single, infinitesimally thin plane. For thicker
samples, the limited DOF and the increased redistribution of probe intensity
caused by scattering make it necessary to use multi-slice ptychography. Figure
1a) shows projected potentials obtained from single-slice and multi-slice
reconstructions performed on simulated 4D-STEM data of a two layer stack of
hexagonal boron nitride (hBN) specimens. The two layers are twisted with
respect to one another, resulting in Moiré patterns at various thicknesses.
Additional details about the hBN specimens and the experimental parameters
used are provided below. Notably, when the sample thickness exceeds the DOF,
single-slice reconstructions show a decrease in resolution and eventually fail
to capture the Moiré pattern expected due to the misorientation of the two hBN
stacks. In contrast, the resolution remains high in the central slice of the
multi-slice reconstruction, and the Moiré patterns are clearly visible. This
outcome confirms that the thickness limit does not constrain ptychography when
multiple slices are integrated into the model, accounting for multiple
scattering effects adequately.

Figure 1: The break-down of the phase object approximation. a) Ptychographic
phase reconstructions using either a single-slice or a multi-slice model for
simulated stacked hBN sample data with different total thicknesses. For the
multi-slice model the central slice is shown. The DOF in this setting is about
$10.1$ nm. b) Beam profile at $200$ kV and a semi-convergence angle
$\theta_{\text{con}}$ of $21$ mrad. This beam is used for the reconstructions
in a).

While multi-slice ptychography offers the capability to reconstruct thick
specimens, it does come with a drawback: its depth resolution is relatively
inferior compared to its lateral resolution. Experimental assessments have
demonstrated that achieving a depth resolution slightly better than the
aperture-limited depth resolution is feasible ([4, 19]), indicating a
correlation between the depth resolution of multi-slice ptychography and the
maximum recorded diffraction angle by the detector ([22]). However, achieving
even this level of depth resolution can be challenging when dealing with
experimental 4D-STEM data ([6]). In this case, additional experimental
parameters, such as the shape and the partial coherence of the probe and the
scan positions, often have to be retrieved alongside the object. As a
consequence, the number of unknown variables (unknowns) that have to be
recovered in the ptychographic iterative phase retrieval process increase for
a remaining number of known variables (knowns). This decrease in
overdetermination is depicted in the red line in Fig. 2. An inadequate ratio
of this overdetermination negatively effects the computational stability of
the iterative process and leads to incorrect reconstruction results. Measures
to alter the ratio of overdetermination favorably are therefore crucial for
successfully performing multi-slice ptychography on realistic experimental
data.

Incorporating regularization methods into the reconstruction algorithm has
yielded improvements in the conditioning of the reconstruction problem ([26,
24, 28]). One particularly promising approach is the application of a missing-
wedge (MW) regularization, which penalizes high axial frequencies at low
lateral frequencies in the reconstructed atomic potential in reciprocal space
([4]). This regularization thus takes into account the limited influence of
the propagation operator in the multi-slice model on the phase of the
transmitted electron wave at low lateral frequencies. The effect on the
reconstruction result is an elongation of the atoms, a phenomenon also
observed in annular dark-field STEM depth-sectioning ([30]) or tilt-series
tomography ([16]). While the MW regularization helps the algorithm in
addressing the challenges in multi-slice ptychography that arise from the
increased number of unknowns in a 3D reconstruction, it is, in principle, less
favorable for 3D reconstructions that prioritize high depth-resolution. In an
extremely regularized case, the impact of the MW regularization on the
reconstruction is akin to a computational approach where all slices are
constrained to be identical ([24]), resulting in a complete loss of depth
information but an enhanced overdetermination ratio.

Another way of improving the ratio of overdetermination is to obtain
additional experimental data. In this paper, we present an extension of multi-
slice electron ptychography that involves incorporating a 4D-STEM defocus
series. The green line in Fig. 2 illustrates the ameliorating effect of either
applying a regularization strategy or including a 4D-STEM defocus series on
the overdetermination ratio in multi-slice ptychography. In contrast to the
proposed MW regularization strategy, however, the multi-focus ptychography
approach does not negatively impact the achievable depth resolution while
keeping the convergence behaviour of the reconstruction algorithm stable. Our
implementation of multi-focus ptychography leverages the gradient-based
ptychography reconstruction algorithm ROP ([24]). Handling the same type of
data, our method serves as an alternative reconstruction scheme to the
recently introduced $\mathcal{S}$-matrix approach for 3D reconstructions from
multi-focus 4D-STEM data ([2]).

Figure 2: Schematic illustrating the overdetermination ratio for various
multi-slice ptychographic phase retrieval settings. Adding more reconstruction
tasks, such as adjusting probe, scan positions, or multiple probe modes, can
complicate algorithm convergence towards a satisfactory solution. However,
incorporating regularization strategies and additional data from defocus
series measurements can improve conditioning.

## Materials and methods

### Multi-focus 3D electron ptychography

Here, we present a novel adaptation of multi-slice ptychography utilizing a
set of 4D-STEM datasets acquired at different focal planes within a thick
specimen. This experimental approach is akin to conventional STEM depth-
sectioning ([29]), with the distinction that diffraction patterns are recorded
at each scan position using a pixelated detector. The various focal lengths in
the dataset are chosen to ensure adequate overlap between the DOF of
consecutive measurements. By setting a sufficiently large convergence angle to
form an electron beam with a small DOF, one can acquire an extensive defocus
series, introducing robust redundancy along the axial direction. In the
ptychographic reconstruction algorithm, the entire 4D-STEM defocus series is
subsequently employed for multi-slice reconstruction. This entails a
replacement of the batched loss function, as given in Eq. (15) in Ref. ([24]),
with:

$\mathcal{L}(V,\psi^{0}_{0},...,\psi^{0}_{F},\vec{\textit{{R}}})=\frac{1}{P\times
F}\sum^{F}_{f=1}\sum^{P}_{p=1}\ell(V,\psi^{0}_{f},\textbf{R}_{p}),$ (1)

where the index $f$ corresponds to one of the $F$ defoci used to obtain the
series. The gradients of the batched loss given by Eq. (16 - 18) in Ref.
([24]) follow accordingly. Note that the gradient of the probe is calculated
with respect to each individual defocused probe
$\nabla_{\psi^{0}_{f}}\mathcal{L}(V,\psi^{0}_{0},...,\psi^{0}_{F},\vec{\bm{R}})$.
The wavefunction $\psi^{0}_{f}$ of each 4D-STEM measurement passes through the
network model separately and thus the partial derivatives formed by the
backward propagation through the model derived in Eq. (12 - 13) in Ref. ([24])
independently apply to the gradient of each defocused probe.

In traditional STEM depth-sectioning, a reduced DOF enhances depth resolution
by localizing the electron probe, and hence the bulk of the diffraction
pattern’s information content, to a specific depth range within the specimen.
This increased sensitivity to depth variations is particularly useful for
detailed investigations of the specimen’s structure along the optical axis.
Conversely, ptychography does not depend on traditional depth-sectioning to
achieve significant depth sensitivity. Since some 3D information is always
encoded in the probe wavefunction, in principle, pytchography should be able
to reconstruct 3D information about the specimen even from a scan at a single
defocus value.

Nevertheless, there are several reasons why multi-focus ptychography could
enhance the quality of reconstructions. One contributing factor is the
increased number of knowns or rather constraints in this approach, which plays
a crucial role in improving reconstruction quality along the optical axis.
This becomes particularly significant when paired with enhanced slice sampling
in the multi-slice reconstruction, leading to an increase in the number of
unknowns. An improved overdetermination ratio could help mitigate any
potential ambiguity in the reconstructed atomic 3D potential.

Another factor lies in the utilization of 4D-STEM data in multi-focus
ptychography, ensuring that all regions of the specimen have been in-focus at
least once during the measurement. This aspect is critical as focused-probe
ptychography demonstrates greater resilience in scenarios where the
illuminating probe is affected by incoherence and/or aberrations. In specimens
with a thickness of several tens of nanometers, conventional ptychography
leads to out-of-focus regions of the specimen. In the appendix A of this
paper, we demonstrate the necessity of performing reconstructions with
multiple probe modes to achieve high reconstruction quality in defocused-probe
ptychography. However, optimizing for multiple modes also increases the number
of unknowns, making it more challenging for the reconstruction algorithm to
generate a reliable solution. This can again be mitigated by the increase in
the number of knowns when 4D STEM data from multiple defocus values are
available

###  $\mathcal{S}$-matrix optical sectioning

This section explores an alternative to multi-focus ptychography for
reconstructing 3D structures using a 4D-STEM defocus series: the parallax
reconstruction based on the $\mathcal{S}$-matrix technique ([20, 2, 25]). This
approach retrieves the atomic potential of thick specimens, which may exhibit
heterogeneity along the optical axis, through a two-step process. Initially,
the complex elements of the $\mathcal{S}$-matrix are computed from a series of
intensity measurements through iterative phase retrieval. Subsequently, the
calculated $\mathcal{S}$-matrix is used to generate a virtual through-focal
phase series, enabling the determination of a sample’s structure in 3D.

Suppose we have reconstructed the matrix components
$\mathcal{S}_{\textbf{r}_{\perp},\textbf{h}}$, an $\mathcal{S}$-matrix
representation transforming an input plane wave into a real-space exit
wavefunction. Assuming the projected atomic potential
$V^{z_{0}}(\textbf{r}_{\perp})$ to be in a layer at depth $z_{0}$ of the
specimen, we can describe the scattering process as a free-space propagation
to $z_{0}$, followed by a phase object interaction and further propagation
through the remaining distance $\delta z=z_{1}-z_{0}$ gives111In the special
case where the projected potential is located in the mid-plane of the
specimen, i.e. $z_{0}=z_{1}/2$, Eq. (2) is identical to a variant of the phase
object approximation obtained by using the second-order Strang splitting
method ([9]).

$\mathcal{S}_{\textbf{r}_{\perp},\textbf{h}}=\mathcal{P}(\textbf{r}_{\perp},\delta
z)\otimes_{\textbf{r}_{\perp}}\left[e^{i\sigma
V^{z_{0}}(\textbf{r}_{\perp})}e^{-i\pi\lambda h^{2}z_{0}}e^{2\pi
i\textbf{h}\cdot\textbf{r}_{\perp}}\right].$ (2)

Having obtained the $\mathcal{S}$-matrix through phase retrieval, we can
calculate the transmission function at depth $z_{0}$ by inverting Eq. (2) to:

$e^{i\sigma V^{z_{0}}(\textbf{r}_{\perp})}=\sum_{\textbf{g},\textbf{h}}e^{2\pi
i\textbf{g}\cdot\textbf{r}_{\perp}}e^{i\pi\lambda\delta
zg^{2}}\mathcal{S}_{\textbf{g},\textbf{h}}e^{i\pi\lambda z_{0}h^{2}}e^{-2\pi
i\textbf{h}\cdot\textbf{r}_{\perp}}.$ (3)

While formulated with the assumption of a sample confined to a single plane
$z_{0}$, the parallax method suggests that the phase component in Eq. (3)
provides an approximation of the structure at various depths $z$ and for
various sample thicknesses $\Delta z$. It is important to note that this
approximation is not exact. Each slice is treated as a strong phase object and
is assumed to affect the measured 4D-STEM intensity independently ([1]).
Despite this assumption of independent contributions from each plane, the
parallax reconstruction process at a specific depth, $z$, actually includes
elements from multiple planes. However, the influence of elements outside this
plane becomes progressively less significant ([2]). It is also worth noting
that errors in determining the scattering matrix (an iterative process) and
limitations due to the approximation of Eq. (3) will both be present in
practice and are hard to distinguish.

### Settings for simulated and experimental data

We explored and compared various multi-slice ptychography approaches
(conventional multi-slice ptychography, regularized multi-slice ptychography,
and multi-focus ptychography) along with the $\mathcal{S}$-matrix optical
sectioning method. We initiated the comparison by applying these multi-slice
ptychography methods to an experimental 4D-STEM defocus series dataset
obtained previously as described in ([2]) from a hetero-structure sample
comprising lead iridate Pb2Ir2O7 (PIO) and yttrium-stabilized zirconia
Y0.095Zr0.905O2 (YSZ). The estimated thickness of the sample is $200$ Å,
composed of approximately $50$ Å of PIO and $150$ Å of YSZ. Experimental
conditions at the double aberration-corrected Thermo Fisher Scientific Titan
80-300 microscope involved an accelerating voltage of $300$ kV, a $20$ mrad
condenser aperture semi-angle, probe steps of $0.21$ Å with a dwell time of
$0.874$ ms, and a beam current of $2.01$ pA. Four datasets were acquired using
a Gatan K3 direct-electron detector with defocus values $\Delta f$ set to
$5.3$ nm, $0$ nm, $-6.2$ nm, and $-14.3$ nm relative to the specimen surface
and where we use a positive value of $\Delta f$ to indicate “overfocus”, i.e.
a focus point closer to the electron source. The $\mathcal{S}$-matrix approach
has previously been applied to reconstruct both light and heavy atoms in this
thick sample ([2]). In this study, alignment of the 4D-STEM datasets was
performed using simultaneously acquired HAADF STEM images. The atomic
positions of Pb and Ir atoms, as determined from these images, were adjusted
to form a distortion-free grid, and this transformation was subsequently
applied to the scan positions $\textbf{R}_{p}$ of the simultaneously-acquired
4D-STEM scan.

Besides the PIO-YSZ sample, we conducted both a simulation and an experiment
on a Van-der-Waals structure comprising two twisted multi-layered hBN units
stacked on top of each other. The simulated structure has a total thickness of
approximately $220$ Å and the two stacks were rotated by $10^{\circ}$, while
the estimated thickness of the real sample is around $170$ Å and the estimated
rotation angle is approximately $12^{\circ}$. For the experimental sample, the
twisted structure was created by folding the multilayer hBN unit onto itself,
ensuring that each hBN stack has the same thickness. Figure 3 provides a
schematic representation of this structure. This structure serves as an ideal
test sample due to its distinct layers of light atoms. However, due to the
sensitivity of hBN to knock-on damage ([13]), experiments were preferably
performed at a limited electron energy. For the experimental investigation,
ptychographic datasets were acquired using a NION HERMES microscope
(aberration corrected STEM) at $60$ kV (wavelength of $4.87$ pm) acceleration
voltage, a $40$ mrad convergence angle $\theta_{\text{con}}$, and a Dectris
ELA direct electron detector mounted at the electron energy loss spectroscopy
(EELS) camera port. Distortions induced by the EEL spectrometer were corrected
using in-house developed software. During data acquisition, a dwell time of
$2.0$ ms and a beam current of $19.0$ pA were employed. The defoci of the
series ranged from $-12$ nm to $12$ nm with an increment of $4$ nm. For each
dataset in the series, a conventional grid scanning procedure with a scanning
step size of $0.3$ Å was used. Similar to the PIO-YSZ sample, alignment of the
4D-STEM defocus series was performed using simultaneously acquired HAADF-STEM
images. However, in this case a method for non-rigid registration of images
([12]) was extended to a series of 4D-STEM data sets ([18]). The deformations
found from synchronously acquired HAADF data were applied to the 4D-STEM data
sets by flattening the diffraction space, registering each real-space slice
and transforming back to 4D. For the simulation, datasets were generated using
abTEM ([15]), choosing an acceleration voltage of $60$ kV and a convergence
semi-angle $\theta_{\text{con}}$ of $36$ mrad. $106\times 106$ pixel
diffraction patterns were recorded, and a dense grid scan of $100\times 100$
probe positions with a step size $\Delta x$ of $0.2$ Å was chosen. The defocus
of the series ranged from $-10$ nm to $10$ nm with an increment of $5$ nm,
which sampled the transition between layers sufficiently finely to avoid jumps
in the registration due to the change in structure. For the simulated data
used in Fig. 1, the acceleration voltage has been set to $200$ kV and the
convergence angle to $21$ mrad, corresponding to a DOF of approximately $10.1$
nm. The diffraction patterns were of size $144\times 144$ pixel, and a dense
grid scan of $67\times 67$ probe positions with a step size $\Delta x$ of
$0.3$ Å was used. In all multi-slice reconstructions, 30 slices with a slice
distance of $1$ nm were employed. For the simulation of both datasets, a
consideration of thermal diffuse scattering (TDS) by using a frozen phonon
model in the simulation tool was not included. To ensure comparability with
the $\mathcal{S}$-matrix method, the probe positions in the reconstructions
were not optimised using the ptychography algorithm.

Figure 3: Model of the stacked multi-layer hBN structure. The simulated hBN
specimen shown from top and side, respectively. The 4D-STEM dataset used for
the reconstructions is generated from only a part of this specimen as
indicated by the scan area, shown as a red line. To better illustrate the
formation of the Moiré pattern in twisted hBN stacks, a VESTA model ([17]) has
been added to the top-view representation.

## Results

###  Reconstruction from experimental PIO-YSZ data

In this analysis, we conducted a comparative study of the phase
reconstructions obtained through conventional multi-slice ptychography,
regularized ptychography, multi-focus ptychography, multi-mode ptychography
and the $\mathcal{S}$-matrix optical sectioning method. For all ptychography
techniques, the atomic potential was reconstructed in ten distinct slices
separated by $2.5$ nm in the z-direction. In the conventional multi-slice,
regularized and multi-mode approach, only the 4D-STEM dataset acquired at a
defocus of $5.3$ nm was utilized. Figure 4 displays four out of the ten slices
for all ptychography reconstructions. In the same figure four slices generated
by the $\mathcal{S}$-matrix optical sectioning approach that correspond to the
depth localization of the ptychography slices are presented.

In the conventional ptychography reconstruction, Fig. 4a) clearly shows the
termination of the PIO layer, Fig. 4b) indicates the transition between PIO
and YSZ with low phase contrast, and Fig. 4c) reveals the uniform crystal
structure of YSZ. However, the reduced quality of the reconstructed YSZ
crystal lattice in Fig. 4d) suggests challenges in retrieving slices far from
the focal plane. Fig 4g) and h) demonstrate an enhancement in the
reconstruction quality for YSZ crystal lattice slices through MW-
regularization, although at the expense of a less distinct transition between
the compounds in Fig. 4f). Multi-focus ptychography, without MW-
regularization, overcomes the challenges present in both multi-slice
ptychography reconstructions. As shown in Fig. 4i)-l), the high-quality
reconstruction reveals detailed structures of each compound and significantly
improves the visibility of the transition between the compounds in the hetero-
structure. Reconstruction results of multi-mode ptychography are displayed in
Fig. 4m)-p). Compared to conventional multi-slice ptychography, this method
better resolves the atomic structure of both compounds. However, similar to
regularized ptychography, the interface layer shown in n) clearly reveals
atomic features, which complicates identifying the transition between the two
compounds in the 3D reconstruction. Fig. 4q)-t) displays the reconstruction
results obtained with the optical sectioning $\mathcal{S}$-matrix technique.
Notably, atoms in all slices appear broader than in the three ptychography
results, and some slices exhibit artificial blurring among the atoms. Fig. 4q)
representing the crystal structure of the PIO compound is unclear, as it may
correspond to the structure of the YSZ compound. However, the appearance of a
dip in the reconstructed phase close to the atomic position suggests the
presence of high Z atomic columns, i.e. Pb and Ir atoms (Z=82 and Z=77),
rather than the lighter Y and Zr atoms (Z=39 and Z=40). Nevertheless, a
potential elongation of the YSZ structure to the slice in Fig. 4q) might be
possible, given the presence of this structure in the transition slice in Fig.
4r). Thus, while the $\mathcal{S}$-matrix optical sectioning method reveals
qualitative differences between the layers, as shown previously ([2]), the
depth resolution is lower than that achieved in multi-slice ptychography.

Figure 4: Ptychographic 3D reconstructions from experimental data of the PIO-
YSZ hetero-structure. a)-d) Reconstructions from multi-slice ptychography
using a single dataset with a defocus $\Delta f$ of 5.3 nm. The scanned area
is outlined by a red dashed line in a). e)-h) Reconstructions from multi-slice
ptychography with MW-regularization applied, using the same dataset as in the
non-regularized case. i)-l) Reconstructions from multi-focus, multi-slice
ptychography using all four datasets and no MW-regularization. m)-p)
Reconstructions from multi-mode, multi-slice ptychography, using again just
the same single defocused dataset as in the first two cases. q)-t)
Reconstructions obtained from the optical sectioning $\mathcal{S}$-matrix
method, also using the entire 4D-STEM defocus series. Adapted with permission
from Ref. ([23]) © Cambridge University Press.

To further explore the enhancement in quality of ptychographic reconstructions
along the beam direction through the inclusion of multiple defocus 4D-STEM
datasets, we generated a depth profile at a diffraction spot characteristic of
both material compounds, yet absent in the interface layer. Figure 5a)
displays the depth profiles for various ptychography techniques, alongside
three diffractograms corresponding to the reconstructed PIO, interface, and
YSZ slices using the multi-focus technique. Notably, the intensity of the
diffraction spot peaks at slice 12, the center of the YSZ compound, across all
techniques. However, whereas the depth profiles for conventional, regularized,
and multi-mode ptychography show a leveling off towards the PIO compound, the
multi-focus technique exhibits a distinct peak, providing a clearer
delineation of the PIO compound. For this analysis, the reconstructions were
performed using 20 slices, each separated by 1 nm, to enhance the sampling
resolution along the optical axis.

Figure 5: Depth profile of the PIO-YSZ hetero-structure taken along a
diffraction spot in the diffractogram of each 2D slice. a) Depth profiles for
the four different multi-slice ptychography approaches investigated in this
paper (i.e. conventional, regularized, multi-mode and multi-focus
ptychography). For multi-focus ptychography, three diffractograms are shown
that correspond to b) one of the PIO slices, c) the interface slice and d) one
of the YSZ slices. The diffraction spot used to generate the depth profile is
encircled in red.

###  Reconstruction from simulated and experimental hBN

Figure 6 illustrates the reconstruction of simulated hBN using multi-slice
ptychography with a) a single 4D-STEM dataset at a defocus of $0$ nm and b) a
4D-STEM defocus series, along with c) the $\mathcal{S}$-matrix-based optical
sectioning that also uses the same defocus series. The 4D-STEM datasets are
generated from a scan area depicted in Fig. 3, and the three-dimensional
reconstructions are presented in Fig. 6. In all three methods, the two hBN
stacks are distinctly separated, and the Moiré pattern resulting from the
rotation of the two stacks is visible in the transition slices.

The reconstructions from single-focus and multi-focus ptychography methods
exhibit almost identical results, with a slight (3 nm) shift along the optical
axis in the former and centering in the latter case. This similarity arises
from the full coherence of the beam and the absence of TDS within the sample
during the generation of simulated data. With a correct scattering model,
there is only one configuration of the potential that can explain the
scattered electrons in the data. However, the shift in the reconstruction when
a single focus is used indicates that the optimization problem is not
adequately constrained. Thus, the voxels corresponding to vacuum can be
arbitrarily distributed before or after the reconstructed potential along the
optical axis. The fact that the reconstructed specimen is centered along the
optical axis in cases where a 4D-STEM defocus series is used supports this
assumption.

Comparing the multi-slice ptychography results to the outcome obtained from
the $\mathcal{S}$-matrix optical sectioning reveals that the reconstructed
atoms are less sharp in the latter case, indicating lower resolution. This
effect could stem from the model’s inability to handle multiple scattering, a
limitation observed in the earlier optical sectioning $\mathcal{S}$-matrix
reconstruction of the PIO-YSZ specimen. An alternative interpretation could be
based on the fact that ptychography more effectively deconvolves the probe
spread function.

Figure 6: Ptychographic 3D reconstructions from simulated data of a stacked
multi-layer hBN structure. A cut-through of a) a multi-slice ptychography
reconstruction using a single 4D-STEM dataset, b) a multi-slice ptychography
reconstruction using a 4D-STEM defocus series and c) the $\mathcal{S}$-matrix
optical sectioning reconstruction using the same 4D-STEM defocus series as in
b). For each method, three slices are additionally shown that contain only the
structure of the first hBN stack, only the structure of the second hBN stack
and only the interface of the two stacks.

The comparison of the three different 3D reconstruction methods is also
conducted on the same hBN specimen using experimentally acquired data.

Figure 7: Ptychographic 3D reconstructions from experimental data of a stacked
multi-layer hBN structure. A cut-through of the reconstruction generated with
a) the conventional ptychographic multi-slice method, b) the ptychographic
multi-focus and multi-slice method and c) the $\mathcal{S}$-matrix optical
sectioning method. For each method, three slices are additionally shown that
best match to the expected structure of the first hBN stack, the expected
structure of the second hBN stack and the expected interface of the two
stacks.

Figure 7 presents the reconstruction results of the methods. This time, a
substantial difference exists between the reconstructions generated by the
single-focus and the multi-focus method along the optical axis. In the single-
focus reconstruction, the potential fluctuates strongly along the beam
direction, whereas this is not the case for the multi-focus method. The
increased number of degrees of freedom in the reconstruction algorithm due to
the inclusion of multiple slices, combined with realistic experimental
conditions, results in incorrect solutions. It becomes evident why MW-
regularization is a suitable regularization precisely in this case - high
fluctuations along the optical axis are penalized, forcing the reconstruction
to closely resemble the one obtained from multi-focus ptychography, where the
additional 4D-STEM datasets acquired at different focal planes act as
constraints. Similar to Fig. 6, we observe that the Moiré pattern is not
resolved in the central slice of the single-focus reconstruction but rather
the $z$-position at which the Moiré occurs is slightly shifted along the optic
axis. Moreover, one of the two independent layers of the hBN stack is clearly
resolved. In contrast, the structures of each layer in the twisted hBN stack
are resolved in the multi-focus reconstruction, and the Moiré pattern is
visible in the central slice, where the transition between the two stacks is
expected. In the $\mathcal{S}$-matrix optical sectioning result, the
transition of the hBN stacks is also clearly visible in the center of the
reconstruction. However, the structure of the individual hBN stacks is not
resolved since the Moiré pattern dominates the reconstruction along the entire
z-direction. This is consistent with a lower axial resolution in the
reconstruction which contributes towards an inability to clearly resolve each
layer in the twisted structure.

We conducted another analysis on the experimental hBN sample, comparing the
intensity ratio of diffraction spots along the beam direction between the two
twisted hBN stacks. Figure 8a) presents the averaged intensity of six
diffraction spots across three techniques: conventional ptychography using a
single focused-probe 4D-STEM dataset, multi-focus ptychography, and
$\mathcal{S}$-matrix optical sectioning, both utilizing the 4D-STEM defocus
series. Figure 8b)-d) depict diffractograms of the reconstructed slices from
multi-focus ptychography, specifically at the onset of the first hBN stack,
the interface, and the onset of the second hBN stack. These images illustrate
the contributions from each stack. The depth profile from conventional
ptychography reveals a relatively uneven intensity trend across the
diffraction spot pairs. As previously noted in Figure 7, the interface layer,
characterized by its distinct Moiré pattern, is not perfectly centered within
the reconstructed volume. In contrast, the depth profile from the multi-focus
ptychography demonstrates a smoother transition between the two stacks, with a
significantly higher ratio at the beginning and end of the reconstruction,
indicating a more accurately reconstructed separation. Meanwhile, the depth
profile from the $\mathcal{S}$-matrix optical sectioning reconstruction shows
a highly linear trend. However, the ratio towards the surfaces of the
reconstructed sample remains relatively low, making it less effective in
pinpointing the exact locations of the hBN stacks within the volume.

Figure 8: Depth profiles of the two twisted hBN stacks in the experimental
sample, calculated from the average of six diffraction spots in each
reconstructed 2D slice’s diffractogram. a) shows intensity profiles of the two
equally sized hBN stacks using focused-probe ptychography, multi-focus
ptychography, and $\mathcal{S}$-matrix optical sectioning. For multi-focus
ptychography, b), c), and d) depict diffractograms for specific slices: b) the
initial slice of the first hBN stack, c) the interface slice with its Moiré
pattern, and d) the final slice of the second hBN stack. A pair of diffraction
spots, which contribute to the average used to generate these profiles, is
highlighted in red and blue. These colors correspond to the first and second
hBN stacks, respectively. The remaining five diffraction spot pairs that form
the average are marked with white arrows in c).

The 3D reconstruction results obtained with the multi-focus and multi-slice
ptychography method and the $\mathcal{S}$-matrix optical sectioning method
from the experimental data could potentially be enhanced through an
alternative postprocessing of the 4D-STEM defocus series. The currently used
HAADF-STEM-based approach for dataset alignment is sub-optimal because HAADF
images are low in contrast when the specimen is thin and primarily composed of
atoms with a low atomic number, especially when the focal plane does not lie
within the specimen. However, restricting the defocus range to only fall
within the specimen thickness would make it more challenging for the
reconstruction algorithm to correctly recover the surface boundaries.
Therefore, we suggest that an alternative alignment of the defocus series
through a multi-slice ptychography reconstruction of each individual dataset
may potentially help address the limitations of the current alignment
procedure.

## Conclusion

An extension of multi-slice ptychography has been proposed, leveraging a
4D-STEM defocus series from a single sample orientation to improve the three-
dimensional phase reconstruction. The method was validated through simulations
and experiments involving various samples, each with a thickness ranging up to
tens of nanometers and atomic numbers ranging from 5 (B) to 82 (Pb). More
specifically, the method was applied to experimental data involving a $20$ nm
thick PIO-YSZ hetero-structure and simulated and experimental data for a $22$
nm and $17$ nm thick multi-layered hBN sample, respectively. Incorporating of
a defocus series into the ptychographic reconstruction algorithm mitigates the
impact of partial incoherence on the reconstruction outcome, rendering the use
of regularization methods unnecessary for achieving a correct solution. This
innovation has demonstrated higher fidelity of reconstructions along the beam
direction as evidenced by the smoothness and overall value of reconstructed
potential in homogeneous areas. The capacity to perform three-dimensional
phase reconstructions using a 4D-STEM defocus series is also achievable with
an $\mathcal{S}$-matrix-based optical sectioning approach. We conducted a
comparative analysis of the two approaches, revealing that the implicit
treatment of multiple scattering through averaging over different momentum
components, as employed in the $\mathcal{S}$-matrix approach, is inadequate.
Instead, an explicit inclusion of multiple scattering in the model,
characteristic of multi-slice ptychography, proves essential for a meaningful
reconstruction.

##  Appendix A Analysis of multi-mode ptychography

In this study, we investigate the impact of partial spatial coherence in the
illuminating probe on ptychographic reconstructions when using just a single
4D-STEM dataset and explore the extent to which a mixed-state model can
counteract it. To achieve this, we conducted a comparative analysis between a
single-mode and a multi-mode ptychographic reconstruction using an
experimentally acquired dataset of the stacked multi-layered hBN sample,
previously detailed in this paper. In this case, the experiment was conducted
with a Thermo Fisher Scientific Spectra $\phi$ FEGTEM microscope fitted with a
CEOS SCORR+ probe spherical aberration corrector, operating at $300$ kV
(wavelength of $1.97$ pm) acceleration voltage, a semi-convergence angle
$\theta_{\text{con}}$ of $25$ mrad, and an EMPAD direct electron detector. The
selected sample region exhibited a Moiré-like pattern arising from a slight
rotation of the two hBN stacks, covering the entire field of view (FOV) of
$3.33$ nm $\times$ $3.33$ nm, corresponding to $256\times 256$ scan positions.
Data acquisition involved a dwell time of $1.0$ ms and a beam current of $6.1$
pA. In the multi-mode reconstruction, we employed four different Hermite-
Gaussian probe modes, orthogonalized using the Gram-Schmidt orthogonalization
method. A multi-slice model with five slices and a slice distance of $3.6$ nm
was utilized to accommodate the estimated sample thickness of approximately
$17$ nm. The reconstruction algorithm optimized the object and probe modes
over 400 iterations. Figure 9 displays the reconstructed slice of the
interface between the two hBN stacks and the illuminating probe for both the
single-state and mixed-state models.

Figure 9: Comparison between the phase reconstructions of a stacked multilayer
hBN sample obtained from single-mode and multi-mode ptychography. The central
slice of the multi-slice reconstructions and their respective diffractograms
are shown. The focused probes used in both reconstructions are illustrated and
the four probe modes that form the multi-mode probe with their respective
contribution are given.

The retrieved probes are depicted for both scenarios, alongside the four
individual probe modes and their respective contributions to the multi-mode
probe, expressed as a percentage. A comparison between single-mode
ptychography and multi-mode ptychography reveals an enhancement in
reconstruction quality with the latter. Specifically, the hBN structure,
particularly the unit cells, is more clearly resolved. However, when examining
their diffractograms, the resolutions of the two reconstructions appear
generally similar. Consequently, while the present findings show some
advantage of multi-mode ptychography over convetional single-mode
ptychography, the substantial difference in reconstructions reported in Ref.
([5]) and ([4]) was not replicated in this study.

To further investigate, we conducted an additional analysis using the bilayer
MoSe2/WS2 data from Ref. ([5]). Once again, we performed both single-mode and
multi-mode ptychographic reconstructions, employing four orthogonalized probe
modes for the latter. Figure 10 displays the single-slice reconstruction of
the sample alongside the single-mode probe used, and the single-slice
reconstruction with the multi-mode probe and its probe modes. In this
comparison, the difference between the reconstruction results is quite
pronounced, consistent with the reported results. Structural features, such as
a monolayer of WS2 and well-aligned and misaligned stacking in bilayer
MoSe2/WS2 regions, are much more clearly resolved with multi-mode
ptychography. Notably, the experiment involved a strongly defocused probe of
approximately $55$ nm and a large scan step size of $2.36$ Å , in contrast to
the focused-probe experiment in the previous analysis. This implies that
partial spatial incoherence has a more pronounced impact on reconstruction
quality when using a defocused probe. This finding is in agreement with what
has been reported in Ref. ([5]).

Figure 10: Improvement of the phase reconstruction of a bilayer MoSe2/WS2
sample through multi-mode ptychography. The reconstructions are performed with
single-mode and multi-mode ptychography and the corresponding defocused probes
of the two reconstructions are illustrated. The four probe modes with their
respective contribution to the multi-mode probe are shown.

## Acknowledgments

M.S., T.C.P., and C.T.K acknowledge support from the Deutsche
Forschungsgemeinschaft (DFG, German Research Foundation) - Project-ID
414984028 - SFB 1404. Work at the Molecular Foundry was supported by the
Office of Science, Office of Basic Energy Sciences, of the U.S. Department of
Energy under Contract No. DE-AC02-05CH11231. D.O.B. also acknowledges funding
from the Department of Defense through the National Defense Science &
Engineering Graduate (NDSEG) Fellowship Program. J.C. acknowledges additional
support from the Presidential Early Career Award for Scientists and Engineers
(PECASE) through the U.S. Department of Energy. This research was supported
under the Discovery Projects funding scheme of the Australian Research Council
(project no. FT190100619). This research was supported by an Australian
Government Research Training Program Scholarship.

## References

## References

  * Bosch and Lazić [2019] Bosch, E. G.and Lazić, I., “Analysis of depth-sectioning stem for thick samples and 3d imaging,” Ultramicroscopy 207, 112831 (2019).
  * Brown  _et al._ [2022] Brown, H. G., Pelz, P. M., Hsu, S. L., Zhang, Z., Ramesh, R., Inzani, K., Sheridan, E., Griffin, S. M., Schloz, M., Pekin, T. C., _et al._ , “A three-dimensional reconstruction algorithm for scanning transmission electron microscopy data from a single sample orientation,” Microscopy and Microanalysis 28, 1632–1640 (2022).
  * Chang  _et al._ [2020] Chang, D. J., Kim, D. S., Rana, A., Tian, X., Zhou, J., Ercius, P., and Miao, J., “Ptychographic atomic electron tomography: Towards three-dimensional imaging of individual light atoms in materials,” Physical Review B 102, 174101 (2020).
  * Chen  _et al._ [2021] Chen, Z., Jiang, Y., Shao, Y. T., Holtz, M. E., Odstrčil, M., Guizar-Sicairos, M., Hanke, I., Ganschow, S., Schlom, D. G., and Muller, D. A., “Electron ptychography achieves atomic-resolution limits set by lattice vibrations,” Science 372, 826–831 (2021).
  * Chen  _et al._ [2020] Chen, Z., Odstrcil, M., Jiang, Y., Han, Y., Chiu, M. H., Li, L. J., and Muller, D. A., “Mixed-state electron ptychography enables sub-angstrom resolution imaging with picometer precision at low dose,” Nature Communications 11, 1–10 (2020).
  * Chen  _et al._ [2016] Chen, Z., Weyland, M., Ercius, P., Ciston, J., Zheng, C., Fuhrer, M., D’Alfonso, A., Allen, L., and Findlay, S., “Practical aspects of diffractive imaging using an atomic-scale coherent electron probe,” Ultramicroscopy 169, 107–121 (2016).
  * Ding  _et al._ [2022] Ding, Z., Gao, S., Fang, W., Huang, C., Zhou, L., Pei, X., Liu, X., Pan, X., Fan, C., Kirkland, A. I., _et al._ , “Three-dimensional electron ptychography of organic–inorganic hybrid nanostructures,” Nature Communications 13, 4787 (2022).
  * Fang  _et al._ [2019] Fang, S., Wen, Y., Allen, C. S., Ophus, C., Han, G. G., Kirkland, A. I., Kaxiras, E., and Warner, J. H., “Atomic electrostatic maps of 1d channels in 2d semiconductors using 4d scanning transmission electron microscopy,” Nature communications 10, 1127 (2019).
  * Findlay  _et al._ [2021] Findlay, S. D., Brown, H. G., Pelz, P. M., Ophus, C., Ciston, J., and Allen, L. J., “Scattering matrix determination in crystalline materials from 4d scanning transmission electron microscopy at a single defocus value,” Microscopy and microanalysis 27, 744–757 (2021).
  * Jacobsen [2018] Jacobsen, C., “Relaxation of the crowther criterion in multislice tomography,” Optics letters 43, 4811–4814 (2018).
  * Jiang  _et al._ [2018] Jiang, Y., Chen, Z., Han, Y., Deb, P., Gao, H., Xie, S., Purohit, P., Tate, M. W., Park, J., Gruner, S. M., _et al._ , “Electron ptychography of 2d materials to deep sub-ångström resolution,” Nature 559, 343–349 (2018).
  * Jones  _et al._ [2015] Jones, L., Yang, H., Pennycook, T. J., Marshall, M. S. J., Van Aert, S., Browning, N. D., Castell, M. R., and Nellist, P. D., “Smart align—a new tool for robust non-rigid registration of scanning microscope data,” Advanced Structural and Chemical Imaging 1, 1–16 (2015).
  * Kotakoski  _et al._ [2010] Kotakoski, J., Jin, C. H., Lehtinen, O., Suenaga, K., and Krasheninnikov, A. V., “Electron knock-on damage in hexagonal boron nitride monolayers,” Physical Review B 82, 113404 (2010).
  * Lozano  _et al._ [2018] Lozano, J. G., Martinez, G. T., Jin, L., Nellist, P. D., and Bruce, P. G., “Low-dose aberration-free imaging of li-rich cathode materials at various states of charge using electron ptychography,” Nano letters 18, 6850–6855 (2018).
  * Madsen and Susi [2021] Madsen, J.and Susi, T., “The abtem code: transmission electron microscopy from first principles,” Open Research Europe 1, 24 (2021).
  * Midgley and Dunin-Borkowski [2009] Midgley, P. A.and Dunin-Borkowski, R. E., “Electron tomography and holography in materials science,” Nature materials 8, 271–280 (2009).
  * Momma and Izumi [2011] Momma, K.and Izumi, F., “Vesta 3 for three-dimensional visualization of crystal, volumetric and morphology data,” Journal of applied crystallography 44, 1272–1276 (2011).
  * O’Leary  _et al._ [2022] O’Leary, C. M., Haas, B., Koch, C. T., Nellist, P. D., and Jones, L., “Increasing spatial fidelity and snr of 4d-stem using multi-frame data fusion,” Microscopy and Microanalysis 28, 1417–1427 (2022).
  * O’Leary  _et al._ [2023] O’Leary, C. M., Zhang, J., Su, C., Kahn, S., Jiang, H., Zettl, A., Ciston, J., and Miao, J., “Three-dimensional imaging of embedded interfaces with multislice ptychography,” arXiv preprint arXiv:2308.15471 (2023).
  * Ophus  _et al._ [2019] Ophus, C., Harvey, T. R., Yasin, F. S., Brown, H. G., Pelz, P. M., Savitzky, B. H., Ciston, J., and McMorran, B. J., “Advanced phase reconstruction methods enabled by four-dimensional scanning transmission electron microscopy,” Microscopy and Microanalysis 25, 10–11 (2019).
  * Pelz  _et al._ [2023] Pelz, P. M., Griffin, S. M., Stonemeyer, S., Popple, D., DeVyldere, H., Ercius, P., Zettl, A., Scott, M. C., and Ophus, C., “Solving complex nanostructures with ptychographic atomic electron tomography,” Nature Communications 14, 7906 (2023).
  * Raines  _et al._ [2010] Raines, K. S., Salha, S., Sandberg, R. L., Jiang, H., Rodríguez, J. A., Fahimian, B. P., Kapteyn, H. C., Du, J., and Miao, J., “Three-dimensional structure determination from a single view,” Nature 463, 214–217 (2010).
  * Schloz  _et al._ [2022] Schloz, M., Pekin, T. C., Brown, H. G., Ciston, J., and Koch, C. T., “High resolution three-dimensional reconstructions in electron microscopy through multifocus ptychography,” Microscopy and Microanalysis 28, 364–366 (2022).
  * Schloz  _et al._ [2020] Schloz, M., Pekin, T. C., Chen, Z., Van den Broek, W., Muller, D. A., and Koch, C. T., “Overcoming information reduced data and experimentally uncertain parameters in ptychography with regularized optimization,” Optics Express 28, 28306–28323 (2020).
  * Terzoudis-Lumsden  _et al._ [2023] Terzoudis-Lumsden, E. W. C., Petersen, T. C., Brown, H. G., Pelz, P. M., Ophus, C., and Findlay, S. D., “Resolution of virtual depth sectioning from four-dimensional scanning transmission electron microscopy,” Microscopy and Microanalysis 29, 1409–1421 (2023).
  * Thibault and Guizar-Sicairos [2012] Thibault, P.and Guizar-Sicairos, M., “Maximum-likelihood refinement for coherent diffractive imaging,” New Journal of Physics 14, 063004 (2012).
  * Tsai  _et al._ [2016] Tsai, E. H. R., Usov, I., Diaz, A., Menzel, A., and Guizar-Sicairos, M., “X-ray ptychography with extended depth of field,” Optics Express 24, 29089–29108 (2016).
  * Varnavides  _et al._ [2023] Varnavides, G., Ribet, S. M., Zeltmann, S. E., Yu, Y., Savitzky, B. H., Dravid, V. P., Scott, M. C., and Ophus, C., “Iterative phase retrieval algorithms for scanning transmission electron microscopy,” arXiv preprint arXiv:2309.05250 (2023).
  * Voyles, Grazul, and Muller [2003] Voyles, P., Grazul, J., and Muller, D., “Imaging individual atoms inside crystals with adf-stem,” Ultramicroscopy 96, 251–273 (2003).
  * Xin and Muller [2009] Xin, H. L.and Muller, D. A., “Aberration-corrected adf-stem depth sectioning and prospects for reliable 3d imaging in s/tem,” Journal of Electron Microscopy 58, 157–165 (2009).
  * Zhang  _et al._ [2023] Zhang, H., Li, G., Zhang, J., Zhang, D., Chen, Z., Liu, X., Guo, P., Zhu, Y., Chen, C., Liu, L., _et al._ , “Three-dimensional inhomogeneity of zeolite structure and composition revealed by electron ptychography,” Science 380, 633–638 (2023).


# Locally Regularized Neural Differential Equations:
Some Black Boxes Were Meant to Remain Closed!

Avik Pal  Alan Edelman  Christopher Rackauckas

###### Abstract

Implicit layer deep learning techniques, like Neural Differential Equations,
have become an important modeling framework due to their ability to adapt to
new problems automatically. Training a neural differential equation is
effectively a search over a space of plausible dynamical systems. However,
controlling the computational cost for these models is difficult since it
relies on the number of steps the adaptive solver takes. Most prior works have
used higher-order methods to reduce prediction timings while greatly
increasing training time or reducing both training and prediction timings by
relying on specific training algorithms, which are harder to use as a drop-in
replacement due to strict requirements on automatic differentiation. In this
manuscript, we use internal cost heuristics of adaptive differential equation
solvers at stochastic time-points to guide the training towards learning a
dynamical system that is easier to integrate. We “close the blackbox” and
allow the use of our method with any adjoint technique for gradient
calculations of the differential equation solution. We perform experimental
studies to compare our method to global regularization to show that we attain
similar performance numbers without compromising on the flexibility of
implementation on ordinary differential equations (ODEs) and stochastic
differential equations (SDEs). We develop two sampling strategies to trade-off
between performance and training time. Our method reduces the number of
function evaluations to $\mathit{0.556\times-0.733\times}$ and accelerates
predictions by $\mathit{1.3\times-2\times}$.

Machine Learning, Implicit Neural Networks, Neural ODEs, ICML

Sensitivity Algorithm Memory Requirement Memory Requirement with Local
Regularization Backsolve Adjoint (Chen et al., 2018)
$\mathcal{O}\left(s\right)$
$\mathcal{O}\left(s\times(1+\mathrm{stages})\right)$ Backsolve Adjoint with
Checkpointing (Chen et al., 2018) $\mathcal{O}\left(s\times c\right)$
$\mathcal{O}\left(s\times(c+\mathrm{stages})\right)$ Interpolating Adjoint
(Hindmarsh et al., 2005) $\mathcal{O}\left(s\times t\right)$
$\mathcal{O}\left(s\times(t+\mathrm{stages})\right)$ Interpolating Adjoint
with Checkpointing (Hindmarsh et al., 2005) $\mathcal{O}\left(s\times
c\right)$ $\mathcal{O}\left(s\times(c+\times\mathrm{stages})\right)$
Quadrature Adjoint (Kim et al., 2021) $\mathcal{O}\left(\left(s+p\right)\times
t\right)$ $\mathcal{O}\left(\left(s+p\right)\times
t+s\times\mathrm{stages}\right)$ Direct Reverse Mode Differentiation
$\mathcal{O}\left(s\times t\times\mathrm{stages}\right)$
$\mathcal{O}\left(s\times\left(t+1\right)\times\mathrm{stages}\right)$

Table 1: Memory Requirements for various Sensitivity Algorithms for ODEs with
Local Regularization: Local Regularization requires
$\mathcal{O}\left(s\times(c+\mathrm{stages})\right)$ memory in the best case,
while Global Regularization (Pal et al., 2021) always requires
$\mathcal{O}\left(s\times t\times\mathrm{stages}\right)$ memory. Figure 1:
Locally Regularized NDE leads to faster predictions and faster training
compared to vanilla NDE.

##  1 Introduction

Implicit Models, such as Neural Ordinary Differential Equations (Chen et al.,
2018) and Deep Equilibrium Models (Bai et al., 2019; Pal et al., 2022), have
emerged as a promising technique to determine the depth of neural networks
automatically. To maximize performance on a dataset, explicit models are tuned
to the “hardest” training sample, which hurts the inference timings for
“easier” – more abundant – samples. Using adaptive differential equation
solvers allow these implicit models to choose the number of steps they need
effectively. This idea of representing neural networks as ODEs has since been
generalized to Stochastic Differential Equations (Liu et al., 2019; Rackauckas
et al., 2020) and other architectures to improve robustness.

Despite the rapid progress in these methods, the core problem of the
scalability of these models is still persistent. Several solutions to these
have been proposed:

  * •

Kelly et al. (2020); Finlay et al. (2020) use higher order derivatives for
regularization.

  * •

Poli et al. (2020) learn neural solvers to solve Neural ODEs faster.

  * •

Pal et al. (2021) proposed a “zero-cost” global regularization scheme.

  * •

Behl et al. (2020) randomize the integration stop time to “smoothen” the
dynamics.

All these methods have definite tradeoffs (See Section 5 for more details).

This paper presents a generally applicable method to force the neural
differential equation training process to choose the least expensive option.
We build upon the global regularization scheme proposed in Pal et al. (2021)
and “close” the blackbox allowing our method to work across various
sensitivity algorithms. Our main contributions include the following111Our
code is publicly available at https://github.com/avik-pal/LocalRegNeuralDE.jl:

  1. 1.

We show that our local regularization method – building upon the primitives
proposed in Pal et al. (2021) – performs at par with global regularization
(See Section 4).

  2. 2.

We present two sampling methods that trade-off small computational costs for
consistently better performance (See Section 3).

  3. 3.

Using local regularization allows our models to leverage optimize-then-
discretize in the backward pass (in addition to discretize-then-optimize
methods). Our method works around the several engineering limitations of
automatic differentiation (AD) systems (Rackauckas, 2022) that are needed to
make global regularization work.

  4. 4.

We empirically show that regularizing solver heuristics with biased sampling
stabilizes the training of larger neural ODEs (See Section 4.3).

##  2 Background

###  2.1 Neural Ordinary Differential Equations

Neural ODEs use an explicit neural network to parameterize the dynamical
system. These involve finding the state at a later time $t_{1}$, given the
value $z_{0}$ at time $t_{0}$. The final state:

$z(t_{1})=z_{0}+\int_{t_{0}}^{t_{1}}f_{\theta}\left(z(t),t\right)dt$ (1)

generally cannot be computed analytically and requires numerical solvers.
Prior works had observed the similarity between fixed time-step discretization
of ODEs and Residual Neural Networks Lu et al. (2018), which was later
extended in the Neural ODE framework by Chen et al. (2018).

$\frac{dz(t)}{dt}=f_{\theta}\left(z(t),t\right)$ (2)

Using adaptive time stepping allows the model to operate at a variable
continuous depth depending on the inputs. Removal of the fixed depth
constraint of Residual Networks provides a more expressive framework and
offers several advantages in problems like density estimation (Grathwohl et
al., 2018), irregularly spaced time series problems (Rubanova et al., 2019),
etc.

###  2.2 Neural Stochastic Differential Equations

Stochastic Differential Equations (SDEs) couple the effect of noise to a
deterministic system of equations. SDE noise can have different variations;
however, we exclusively focus on diagonal multiplicative noise for this paper.
Liu et al. (2019) modified Neural ODEs by stochastic noise injection in the
form of Neural SDEs. Liu et al. (2019) empirically showed that stochastic
noise injection improves the robustness and generalization performance of
neural ODEs. In Neural SDEs, we simultaneously train two explicit neural
networks – drift $f_{\theta}(z(t),t)$ and diffusion $g_{\phi}(z(t),t))$ s.t.:

$dz(t)=f_{\theta}(z(t),t)dt+g_{\phi}(z(t),t))dW$ (3)

###  2.3 Adaptive Time Stepping for Differential Equations

Runge-Kutta (RK) Methods (Runge, 1895; Kutta, 1901) are widely used to
approximate the solutions of ordinary differential equations numerically.
Given a tableau of coefficients $\left\\{A,c,b\right\\}$, these methods
combine $s$ stages to obtain the estimate at $t+dt$.

$\displaystyle k_{s}=f\left(t+c_{s}\cdot dt,z(t)+\sum_{i=1}^{s-1}a_{si}\cdot
k_{i}\right)$ (4) $\displaystyle
z(t+dt)=z(t)+dt\cdot\left(\sum_{i=1}^{s}b_{i}\cdot k_{i}\right)$ (5)

Adaptive solvers need to maximize the step size ($dt$) while keeping the error
estimate below the user-specified tolerances, i.e., they need to satisfy:

$\texttt{E}_{\texttt{Est}}\leq\texttt{atol}+\texttt{max}\left(|z(t)|,|z(t+dt)|\right)\cdot\texttt{rtol}$
(6)

where $\texttt{E}_{\texttt{Est}}$ is the local error estimate. Adaptive
solvers utilize an additional linear combiner $\tilde{b}_{i}$ to get an
alternate solution, typically with one order less convergence (Wanner &
Hairer, 1996; Fehlberg, 1968; Dormand & Prince, 1980; Tsitouras, 2011).

$\tilde{z}(t+dt)=z(t)+dt\cdot\left(\sum_{i=1}^{s}\tilde{b}_{i}\cdot
k_{i}\right)$ (7)

A classic result from Richardson extrapolation shows that
$\texttt{E}_{\texttt{Est}}=\|\tilde{z}(t+dt)-z(t+dt)\|$ is an estimate of the
local truncation error (Hairer et al., 1993; Ascher & Petzold, 1998). The new
step size is determined using the following:

$q=\left\|\frac{\texttt{E}_{\texttt{Est}}}{\texttt{atol}+\texttt{max}\left(|z(t)|,|z(t+dt)|\right)\cdot\texttt{rtol}}\right\|$
(8)

  * •

If $q<1$, $dt$ is accepted.

  * •

Otherwise, it is rejected and reduced. A standard PI controller proposes the
new step size to be:

$dt_{new}=\eta\cdot q_{n-1}^{\alpha}\cdot q_{n}^{\beta}\cdot dt$ (9)

where $\eta$ is the safety factor, $q_{n-1}$ is the error proportion of the
previous step, and $(\alpha,\beta)$ are tunable PI-gain hyperparameters
(Wanner & Hairer, 1996).

We defer the discussion of error estimation schemes for stochastic RK
integrators to Rackauckas & Nie (2017, 2020).

Data: x, $f_{\theta}$, $t_{span}$

Result: sol, r

Define $\frac{du}{dt}=f_{\theta}(u,t)$

$t_{0},t_{1}\leftarrow t_{span}$

$\texttt{t}_{\texttt{reg}}\sim\mathbb{U}\left[t_{0},t_{1}\right]$

$\texttt{sol}\leftarrow\texttt{solve}(\frac{du}{dt},\texttt{ DE
Solver},~{}t_{span})$

$u_{\texttt{t}_{\texttt{reg}}}\leftarrow\texttt{sol}(\texttt{t}_{\texttt{reg}})$

Run single step for the solver with time-span
$(\texttt{t}_{\texttt{reg}},t_{1})$

$\texttt{r}\leftarrow$ Local Error Estimate $@t=\texttt{t}_{\texttt{reg}}$

Algorithm 1 Unbiased Sampling: Training Phase

Data: x, $f_{\theta}$, $t_{span}$

Result: sol, r

Define $\frac{du}{dt}=f_{\theta}(u,t)$

$t_{0},t_{1}\leftarrow t_{span}$

$\texttt{sol}\leftarrow\texttt{solve}(\frac{du}{dt},\texttt{ DE
Solver},~{}t_{span})$

$\texttt{t}_{\texttt{reg}}\sim\mathbb{U}\left(\texttt{sol}.t\right)$

$u_{\texttt{t}_{\texttt{reg}}}\leftarrow\texttt{sol}(\texttt{t}_{\texttt{reg}})$

Run single step for the solver with time-span
$(\texttt{t}_{\texttt{reg}},t_{1})$

$\texttt{r}\leftarrow$ Local Error Estimate $@t=\texttt{t}_{\texttt{reg}}$

Algorithm 2 Biased Sampling: Training Phase

###  2.4 Global Regularization using Local Error Estimates

In Section 2.3, we observe how larger local error estimates
$\texttt{E}_{\texttt{Est}}$ contribute to reduced step sizes and, in turn,
higher training and prediction times for Neural Differential Equations. Pal et
al. (2021) proposed a regularization scheme to minimize the “total local error
in order to learn Neural ODEs with as large step sizes as possible.” They
compute the regularization term $\left(\mathcal{R}_{E}\right)_{g}$ as:

$\left(\mathcal{R}_{E}\right)_{g}=\sum_{j}\left(\texttt{E}_{\texttt{Est}}\right)_{j}\cdot|dt_{j}|$
(10)

where the summation is over the $j$ time steps of the solution. All the local
error estimates $\texttt{E}_{\texttt{Est}}$ are precomputed by adaptive
differential equation solvers, which makes the forward pass effectively have
zero additional overhead.

Continuous adjoint methods (Chen et al., 2018) define the derivatives in terms
of the ODE quantities. The terms
$\left(\texttt{E}_{\texttt{Est}}\right)_{j}=\sum_{i=1}^{s}(b_{i}-\tilde{b}_{i})\cdot
k_{i}$ cannot be constructed directly from the trajectory of the ODE solution,
since $k_{i}$ terms are defined by the solver method (and not by the
continuous ODE solution). As a result Pal et al. (2021) performs discrete
sensitivity analysis to compute the derivatives w.r.t. the regularization
terms.

Discrete Sensitivity analysis is more stable than continuous adjoints (Zhang &
Sandu, 2014) and stabilizes the training process of Neural ODEs (Onken &
Ruthotto, 2020). However, these methods are more memory intensive than
continuous adjoints. Another shortcoming is the advanced tooling mechanism
needed to differentiate through differential equation solvers (Rackauckas,
2022), making this method hard to adopt. Since adaptive ODE solvers change the
number of steps based on the values in the differential equation, discrete
sensitivity analysis requires dynamic AD tooling which is generally less
optimized than those built for static graph cases. We note that there are
checkpointing methods to reduce the memory overhead (Dauvergne & Hascoët,
2006), though these generally have focused on static graph AD systems.

###  2.5 Multiscale Neural ODE with Input Injection

Multiscale modeling (Burt & Adelson, 1987) has been the central theme for
several deep computer vision applications (Farabet et al., 2012; Yu & Koltun,
2015; Chen et al., 2016, 2017). Standard Neural ODEs set the initial condition
as input $x$ from the previous layer. Multiscale Neural ODE with Input
Injection is based on the Multiscale Deep Equilibrium Models (DEQs) (Bai et
al., 2019, 2020; Pal et al., 2022). If our model operates at $n$ scales, the
initial condition $\\{(z_{0})_{i}\\}_{i\in[n-1]}$ is defined as
$(z_{0})_{0}=x$ and $(z_{0})_{i\in\\{1,\dots n-1\\}}=0$. Every feature set is
upsampled or downsampled and fed into the other scales using a neural network.
In line with Multiscale DEQs, we also repeatedly inject the original input $x$
at every step of the dynamical system.

##  3 Methods

In Section 2.4, we discussed the downsides of using global regularization with
local error estimates. To summarize:

  1. 1.

Global Regularization relies on discrete sensitivity analysis, which is more
memory intensive.

  2. 2.

Global Regularization depends on AD tooling to support dynamic compute graphs
in an efficient way, making it hard to incorporate into existing codebases.

To get around these limitations, we developed a new technique using local
sampling of error estimates at specific time points, rather than globally over
the full interval. We deal with sampling the “appropriate” time point for
regularization by two strategies:

  * •

Algorithm 1 Unbiased Sampling: We random uniformly sample the time-point in
the integration time span. Intuitively, since we will perform the training for
“a large number of steps,” the learned dynamical system would end up being
faster to solve “everywhere” over the time span.

  * •

Algorithm 2 Biased Sampling: Adaptive Time-Stepping Differential Equation
Solvers naturally take more steps around the area, which is harder to
integrate. We can bias the regularization to operate around parts of the
dynamical system which are “harder” by sampling a time-point from the solution
time points.

###  3.1 Unbiased Sampling

When training a Neural ODE, the integration timespan is fixed. Training any
deep learning model involves several thousand steps. We compute the total
local error estimate over the entire timespan when performing global
regularization. For unbiased sampling, we hypothesize that if we regularize at
random uniformly sampled time points in the timespan, the learned dynamical
system will demonstrate similar properties in terms of NFE compared to global
regularization. Our new regularization term becomes
$\left(\mathcal{R}_{E}\right)_{\texttt{unbiased}}$ as:

$\left(\mathcal{R}_{E}\right)_{\texttt{unbiased}}=\left(\texttt{E}_{\texttt{Est}}\right)_{\texttt{t}_{\texttt{reg}}}\cdot|dt_{\texttt{t}_{\texttt{reg}}}|\quad\texttt{t}_{\texttt{reg}}\sim\mathbb{U}[\texttt{tspan}]$
(11)

###  3.2 Biased Sampling

Figure 2: Robertson Stiff ODE System: Solving stiff systems like Robertson
(Robertson, 1966) (using Rodas5 (Piché, 1995)) involves spending around $75\%$
of the time in $t<5000$ (i.e. $5\%$ of the timespan). The vertical lines
denote the timepoints at which the ODE System was solved.

Consider a simple scenario where the learned dynamics of the DE is harder to
solve in $\left[0.25,0.35\right]$, and we are solving the DE from $t_{0}=0$ to
$t_{1}=1$. Our primary aim is to modify the learned system s.t. it becomes
simpler to solve in $\left[0.25,0.35\right]$. If we use unbiased sampling, the
probability that we regularize at
$\texttt{t}_{\texttt{reg}}\in\left[0.25,0.35\right]$ is $0.1$ (which is low).
The problem gets even more severe if the range is lowered. An extreme version
of this problem is observed for stiff systems like Robertson’s Equations (See
Figure 2) where $75\%$ of the time is spent in solving $5\%$ of the problem.
We note that these extreme scenarios rarely occur for traditional deep
learning tasks since Pal et al. (2021) observed minor speedups using stiffness
regularization. However, the problem that some parts of the dynamical system
are harder to integrate persists, and designing a regularization scheme
targeting those parts is highly desirable.

We considered a simple scenario where the learned dynamical system was fixed.
However, while training NDEs, this system evolves with training, and apriori
predicting the more difficult portions to integrate is not feasible. Adaptive
solvers take more frequent steps in the parts of the DE where it is harder to
integrate. Anantharaman et al. (2020) leveraged this property of adaptive
solvers to learn surrogates for stiff systems. Since these solvers adapt to
concentrate around the most numerically difficult time points, we
automatically obtain the time points where we want to regularize the model.
Hence, for our biased sampling regularize, we uniformly sample the
regularization timepoint $\texttt{t}_{\texttt{reg}}$ from the time points at
which the solver solved the differential equation.

Method Train Accuracy (%) Test Accuracy (%) Training Time (hr) Prediction Time
(s / batch) Testing NFE Vanilla NODE $99.898\pm 0.066$ $97.612\pm 0.163$
$0.54\pm 0.001$ $0.088\pm 0.020$ $303.559\pm 3.194$ STEER
${}^{\text{\textparagraph}}$ $100.00\pm 0.000$ $97.94\hphantom{0}\pm
0.03\hphantom{0}$ $1.31\pm 0.07\hphantom{0}$ $0.092\pm 0.002$
$265.0\hphantom{0}\hphantom{0}\pm 3.46\hphantom{0}$ TayNODE
${}^{\text{\textparagraph}}$ $\hphantom{0}98.98\pm 0.06$ $97.89\hphantom{0}\pm
0.00$ $1.19\pm 0.07$ $0.079\pm 0.007$
$\hphantom{0}80.3\hphantom{0}\hphantom{0}\pm 0.43$ ERNODE
${}^{\text{\textparagraph}}$ $\hphantom{0}99.71\pm 0.28$ $97.32\hphantom{0}\pm
0.06$ $0.82\pm 0.02$ $0.060\pm 0.001$ $177.0\hphantom{00}\pm 0.00$ SRNODE
${}^{\text{\textparagraph}}$ $100.00\pm 0.000$ $98.08\hphantom{0}\pm 0.15$
$1.24\pm 0.06$ $0.094\pm 0.003$ $259.0\hphantom{00}\pm 3.46$ Local Unbiased
ERNODE $99.447\pm 0.039$ $97.526\pm 0.131$ $0.49\pm 0.002$ $0.046\pm 0.002$
$187.961\pm 1.812$ Local Biased ERNODE $99.477\pm 0.051$ $97.488\pm 0.016$
$1.12\pm 0.065$ $0.044\pm 0.002$ $182.849\pm 1.578$ Vanilla NSDE $98.27\pm
0.11$ $96.66\pm 0.16$ $2.70\pm 0.00$ $\hphantom{0}0.51\pm 0.07$ $313.86\pm
2.94$ ERNSDE ${}^{\text{\textparagraph}}$ $98.16\pm 0.11$ $96.27\pm 0.35$
$4.19\pm 0.04$ $\hphantom{0}7.23\pm 0.14$ $184.67\pm 2.31$ SRNSDE
${}^{\text{\textparagraph}}$ $98.79\pm 0.12$ $96.80\pm 0.07$ $8.54\pm 0.37$
$14.50\pm 0.40$ $382.00\pm 4.00$ Local Unbiased ERNSDE $98.05\pm 0.09$
$96.57\pm 0.13$ $2.10\pm 0.01$ $\hphantom{0}0.39\pm 0.10$ $228.93\pm 1.77$
Local Biased ERNSDE $98.02\pm 0.07$ $96.44\pm 0.16$ $1.90\pm 0.00$
$\hphantom{0}0.36\pm 0.03$ $230.10\pm 0.71$

Table 2: MNIST Image Classification using Neural DE: Using local unbiased
regularization on neural ODE speeds up training by $\mathit{1.1\times}$ and
predictions by $\mathit{1.9\times}$ while reducing the total NFEs to
$\mathit{0.619\times}$. Local Biased Regularization tends to slow down
training for smaller models on GPU while it further reduces the NFEs by
$\mathit{0.602\times}$. For Neural SDE, we observe a similar reduction of NFEs
by $\mathit{0.729\times-0.733\times}$ and a training time improvement of
$\mathit{1.28\times-1.42\times}$. The best global regularization method gets
lower NFEs but overall takes more wall clock than the best performing local
regularization method.

##  4 Experiments

In this section, we compare the effectiveness of unbiased and biased local
regularization’s effectiveness on the training and prediction timings of NDEs.
We choose image classification and time series prediction problems in line
with prior works on accelerating NDEs. We consider the following baselines:

  1. 1.

Vanilla Neural ODE with Continuous Interpolating Adjoint.

  2. 2.

Vanilla Neural SDE with discrete sensitivities.

  3. 3.

Global Regularization of Neural Differential Equations using discrete
sensitivity analysis (Pal et al., 2021).

  4. 4.

TayNODE (Kelly et al., 2020) and STEER (Behl et al., 2020) for models reported
in Pal et al. (2021).

Figure 3: MNIST Classification using Neural ODE

We use the DifferentialEquations.jl (Rackauckas et al., 2019) and Lux.jl (Pal,
2022) software stack written in the Julia Programming Language (Bezanson et
al., 2017) for all our experiments.

Some details about the data presented in the tables:

  * •

All experimental results in the tables marked with ¶ were taken directly from
Pal et al. (2021).

  * •

We have tried to match the hardware details presented in the paper and the
corresponding GitHub repository for Pal et al. (2021), but we note that
differences in wall clock timings can be partially attributed to hardware.

  * •

TayNODE (Kelly et al., 2020) uses a different ODE integrator. Hence the NFEs
are not directly comparable.

Method Test Loss ($\times 10^{-3}$) Training Time (hr) Prediction Time (s /
batch) Testing NFE Vanilla NODE $3.41\pm 0.10$ $2.48\pm 0.22$ $0.16\pm 0.01$
$758.0\pm 25.87$ STEER ${}^{\text{\textparagraph}}$ $3.48\pm 0.01$ $1.62\pm
0.26$ $0.54\pm 0.06$ $699.0\pm 141.1$ TayNODE ${}^{\text{\textparagraph}}$
$4.21\pm 0.01$ $12.3\pm 0.32$ $0.22\pm 0.02$ $167.3\pm 11.93$ ERNODE
${}^{\text{\textparagraph}}$ $3.57\pm 0.00$ $0.94\pm 0.13$ $0.21\pm 0.02$
$287.0\pm 17.32$ SRNODE ${}^{\text{\textparagraph}}$ $3.58\pm 0.05$ $0.87\pm
0.09$ $0.20\pm 0.01$ $273.0\pm 0.000$ Local Unbiased ERNODE $3.64\pm 0.07$
$2.31\pm 0.02$ $0.09\pm 0.00$ $422.0\pm 4.580$ Local Biased ERNODE $3.63\pm
0.08$ $2.12\pm 0.24$ $0.10\pm 0.01$ $463.0\pm 63.02$

Table 3: Physionet Time Series Interpolation: Local Regularization reduces
NFEs by $\mathit{0.556\times-0.610\times}$ reducing the prediction timings by
$\mathit{1.6\times-1.78\times}$. Our methods additionally improve training
timings by $\mathit{1.073\times-1.167\times}$. We note that the difference in
training time compared to (E/S)RNODE methods is due to change in the
sensitivity algorithm.

###  4.1 MNIST Image Classification

We train a neural differential equation classifier to map flattened MNIST
(LeCun et al., 1998) images to their corresponding labels.

####  4.1.1 Neural Ordinary Differential Equation

Training Details: We use the same model architecture as described in Kelly et
al. (2020). Our model comprises of single hidden layered explicit model
$f_{\theta}$ modeling the ODE dynamics followed by a linear classifier
$g_{\phi}$. The hidden layer is 100-dimensional. We train with a batch size of
512 for a total of 7500 steps. We use Adam (Kingma & Ba, 2014) with a constant
learning rate of $0.001$. For error estimate regularization, we exponentially
decrease the regularization coefficient from $2.5$ to $1.0$. We use Tsit5
(Tsitouras, 2011) as the ODE integrator with an absolute and relative
tolerance of $10^{-8}$.222We note that this is not a realistic tolerance at
which image classification models are trained. We use this tolerance to allow
a direct comparison to prior works.

Baselines: We consider a Vanilla Neural ODE trained with the exact
aforementioned specifications. All other baselines are directly taken from Pal
et al. (2021).

Results: We summarize the results in Table 2 and Figure 3. Using local
regularization speeds up prediction in all cases while it leads to a minor
slowdown during training for biased sampling.

####  4.1.2 Neural Stochastic Differential Equation

Figure 4: MNIST Classification using Neural SDE

Training Details: We downsample the flattened images to a 32-dimensional
vector before feeding it into the Neural SDE which uses a diffusion model
($f_{\theta}$) having a 64-dimensional hidden layer and a linear drift model
($g_{\phi}$). Finally, a linear classifier ($h_{\gamma}$) predicts the label.
We train our models on CPU with a batch size of $512$ for a total of $4000$
steps. We optimize the weights using Adam (Kingma & Ba, 2014) with a constant
learning rate of $0.01$. We use SOSRI2 SDE solver (Rackauckas & Nie, 2017)
with a tolerance of $0.14$. We fix our regularization coefficient to be
$10^{3}$. For this experiment, we rely on discrete sensitivity analysis.

Baselines: ERNSDE and SRNSDE results were taken from Pal et al. (2021). These
were trained for $40$ epochs, nearly equivalent to training for $4000$
iterations.

Results: We summarize the results in Table 2 and Figure 4. Local
regularization improves training and prediction performance while keeping the
test accuracy nearly constant.

###  4.2 Physionet Time Series Interpolation

Figure 5: Physionet Time Series Interpolation using Latent ODE

Training Details: We use the experimental setup for Physionet 2012 Challenge
Dataset (Citi & Barbieri, 2012) from Kelly et al. (2020). We use a Latent
Neural ODE (Rubanova et al., 2019) to perform time series interpolation on the
dataset. We use the preprocessed dataset from Kelly et al. (2020) to ensure a
fair comparison and independent runs are performed using an 80:20 split of the
dataset.

Configuration Method Train Accuracy (%) Test Accuracy (%) Training Time (s /
batch) Prediction Time (s / batch) Testing NFE Standard Vanilla $83.683\pm
1.450$ $67.394\pm 0.849$ $0.457\pm 0.018$ $0.130\pm 0.013$ $115.315\pm 12.136$
Local Unbiased ER $83.665\pm 0.805$ $67.678\pm 0.874$ $0.399\pm 0.014$
$0.096\pm 0.007$ $\hphantom{0}89.048\pm\hphantom{0}7.335$ Local Biased ER
$83.958\pm 1.032$ $67.745\pm 0.824$ $0.555\pm 0.008$ $0.088\pm 0.003$
$\hphantom{0}81.301\pm\hphantom{0}1.255$ Multi-Scale Vanilla $92.807\pm
12.458$ $80.048\pm 6.740$ $0.572\pm 0.012$ $0.170\pm 0.005$
$\hphantom{0}27.616\pm\hphantom{0}0.905$ Local Unbiased ER
$94.159\pm\hphantom{0}9.694$ $80.432\pm 5.548$ $0.641\pm 0.025$ $0.175\pm
0.019$ $\hphantom{0}27.760\pm\hphantom{0}0.177$ Local Biased ER
$99.987\pm\hphantom{0}0.023$ $83.460\pm 0.727$ $0.774\pm 0.293$ $0.163\pm
0.015$ $\hphantom{0}26.334\pm\hphantom{0}0.992$

Table 4: CIFAR10 Image Classification using Neural DE: For the standard Neural
ODE, local regularization reduces the NFE by
$\mathit{0.705\times-0.772\times}$, thereby improving prediction timings by
$\mathit{1.35\times-1.477\times}$. However, unregularized model training takes
$\mathit{0.823\times}$ the time for the biased model. For multi-scale models,
the NFE and prediction time improvements are marginal and come at the cost of
higher training time.

For specific model architecture details, we refer the readers to Pal et al.
(2021). We train the model for a total of $3000$ iterations using Adamax
(Kingma & Ba, 2014) with a learning rate of $0.01$ with $10^{-5}$ inverse
decay per step. We use a batch size of $512$. We diverge from Pal et al.
(2021), in using the regularization term as
$\left(\texttt{E}_{\texttt{Est}}\right)_{\texttt{t}_{\texttt{reg}}}\cdot|dt|_{\texttt{t}_{\texttt{reg}}}$
instead of the squared regularization term
$\sum_{j}\left(\texttt{E}_{\texttt{Est}}\right)_{j}^{2}$. Additionally, we
decay the regularization coefficient exponentially from $100$ to $10$ over the
$3000$ training iterations.

Baselines: Vanilla NODE was trained with the exact aforementioned
configuration. All the other baselines were trained using discrete sensitivity
analysis, and the exact details are present in Pal et al. (2021).

Results: We summarize the results in Figure 5 and Table 3.

###  4.3 CIFAR10 Image Classification

####  4.3.1 Neural Ordinary Differential Equation

Figure 6: CIFAR10 Image Classification using Standard Neural ODE

Training Details: We use the CNN architecture for CIFAR10 as described in Poli
et al. (2020). We train the models for $31250$ steps with Adam (Kingma & Ba,
2014) using a cosine-annealing learning rate scheduler from $0.003$ to
$0.0001$. We train the models with a batch size of $32$ and keep the
regularization coefficient fixed at $2.5$. We use Tsit5 (Tsitouras, 2011) with
a tolerance of $10^{-4}$.

Results: We summarize the results in Figure 6 and Table 4.

####  4.3.2 Multiscale Neural ODE

Figure 7: CIFAR10 Image Classification using Multi-Scale Neural ODE

Training Details: We modify the Tiny Multiscale DEQ architecture for CIFAR10
from Bai et al. (2020) as Multiscale Neural ODE with Input Injection (See
Section 2.5). To stabilize the training for larger models, we exponentially
increase the regularization coefficient from $0.1$ to $5.0$. We train with a
batch size of $128$ using VCAB3 (Wanner & Hairer, 1996) with a tolerance of
$0.05$.

Results: We summarize the results in Figure 7 and Table 4. The benefits from
regularization for NFEs and prediction timings seem marginal. However,
regularization using biased sampling makes the training dynamics stable as
observed in Figure 7.

##  5 Related Works

Neural Differential Equations (and by extension most Implicit Neural Networks)
are mostly constrained by expensive training and prediction timings (Dupont et
al., 2019; Kelly et al., 2020; Finlay et al., 2020; Behl et al., 2020; Pal et
al., 2021). Traditionally accelerating these models have relied on using
higher-order regularization terms to constrain the space of learnable dynamics
(Finlay et al., 2020; Kelly et al., 2020). These models speed up predictions,
but their benefits are often overshadowed by a massive training slowdown (Pal
et al., 2021).

More recently, quite a few first-order schemes have been proposed. Behl et al.
(2020) randomized the endpoint of Neural ODEs to incentivize simpler dynamics.
However, Pal et al. (2021) didn’t find significant benefits of using STEER in
their experiments. Pal et al. (2021) used internal solver heuristics – local
error and stiffness estimates – to control the learned dynamics in a way that
decreased both prediction and training time. We discuss their work in detail
in Section 2.4. Xia et al. (2021) rewrite Neural ODEs as heavy ball ODEs to
accelerate both forward and backward passes. Djeumou et al. (2022) replace ODE
solvers in the forward with a Taylor-Lagrange expansion and report
significantly better training and prediction times.

##  6 Discussion

In this manuscript, we have shown that we can obtain similar properties to
global regularization by regularizing dynamical systems at randomly sampled
time points. Additionally, this comes with the benefit of not being forced
into a specific sensitivity analysis method. We have taken every experiment in
Pal et al. (2021) and empirically showed that our local regularization works
at par with global regularization. However, our experiments using stiffness
estimate for local regularization did not yield positive results and were not
presented in this manuscript. Thus, we have demonstrated that we can “close
the blackbox” and still leverage all the benefits of internal solver
heuristics to improve training and predictions of neural differential
equations.

###  6.1 Limitations

We note the following limitations of our work:

  * •

Similar to Pal et al. (2021), if the objective is to learn the actual
dynamical system, our method will not yield proper results. Our method is
applicable only when the final state is relevant, i.e., in most classical deep
learning tasks.

  * •

Regularization introduces a new regularization coefficient hyperparameter
which, if not tuned correctly, can lead to unstable dynamics or might negate
the scheme’s usefulness.

##  7 Acknowledgement

The authors acknowledge the MIT SuperCloud and Lincoln Laboratory
Supercomputing Center for providing HPC resources that have contributed to the
research results reported within this paper. This material is based upon work
supported by the National Science Foundation under grant no. OAC-1835443,
grant no. SII-2029670, grant no. ECCS-2029670, grant no. OAC-2103804, and
grant no. PHY-2021825. We also gratefully acknowledge the U.S. Agency for
International Development through Penn State for grant no. S002283-USAID. The
information, data, or work presented herein was funded in part by the Advanced
Research Projects Agency-Energy (ARPA-E), U.S. Department of Energy, under
Award Number DE-AR0001211 and DE-AR0001222. We also gratefully acknowledge the
U.S. Agency for International Development through Penn State for grant no.
S002283-USAID. The views and opinions of authors expressed herein do not
necessarily state or reflect those of the United States Government or any
agency thereof. This material was supported by The Research Council of Norway
and Equinor ASA through Research Council project “308817 - Digital wells for
optimal production and drainage”. Research was sponsored by the United States
Air Force Research Laboratory and the United States Air Force Artificial
Intelligence Accelerator and was accomplished under Cooperative Agreement
Number FA8750-19-2-1000. The views and conclusions contained in this document
are those of the authors and should not be interpreted as representing the
official policies, either expressed or implied, of the United States Air Force
or the U.S. Government. The U.S. Government is authorized to reproduce and
distribute reprints for Government purposes notwithstanding any copyright
notation herein.

## References

  * Anantharaman et al. (2020) Anantharaman, R., Ma, Y., Gowda, S., Laughman, C., Shah, V., Edelman, A., and Rackauckas, C.  Accelerating simulation of stiff nonlinear systems using continuous-time echo state networks.  _arXiv preprint arXiv:2010.04004_ , 2020.
  * Ascher & Petzold (1998) Ascher, U. M. and Petzold, L. R.  _Computer methods for ordinary differential equations and differential-algebraic equations_ , volume 61.  Siam, 1998.
  * Bai et al. (2019) Bai, S., Kolter, J. Z., and Koltun, V.  Deep equilibrium models.  _arXiv preprint arXiv:1909.01377_ , 2019.
  * Bai et al. (2020) Bai, S., Koltun, V., and Kolter, J. Z.  Multiscale Deep Equilibrium Models.  _arXiv:2006.08656 [cs, stat]_ , November 2020.  URL http://arxiv.org/abs/2006.08656.  arXiv: 2006.08656.
  * Behl et al. (2020) Behl, H., Ghosh, A., Dupont, E., Torr, P., and Namboodiri, V.  Steer: simple temporal regularization for neural odes.  pp. 1–13. Neural Information Processing Systems Foundation, Inc., 2020\.
  * Bezanson et al. (2017) Bezanson, J., Edelman, A., Karpinski, S., and Shah, V. B.  Julia: A fresh approach to numerical computing.  _SIAM Review_ , 59(1):65–98, 2017.  doi: 10.1137/141000671.
  * Burt & Adelson (1987) Burt, P. J. and Adelson, E. H.  The laplacian pyramid as a compact image code.  In _Readings in computer vision_ , pp. 671–679. Elsevier, 1987\.
  * Chen et al. (2016) Chen, L.-C., Yang, Y., Wang, J., Xu, W., and Yuille, A. L.  Attention to scale: Scale-aware semantic image segmentation.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 3640–3649, 2016.
  * Chen et al. (2017) Chen, L.-C., Papandreou, G., Kokkinos, I., Murphy, K., and Yuille, A. L.  Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs.  _IEEE transactions on pattern analysis and machine intelligence_ , 40(4):834–848, 2017.
  * Chen et al. (2018) Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K.  Neural ordinary differential equations.  In _Advances in neural information processing systems_ , pp. 6571–6583, 2018.
  * Citi & Barbieri (2012) Citi, L. and Barbieri, R.  Physionet 2012 challenge: Predicting mortality of icu patients using a cascaded svm-glm paradigm.  In _2012 Computing in Cardiology_ , pp. 257–260. IEEE, 2012.
  * Dauvergne & Hascoët (2006) Dauvergne, B. and Hascoët, L.  The data-flow equations of checkpointing in reverse automatic differentiation.  In _International Conference on Computational Science_ , pp. 566–573. Springer, 2006.
  * Djeumou et al. (2022) Djeumou, F., Neary, C., Goubault, E., Putot, S., and Topcu, U.  Taylor-lagrange neural ordinary differential equations: Toward fast training and evaluation of neural odes, 2022.
  * Dormand & Prince (1980) Dormand, J. R. and Prince, P. J.  A family of embedded runge-kutta formulae.  _Journal of computational and applied mathematics_ , 6(1):19–26, 1980.
  * Dupont et al. (2019) Dupont, E., Doucet, A., and Teh, Y. W.  Augmented neural odes.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Farabet et al. (2012) Farabet, C., Couprie, C., Najman, L., and LeCun, Y.  Learning hierarchical features for scene labeling.  _IEEE transactions on pattern analysis and machine intelligence_ , 35(8):1915–1929, 2012.
  * Fehlberg (1968) Fehlberg, E.  _Classical fifth-, sixth-, seventh-, and eighth-order Runge-Kutta formulas with stepsize control_.  National Aeronautics and Space Administration, 1968.
  * Finlay et al. (2020) Finlay, C., Jacobsen, J.-H., Nurbekyan, L., and Oberman, A. M.  How to train your neural ode.  _arXiv preprint arXiv:2002.02798_ , 2020.
  * Grathwohl et al. (2018) Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D.  Ffjord: Free-form continuous dynamics for scalable reversible generative models.  _arXiv preprint arXiv:1810.01367_ , 2018.
  * Hairer et al. (1993) Hairer, E., Norsett, S., and Wanner, G.  _Solving Ordinary Differential Equations I: Nonstiff Problems_ , volume 8.  01 1993.  ISBN 978-3-540-56670-0.  doi: 10.1007/978-3-540-78862-1.
  * Hindmarsh et al. (2005) Hindmarsh, A. C., Brown, P. N., Grant, K. E., Lee, S. L., Serban, R., Shumaker, D. E., and Woodward, C. S.  Sundials: Suite of nonlinear and differential/algebraic equation solvers.  _ACM Trans. Math. Softw._ , 31(3):363–396, sep 2005.  ISSN 0098-3500.  doi: 10.1145/1089014.1089020.  URL https://doi.org/10.1145/1089014.1089020.
  * Kelly et al. (2020) Kelly, J., Bettencourt, J., Johnson, M. J., and Duvenaud, D.  Learning differential equations that are easy to solve.  _arXiv preprint arXiv:2007.04504_ , 2020.
  * Kim et al. (2021) Kim, S., Ji, W., Deng, S., Ma, Y., and Rackauckas, C.  Stiff neural ordinary differential equations.  _Chaos: An Interdisciplinary Journal of Nonlinear Science_ , 31(9):093122, 2021.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization, 2014.  URL http://arxiv.org/abs/1412.6980.  cite arxiv:1412.6980Comment: Published as a conference paper at the 3rd International Conference for Learning Representations, San Diego, 2015.
  * Kutta (1901) Kutta, W.  Beitrag zur naherungsweisen integration totaler differentialgleichungen.  _Z. Math. Phys._ , 46:435–453, 1901.
  * LeCun et al. (1998) LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Liu et al. (2019) Liu, X., Xiao, T., Si, S., Cao, Q., Kumar, S., and Hsieh, C.-J.  Neural sde: Stabilizing neural ode networks with stochastic noise, 2019\.
  * Lu et al. (2018) Lu, Y., Zhong, A., Li, Q., and Dong, B.  Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations.  In _International Conference on Machine Learning_ , pp. 3276–3285. PMLR, 2018.
  * Onken & Ruthotto (2020) Onken, D. and Ruthotto, L.  Discretize-optimize vs. optimize-discretize for time-series regression and continuous normalizing flows.  _arXiv preprint arXiv:2005.13420_ , 2020.
  * Pal (2022) Pal, A.  Lux: Explicit parameterization of deep neural networks in julia.  https://github.com/avik-pal/Lux.jl/, 2022.
  * Pal et al. (2021) Pal, A., Ma, Y., Shah, V., and Rackauckas, C. V.  Opening the blackbox: Accelerating neural differential equations by regularizing internal solver heuristics.  In _International Conference on Machine Learning_ , pp. 8325–8335. PMLR, 2021.
  * Pal et al. (2022) Pal, A., Edelman, A., and Rackauckas, C.  Mixing implicit and explicit deep learning with skip deqs and infinite time neural odes (continuous deqs).  _arXiv preprint arXiv:2201.12240_ , 2022.
  * Piché (1995) Piché, R.  An l-stable rosenbrock method for step-by-step time integration in structural dynamics.  _Computer Methods in Applied Mechanics and Engineering_ , 126(3-4):343–354, 1995.
  * Poli et al. (2020) Poli, M., Massaroli, S., Yamashita, A., Asama, H., and Park, J.  Hypersolvers: Toward fast continuous-depth models.  _arXiv preprint arXiv:2007.09601_ , 2020.
  * Rackauckas (2022) Rackauckas, C.  Engineering trade-offs in automatic differentiation: From tensorflow and pytorch to jax and julia, Jan 2022.  URL https://tinyurl.com/2j94bhc9.
  * Rackauckas & Nie (2017) Rackauckas, C. and Nie, Q.  Adaptive methods for stochastic differential equations via natural embeddings and rejection sampling with memory.  _Discrete and continuous dynamical systems. Series B_ , 22(7):2731, 2017.
  * Rackauckas & Nie (2020) Rackauckas, C. and Nie, Q.  Stability-optimized high order methods and stiffness detection for pathwise stiff stochastic differential equations.  In _2020 IEEE High Performance Extreme Computing Conference (HPEC)_ , pp. 1–8. IEEE, 2020.
  * Rackauckas et al. (2019) Rackauckas, C., Innes, M., Ma, Y., Bettencourt, J., White, L., and Dixit, V.  Diffeqflux.jl - A julia library for neural differential equations.  _CoRR_ , abs/1902.02376, 2019.  URL http://arxiv.org/abs/1902.02376.
  * Rackauckas et al. (2020) Rackauckas, C., Ma, Y., Martensen, J., Warner, C., Zubov, K., Supekar, R., Skinner, D., Ramadhan, A., and Edelman, A.  Universal differential equations for scientific machine learning, 2020\.
  * Robertson (1966) Robertson, H.  The solution of a set of reaction rate equations.  _Numerical analysis: an introduction_ , 178182, 1966.
  * Rubanova et al. (2019) Rubanova, Y., Chen, R. T., and Duvenaud, D.  Latent odes for irregularly-sampled time series.  _arXiv preprint arXiv:1907.03907_ , 2019.
  * Runge (1895) Runge, C.  Über die numerische auflösung von differentialgleichungen.  _Mathematische Annalen_ , 46(2):167–178, 1895\.
  * Tsitouras (2011) Tsitouras, C.  Runge–kutta pairs of order 5(4) satisfying only the first column simplifying assumption.  _Computers & Mathematics with Applications_, 62(2):770 – 775, 2011.  ISSN 0898-1221.  doi: https://doi.org/10.1016/j.camwa.2011.06.002.  URL http://www.sciencedirect.com/science/article/pii/S0898122111004706.
  * Wanner & Hairer (1996) Wanner, G. and Hairer, E.  _Solving ordinary differential equations II_.  Springer Berlin Heidelberg, 1996.
  * Xia et al. (2021) Xia, H., Suliafu, V., Ji, H., Nguyen, T., Bertozzi, A., Osher, S., and Wang, B.  Heavy ball neural ordinary differential equations.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Yu & Koltun (2015) Yu, F. and Koltun, V.  Multi-scale context aggregation by dilated convolutions.  _arXiv preprint arXiv:1511.07122_ , 2015.
  * Zhang & Sandu (2014) Zhang, H. and Sandu, A.  Fatode: a library for forward, adjoint, and tangent linear integration of odes.  _SIAM Journal on Scientific Computing_ , 36(5):C504–C523, 2014.