11institutetext: Inria 11email:
<EMAIL_ADDRESS>
# Enhanced Distribution Modelling via Augmented Architectures For Neural ODE
Flows

Etrit Haxholli  Marco Lorenzi

###### Abstract

While the neural ODE formulation of normalizing flows such as in FFJORD
enables us to calculate the determinants of free form Jacobians in
$\mathcal{O}(D)$ time, the flexibility of the transformation underlying neural
ODEs has been shown to be suboptimal. In this paper, we present AFFJORD, a
neural ODE-based normalizing flow which enhances the representation power of
FFJORD by defining the neural ODE through special augmented transformation
dynamics which preserve the topology of the space. Furthermore, we derive the
Jacobian determinant of the general augmented form by generalizing the chain
rule in the continuous sense into the _cable rule_ , which expresses the
forward sensitivity of ODEs with respect to their initial conditions. The
cable rule gives an explicit expression for the Jacobian of a neural ODE
transformation, and provides an elegant proof of the instantaneous change of
variable. Our experimental results on density estimation in synthetic and high
dimensional data, such as MNIST, CIFAR-10 and CelebA ($32\times 32$), show
that AFFJORD outperforms the baseline FFJORD through the improved flexibility
of the underlying vector field.

##  1 Introduction

Normalizing flows are diffeomorphic random variable transformations providing
a powerful theoretical framework for generative modeling and probability
density estimation [1]. While the practical application of normalizing flows
is generally challenging due to computational bottlenecks, most notably
regarding the $\mathcal{O}(D^{3})$ computation cost of the Jacobian
determinant, different architectures have been proposed in order to scale
normalizing flows to high dimensions while at the same time ensuring the
flexibility and bijectivity of the transformations [1, 2, 3]. The common
strategy consists of placing different architectural restrictions on the
model, to enforce special Jacobian forms, with less computationally demanding
determinants.

A noteworthy approach is based on neural ODEs [4], as they enable us to
calculate the determinants of free form Jacobians in $\mathcal{O}(D)$ time
[5]. More specifically, the rule for the instantaneous change of variable [4]
provides an important theoretical contribution to normalizing flows, as it
yields a closed form expression of the Jacobian determinant of a neural ODE
transformation.

In this case, calculating the Jacobian determinant simplifies to calculating
the integral of the divergence of the vector field along the transformation
trajectory. Such models are known as Continuous Normalizing Flows (CNFs). In
[5], these ideas are further explored and computational simplifications are
introduced, notably the use of Hutchinson’s trace estimator [11]. The
resulting model is named FFJORD.

---
Figure 1: Vector field evolution in AFFJORD as a function of the augmented
dimensions $\boldsymbol{z}^{*}_{t}$.

In [6], it is shown that there exist functions which neural ODEs are not
capable of representing. To tackle this issue and enhance the expressiveness
of the neural ODE transformation, they propose to lift the data into a higher
dimensional space, on which the neural ODE is applied, to subsequently project
the output back to the original space. Such augmented neural ODEs (ANODEs)
have been experimentally shown to lead to improved flexibility and
generalisation properties than the non-augmented counterpart.

Inspired by [6], in this paper we develop a theoretical framework to increase
the flexibility of ODE flows, such as FFJORD, through dimension augmentation.
To this end, we derive an explicit formula for Jacobians corresponding to
neural ODE transformations. This formula represents the continuous
generalization of the chain rule, which we name _the cable rule_ , that
ultimately allows the derivation of the Jacobian expression and determinant
for the composition of the operations defining the ANODE flow: augmentation,
neural ODE transformation, and projection.

To enable computational feasibility and ensure that the ANODE transformation
is diffeomorphic, we allow the augmented dimensions to parameterize the vector
field acting on the original dimensions (but not vice-versa). We name the
resulting model augmented FFJORD (AFFJORD). In AFFJORD, the evolution of time
is high dimensional (Figure 1), and is learnt via the vector field defined by
the augmented component, contrasting the linear time evolution of FFJORD. This
setup coincides with the framework introduced by [18], but in the context of
normalizing flows. Our experiments on 2D data of toy distributions and image
datasets such as MNIST, CIFAR-10 and CelebA ($32\times 32$), show that the
proposed ANODE flow defined by AFFJORD outperforms FFJORD in terms of density
estimation, thus highlighting the improved flexibility and representation
properties of the underlying vector field.

##  2 Background and Related Work

Normalizing Flows: The Normalizing Flow framework was previously defined in
[7, 8], and was popularised by [1] and by [9] respectively in the context of
variational inference and density estimation.
A Normalizing Flow is a transformation defined by a sequence of invertible and
differentiable functions mapping a simple base probability distribution (e.g.,
a standard normal) into a more complex one. Let $\boldsymbol{Z}$ and
$\boldsymbol{X}=g(\boldsymbol{Z})$ be random variables where $g$ is a
diffeomorphism with inverse $h$. If we denote their probability density
functions by $f_{\boldsymbol{Z}}$ and $f_{\boldsymbol{X}}$, based on the
change of variable theorem we get:

$f_{\boldsymbol{X}}(\boldsymbol{x})=f_{\boldsymbol{Z}}(\boldsymbol{z})|\frac{d\boldsymbol{z}}{d\boldsymbol{x}}|=f_{\boldsymbol{Z}}(h(\boldsymbol{x}))|\frac{dh(\boldsymbol{x})}{d\boldsymbol{x}}|.$
(1)

In general, we want to optimize the parameters of $h$ such that we maximize
the likelihood of sampled points
${\boldsymbol{x}}_{1},...,{\boldsymbol{x}}_{n}$. Once these parameters are
optimized, then we can give as input any test point ${\boldsymbol{x}}$ on the
right hand side of Equation 1, and calculate its likelihood. For the
generative task, being able to easily recover $g$ from $h$ is essential, as
the generated point ${\boldsymbol{x}}_{g}$ will take form
${\boldsymbol{x}}_{g}$=$g(\boldsymbol{z}_{s})$, where $\boldsymbol{z}_{s}$ is
a sampled point from the base distribution $f_{\boldsymbol{Z}}$.
For increased modeling flexibility, we can use a chain (flow) of
transformations, $\boldsymbol{z}_{i}=g_{i}(\boldsymbol{z}_{i-1}),i\in[n]$. In
this case due to chain rule we have:

$f_{\boldsymbol{Z}_{n}}(\boldsymbol{z}_{n})=f_{\boldsymbol{Z}_{0}}(\boldsymbol{z}_{0})|\frac{d\boldsymbol{z}_{0}}{d\boldsymbol{z}_{n}}|=f_{\boldsymbol{Z}_{0}}(h_{0}(...h_{n}(\boldsymbol{z}_{n})))\prod_{i=1}^{n}|\frac{dh_{i}(\boldsymbol{z}_{i})}{d\boldsymbol{z}_{i}}|$
(2)

The interested reader can find a more in-depth review of normalizing flows in
[3] and [10].

Neural ODE Flows: Neural ODEs, [4], are continuous generalizations of residual
networks:

$\boldsymbol{z}_{t_{i+1}}=\boldsymbol{z}_{i}+\epsilon
f(\boldsymbol{z}_{t_{i}},t_{i},\boldsymbol{\theta})\rightarrow\boldsymbol{z}(t)=\boldsymbol{z}(0)+\int_{0}^{t}f(\boldsymbol{z}(\tau),\tau,\boldsymbol{\theta})d\tau\text{,
as }\epsilon\rightarrow 0.$ (3)

In [13, 4], it is shown that the gradients of neural ODEs can be computed via
the adjoint method, with constant memory cost with regards to "depth":

$\frac{dL}{d\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}dt=-\int_{T}^{0}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}dt,$
(4)

where $\frac{\partial L}{\partial\boldsymbol{z}(t)}$ can be calculated
simultaneously by

$\frac{\partial L}{\partial\boldsymbol{z}(t)}=\frac{\partial
L}{\partial\boldsymbol{z}(T)}-\int_{T}^{t}\frac{\partial
L}{\partial\boldsymbol{z}(\tau)}\frac{\partial
f(\boldsymbol{z}(\tau),\tau,\boldsymbol{\theta}(\tau))}{\partial\boldsymbol{z}(\tau)}d\tau.$
(5)

In addition, they also derive the expression for the instantaneous change of
variable, which enables one to train continuous normalizing flows:

$\log{p(\boldsymbol{z}(0))}=\log{p(\boldsymbol{z}(T))}+\int_{0}^{T}tr{\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}dt},$
(6)

where $\boldsymbol{z}(0)$ represents a sample from the data. A surprising
benefit is that one does not need to calculate the determinant of the Jacobian
of the transformation anymore, but simply the trace of a matrix. These models
are collectively called Neural ODE flows (NODEFs) or simply Continuous
Normalizing Flows (CNFs). In [5], these ideas are further explored and
computational simplifications are introduced, notably the use of Hutchinson’s
trace estimator, [11, 12], as an unbiased stochastic estimator of the trace in
the likelihood expression in Equation 6. The resulting model is named FFJORD.

Multiscale Architectures: In [2], a multiscale architecture for normalizing
flows is implemented, which transforms the data shape from $[c,s,s]$ to
$[4c,\frac{s}{2},\frac{s}{2}]$, where $c$ is the number of channels and $s$ is
the height and width of the image. Effectively this operation trades spatial
size for additional channels, and after each transformation a normalizing flow
is applied on half the channels, while the other half are saved in an array.
This process can be repeated as many times as the width and height of the
transformed image remain even numbers. In the end, all saved channels are
concatenated to construct an image with the original dimensions. A visual
description of this process can be found in Figure 2.

Figure 2: Example of the multiscale architecture on MNIST dataset

Augmented Neural ODEs: Considering that transformations by neural ODEs are
diffeomorphisms (hence homeomorphisms), [6] show that Neural Ordinary
Differential Equations (NODEs) learn representations that preserve the
topology of the input space, and prove that this implies the existence of
functions that Neural ODEs cannot represent. To address these limitations,
they introduce the Augmented Neural ODEs:

$\frac{d}{dt}\begin{bmatrix}\boldsymbol{z}(t)\\\
\boldsymbol{z}^{*}(t)\end{bmatrix}=h(\begin{bmatrix}\boldsymbol{z}(t)\\\
\boldsymbol{z}^{*}(t)\end{bmatrix},t)=\begin{bmatrix}f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})\\\
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})\end{bmatrix}$
$\text{ for }\begin{bmatrix}\boldsymbol{z}(0)\\\
\boldsymbol{z}^{*}(0)\end{bmatrix}=\begin{bmatrix}\boldsymbol{x}\\\
0\end{bmatrix},$ (7)

where $\boldsymbol{z}^{*}(t)$ is the augmented component, and $h=[f,g]$ is the
vector field to be learnt.

In addition to being more expressive models, [6] show that augmented neural
ODEs are empirically more stable, generalize better and have a lower
computational cost than Neural ODEs. A schematic of their architecture in the
discrete case can be found in Figure 3.

---
Figure 3: Architecture of discretized augmented neural ODEs (left) and of the
discretized augmented neural ODE flows implemented in AFFJORD (right).

##  3 Proposed Framework

In the first subsection, we introduce our model AFFJORD, which is a special
case of augmented neural ODEs. In the second subsection, we give the
generalisation of the chain rule in the continuous sense which we refer to as
the cable rule, which is analogous to forward sensitivity. Next, we give an
intuitive proof of the continuous backpropagation, by showing its equivalence
to the continuous generalisation of the total derivative decomposition. Then,
using the cable rule, we give a more detailed explanation on why the
instantaneous change of variable still holds in our model. In the final
section, the augmented multiscale architecture is described, as it is
implemented in the experimental section.

###  3.1 Proposed Model: Augmented FFJORD (AFFJORD)

A neural ODE of the type
$f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})$, where
$\boldsymbol{z}^{*}(t)=\boldsymbol{z}^{*}(t)+\int_{0}^{t}g(\boldsymbol{z}^{*}(\tau),\boldsymbol{\phi})d\tau$,
can equivalently be written as
$h_{f}(\boldsymbol{z}(t),t,\boldsymbol{\theta})=f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})$.
Thus, motivated by [6], we propose to lift each data point $\boldsymbol{z}(0)$
$\in\mathbb{R}^{n}$ to a higher dimensional space $\mathbb{R}^{n+m}$, by
augmentation via an $m$ dimensional vector $\boldsymbol{z}^{*}(0)$, which in
practice is set to be the zero vector. Therefore, the joint vector
$[\boldsymbol{z}(0),\boldsymbol{z}^{*}(0)]$ is transformed by the vector field
$h=(f,g)$:

$w(T)=\begin{bmatrix}\boldsymbol{z}(T)\\\
\boldsymbol{z}^{*}(T)\end{bmatrix}=\begin{bmatrix}\boldsymbol{z}(0)\\\
\boldsymbol{z}^{*}(0)\end{bmatrix}+\int_{0}^{T}\begin{bmatrix}f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})\\\
g(\boldsymbol{z}^{*}(t),\boldsymbol{\phi})\end{bmatrix}dt.$ (8)

We are not interested in $z^{*}(T)$ as our interest lies on the transformed
$z(0)$, that is $z(T)$. Since by definition such coupled dynamics are
contained in the formulation of neural ODEs, this implies that the
instantaneous change of formula still holds, and the transformation is
injective. This sort of augmentation can be seen as a special case of
augmentation introduced in [6], where the augmented dimensions depend only on
themselves. The augmented dimensions $\boldsymbol{z}^{*}(t)$ can also be seen
as time dependent weights of the non-autonomous $f$ whose evolution is
determined by the autonomous ODE $g$, hence giving $f$ greater flexibility in
time [18].

###  3.2 The Cable Rule

If we define a chain of transformations

$\boldsymbol{z}_{i}=g_{i}(\boldsymbol{z}_{i-1})\text{ for }i\in\\{1,...,n\\},$
(9)

due to the chain rule we have:

$\frac{d\boldsymbol{z}_{n}}{d\boldsymbol{z}_{0}}=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}_{n-1}}\frac{d\boldsymbol{z}_{n-1}}{d\boldsymbol{z}_{0}}=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}_{n-1}}\frac{\partial\boldsymbol{z}_{n-1}}{\partial\boldsymbol{z}_{n-2}}...\frac{\partial\boldsymbol{z}_{1}}{\partial\boldsymbol{z}_{0}}=$
$=\frac{\partial
g_{n}(\boldsymbol{z}_{n-1})}{\partial\boldsymbol{z}_{n-1}}\frac{\partial
g_{n-1}(\boldsymbol{z}_{n-2})}{\partial\boldsymbol{z}_{n-2}}...\frac{\partial
g_{1}(\boldsymbol{z}_{0})}{\partial\boldsymbol{z}_{0}}.$ (10)

We can choose each $g_{i}$ to infinitesimally modify its input, i.e.,
$g_{i}(\boldsymbol{z}_{i-1})=\boldsymbol{z}_{i-1}+\epsilon
f_{i}(\boldsymbol{z}_{i-1},t_{i-1},\boldsymbol{\theta})$, so that the chain in
Expression 9 transforms $\boldsymbol{z}_{0}$ continuously. It is clear that
$\boldsymbol{z}(t)=\boldsymbol{z}(0)+\int_{0}^{t}f(\boldsymbol{z}(t),t,\boldsymbol{\theta})d\tau$
is the limit of the previous iterative definition when $\epsilon\rightarrow
0$. Then the expression $\frac{d\boldsymbol{z}_{n}}{d\boldsymbol{z}_{0}}$ in
Equation 10 converges to $\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}$,
which as we show in Appendix 0.A, satisfies the differential equation below:

$\frac{d(\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)})}{dt}=\frac{\partial
f(\boldsymbol{z}(t))}{\partial\boldsymbol{z}(t)}\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}.$
(11)

Using the Magnus expansion [16, 17], we conclude that

$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}=e^{\boldsymbol{\Omega}(T)},\text{
for\ \ }\boldsymbol{\Omega}(t)=\sum_{k=1}^{\infty}\boldsymbol{\Omega}_{k}(t),$
(12)

where $\boldsymbol{\Omega}_{i}(t)$ are the terms of the Magnus expansion (See
Appendix 0.A). As this expression gives the generalisation of the chain rule
in the continuous sense, we refer to it as the cable rule. We notice that
Equation 11 gives the dynamics of the Jacobian of the state with respect to
the initial condition $\boldsymbol{z(0)}$. The cable rule is therefore
analogous to the forward sensitivity formula for ODEs which provides the
dynamics of Jacobian of the state with respect to the parameters
$\boldsymbol{\theta}$ of the flow [21]. This relation is highlighted and
explained in more detail in Appendix A. On this note, in Appendix 0.G, we
derive the cable rule via Equation 5. Furthermore, in Appendix 0.E, we derive
the instantaneous change of variables from the cable rule.

###  3.3 The Continuous Generalisation of the Total Derivative Decomposition
and Continuous Backpropagation

In the case that
$f=f(\boldsymbol{x}(\boldsymbol{\theta}),\boldsymbol{y}(\boldsymbol{\theta}))$,
then $\frac{df}{d\boldsymbol{\theta}}=\frac{\partial
f}{\partial\boldsymbol{x}}\frac{d\boldsymbol{x}}{d\boldsymbol{\theta}}+\frac{\partial
f}{\partial\boldsymbol{y}}\frac{d\boldsymbol{y}}{d\boldsymbol{\theta}}$, as
$\boldsymbol{\theta}$ contributes to both $\boldsymbol{x}$ and
$\boldsymbol{y}$, which in turn determine $f$. If
$\boldsymbol{z}(T)=\boldsymbol{z}(0)+\int_{0}^{T}f(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)dt$,
then $\boldsymbol{\theta}$ controls the vector field at each time point during
integration, hence we expect that these infinitesimal contributions of the
transformation from $\boldsymbol{z}(0)$ to $\boldsymbol{z}(T)$ should be
integrated. Indeed, as we prove in Appendix 0.B, the following holds:

$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta}(t))}{\partial\boldsymbol{\theta}(t)}\frac{\partial\boldsymbol{\theta}(t)}{\partial\boldsymbol{\theta}}dt,$
(13)

from which for $\boldsymbol{\theta}(t)=\boldsymbol{\theta}$, and for some
function $L=L(\boldsymbol{z}(T))$, we deduce:

$\frac{dL}{d\boldsymbol{\theta}}=\frac{\partial
L}{\partial\boldsymbol{z}(T)}\frac{d\boldsymbol{z}(T)}{d\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(T)}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}dt=\int_{0}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}dt$ (14)

thus giving an alternative and intuitive proof of continuous backpropagation
[13, 4]. In Appendix H, we show that the adjoint method can be used to prove
Equation 13, hence the continuous total derivative decomposition is equivalent
to continuous backpropagation. An alternate derivation of Equation 13 can be
found in the Appendix of [19]

###  3.4 AFFJORD as a Special Case of Augmented Neural ODE Flows

As discussed in subsection 3.1, the ODE dynamics in Equation 8, can be seen as
a special case of the following joint ODE transformation:

$w(T)=\begin{bmatrix}\boldsymbol{z}(T)\\\
\boldsymbol{z}^{*}(T)\end{bmatrix}=\begin{bmatrix}\boldsymbol{z}(0)\\\
\boldsymbol{z}^{*}(0)\end{bmatrix}+\int_{0}^{T}\begin{bmatrix}f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})\\\
g(\boldsymbol{z}^{*}(t),\boldsymbol{z}(t),\boldsymbol{\phi})\end{bmatrix}dt.$
(15)

In this general form, the model is unsuitable to be used in practice for two
reasons:
1) The transformation of $\boldsymbol{z}(0)$ to $\boldsymbol{z}(T)$ is not
necessarily injective. Indeed, the transformation from
$[\boldsymbol{z}(0),\boldsymbol{z}^{*}(0)]$ to
$[\boldsymbol{z}(T),\boldsymbol{z}^{*}(T)]$ is injective due to the
Picard–Lindelöf Theorem, however, for two data points
$\boldsymbol{z^{\prime}}(0)$ and $\boldsymbol{z^{\prime\prime}}(0)$, their
images $\boldsymbol{z^{\prime}}(T)$, $\boldsymbol{z^{\prime\prime}}(T)$ might
be identical as long as their respective augmented dimensions
$\boldsymbol{z^{\prime}}^{*}(T)$, $\boldsymbol{z^{\prime\prime}}^{*}(T)$
differ.
2) The Jacobian determinant of this general transformation is computationally
intractable. Using the chain rule we can express the Jacobian determinant of
this transformation as

$\big{|}\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}\big{|}=\big{|}\frac{d\boldsymbol{z}(T)}{d\begin{bmatrix}\boldsymbol{z}(T),\boldsymbol{z}^{*}(T)\end{bmatrix}}\frac{d\begin{bmatrix}\boldsymbol{z}(T),\boldsymbol{z}^{*}(T)\end{bmatrix}}{d\begin{bmatrix}\boldsymbol{z}(0),\boldsymbol{z}^{*}(0)\end{bmatrix}}\frac{d\begin{bmatrix}\boldsymbol{z}(0),\boldsymbol{z}^{*}(0)\end{bmatrix}}{d\boldsymbol{z}(0)}\big{|}.$
(16)

The middle term on the RHS of Equation 16 can be further developed via the
cable rule, to give the expression of the determinant of the Jacobian of
augmented neural ODE flows, which is not computationally feasible in general.
As explained in Section 3.1, the special case (AFFJORD) formulated in Equation
8, mitigates the issues mentioned above.
Regarding issue 1), we have the following:

###### Proposition 1

The architecture of AFFJORD ensures that the transformation is injective.

###### Proof

Indeed, as $\boldsymbol{z}^{*}(t)$ is not dependent on external factors, and
since $\boldsymbol{z}^{*}(0)$ is constant regarding $\boldsymbol{z}(0)$, the
end result $\boldsymbol{z}^{*}(T)$ will always be the same. Hence, for
$\boldsymbol{z^{\prime}}(0)$ and $\boldsymbol{z^{\prime\prime}}(0)$ their
images $\boldsymbol{z^{\prime}}(T)$ and $\boldsymbol{z^{\prime\prime}}(T)$
must be different, since their equality would imply
$[\boldsymbol{z^{\prime}}(T),\boldsymbol{z}^{*}(T)]=[\boldsymbol{z^{\prime\prime}}(T),\boldsymbol{z}^{*}(T)]$
contradicting the Picard–Lindelöf Theorem.

Issue 2) is mitigated as Equation 16 simplifies to

$\big{|}\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}\big{|}=\big{|}\begin{bmatrix}I,0\end{bmatrix}{\begin{bmatrix}e^{\int_{0}^{T}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}dt+\boldsymbol{\Omega}^{[z]}_{2}(T)+...}&\bar{\boldsymbol{B}}(T)\\\
0&\bar{\boldsymbol{D}}(T)\end{bmatrix}}\begin{bmatrix}I\\\
0\end{bmatrix}\big{|},$ (17)

for two block matrices $\bar{\boldsymbol{B}}(T)$ and
$\bar{\boldsymbol{D}}(T)$. This is proven in Appendix 0.F. In this case, the
$[z]$ in $\boldsymbol{\Omega}^{[z]}_{2}(T)$ denotes the restriction of the
Magnus expansion to the original dimensions. In case that the base
distribution is multivariate normal, from

$-\log{p(\boldsymbol{z}(0))}=-\log\big{[}{p(\boldsymbol{z}(T))}\
\big{|}e^{\int_{0}^{T}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}dt+...}\big{|}\big{]}$

we derive the following loss function:

$L=\frac{||\boldsymbol{\boldsymbol{Z}}(T)||^{2}}{2}-{\int_{0}^{T}tr\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t))}{\partial\boldsymbol{z}(t)}dt}.$
(18)

###  3.5 Multiscale Architecture in Augmented Neural ODE Flows

Figure 4: Example of the augmented multiscale architecture on MNIST dataset

With reference to Figure 4, the multiscale architecture in the augmented case
performs an Augmented Continuous Flow on the data as described in Subsection
3.1, then squeezes the channels (the augmented as well as the original
channels) as in the original multiscale architecture. However after the second
transformation the augmented channels are removed and stored in a separated
array (Array A). The data channels are treated as before, that is, half of
them are saved (Array B), and the other half are squeezed again. After this
process is finished we add new augmented channels, to repeat the cycle. Note
that in order to generate data by the inverse transformation, we need to
retrieve the transformed augmented dimensions previously stored (i.e., Array
A).

##  4 Experiments

We compare the performance of AFFJORD with respect to the base FFJORD on 2D
data of toy distributions, as well as on standard benchmark datasets such as
MNIST, CIFAR-10 and CelebA($32\times 32$). In the case of the 2D data, we use
the implementation of CNFs provided in [4], since using the Hutchinson’s trace
estimator and GPUs as in FFJORD provides no computational benefits in low
dimensions. In this case we also use the non-adaptive Runge-Kutta 4 ODE
solver, while for image data we use Dopri5, as well as the FFJORD
implementation of [5]. We use a batch size of $200$ for image data and a batch
size of $512$ for the toy datasets. In the case of image data, we use a
learning rate of $6\times 10^{-4}$, while for toy data the learning rate is
set to $10^{-3}$. All experiments were performed on a single GPU.

###  4.1 Toy 2D Datasets

Figure 5: Probability density modeling capabilities of AFFJORD (second column)
and FFJORD (third column) on 2D data of toy distributions.

In order to visualise the performance of the model, we first test FFJORD and
AFFJORD on 2D data of toy distributions, depicted in Figure 5. In both cases
we use the Runge-Kutta 4 solver with 40 time steps (160 function evaluations).
In these examples, we used a hypernet architecture, where the augmented
dimensions were fed to a hypernet (denoted as $hyp$) in order to generate the
weights of the field of the main dimensions. The expression of the vector
field to be learnt is the following:
$\frac{d\boldsymbol{z}}{dt}=f\big{(}[\boldsymbol{z}(t),\boldsymbol{z}^{*}(t)],\boldsymbol{\theta}(t)=hyp(\boldsymbol{z}^{*}(t),\boldsymbol{w})\big{)}$,
where
$\frac{d\boldsymbol{z}^{*}(t)}{dt}=g(\boldsymbol{z}^{*}(t),\boldsymbol{\phi})$.
Thus, the learnable parameters are $\boldsymbol{w}$ and $\boldsymbol{\phi}$.

While both FFJORD and AFFJORD are capable of modelling multi-modal and
discontinuous distributions, Figure 5 shows that AFFJORD has higher
flexibility in modeling the complex data distributions considered, in
comparison to FFJORD. In the first row, the target is the TY distribution,
where the datapoints form letters and a cluster of Gaussian distributions.
AFFJORD is more capable of separating the Gaussian spheres and modelling the
shape of the text.

On the second row AFFJORD is capable of separating the Gaussian distribution
in the center from the square that surrounds it. Furthermore, it separates the
Gaussians in the corners from the hash symbol properly. In Figure 6, we show
the results of the validation loss per iteration for both models. We can
notice that the loss of our model is roughly two standard deviations lower
than the one of FFJORD. For each model the experiment was repeated 30 times.
The experiments provided here show that AFFJORD is characterized by high
flexibility of the vector field. Indeed, in FFJORD the vector field changes
more slowly, whereas in AFFJORD the field is able to change almost abruptly,
due to this greater flexibility in time 111Videos showing the comparison of
dynamics between FFJORD and AFFJORD can be found at
https://imgur.com/gallery/kMGCKve.

|
---|---
Figure 6: Performance of FFJORD and AFFJORD on the Hash-Gaussian toy 2D
dataset (left) and TY toy 2D dataset (right).

It is important to emphasize that AFFJORD retains the ability to generate
samples from the learnt distribution, by simply integrating in the opposite
direction. Indeed, for a given sample $\boldsymbol{z}_{s}(T)$ from the base
distribution, we augment it to $\boldsymbol{z}^{*}(T)$, as
$\boldsymbol{z}^{*}(T)$ is the same for all data points. Then, we can simply
integrate backwards this concatenated vector
$[\boldsymbol{z}_{s}(T),\boldsymbol{z}^{*}(T)]$ to
$[\boldsymbol{z}_{s}(0),\boldsymbol{z}^{*}(0)]$, drop the generated augmented
dimensions $\boldsymbol{z}^{*}(0)$, and simply keep the generated data point
$\boldsymbol{z}_{s}(0)$.

###  4.2 Image Datasets

We show that AFFJORD outperforms FFJORD on MNIST, CIFAR-10 and CelebA
($32\times 32$). There are several architectures of FFJORD that can be used
for this application. Out of the architectures that we tested, the one that
performed best was the multiscale one, with three convolutional layers with 64
channels each. The number of CNF blocks was 1, and time was implemented by
simply concatenating it as a channel into the data.

Table 1:  Experimental Results for Density Estimation Models, in Bits/Dim for MNIST, CIFAR-10 and CelebA ($32\times 32$). Lower Is Better. The Multiscale Architecture Is Used in All Cases. MODEL | MNIST | CIFAR10 | CelebA($32\times 32$)
---|---|---|---
Real NVP | 1.06 | 3.49 | -
Glow | 1.05 | 3.35 | -
FFJORD | 0.96 $\pm.00$ | 3.37 $\pm.00$ | 3.28$\pm.00$
AFFJORD | $\mathbf{0.95}\pm.01$ | $\mathbf{3.32}\pm.01$ | $\mathbf{3.23}\pm.00$

When time was implemented via a hypernet, we observed that the training time
increased and performance decreased, especially in the case of CIFAR-10. For
AFFJORD we use the exact same base architecture, however, as in the case of 2D
toy data, we enable the evolution of the vector fields through time via a
hypernet which takes the augmented dimensions as an input, and outputs the
weights of the main component. The augmented dimensions are concatenated as a
channel to the main channel of the data.

---
Figure 7: Graphs of training and evaluation losses of FFJORD and AFFJORD on
CIFAR-10 as well as CelebA ($32\times 32$). We run the experiments 5 times.
Lower is better.

The formulas for both the main field and the augmented component remain
unchanged from the case of the 2D toy data, that is,
$\frac{d\boldsymbol{z}}{dt}=f\big{(}[\boldsymbol{z}(t),\boldsymbol{z}^{*}(t)],\boldsymbol{\theta}(t)=hyp(\boldsymbol{z}^{*}(t),\boldsymbol{w})\big{)}$
and
$\frac{d\boldsymbol{z}^{*}(t)}{dt}=g(\boldsymbol{z}^{*}(t),\boldsymbol{\phi})$.

The main difference here is that $g(\boldsymbol{z}^{*}(t),\boldsymbol{\phi})$
is a fully connected network with one hidden layer, which does not take as
input all the dimensions of $\boldsymbol{z}^{*}(t)$ but merely $20$ of them.
Number $20$ was chosen as during fine-tuning, the best performance was reached
in this setting. The width of the hidden layer is also $20$, as it is the
output. We fix $10$ of these $20$ dimensions and feed them to a linear
hypernet with weight matrix shape $[10,p]$ to output the $p$ weights for the
main component. It should be emphasized that the architecture of FFJORD in the
main dimensions remains unchanged in AFFJORD for fair comparison, and we only
fine-tuned the augmented structure in addition. Additional details about the
experimental settings can be found in Appendix 0.I.

As we show in Table 1, AFFJORD slightly outperforms FFJORD on MNIST on the
best run, since both models reach optimal performance, as seen from the
generated samples in Figure 8. However, our model outperforms FFJORD on the
CIFAR-10 and CelebA ($32\times 32$) dataset, as illustrated in Table 1, as
well as in Figure 7. Based on the conducted experiments the farther FFJORD is
from optimal performance, the larger the improvements brought by AFFJORD are.
The calculation of results is done as in [5], where for each run the best
evaluation result over epochs is taken. After 5 runs, for each model, the
scores are averaged and reported in Table 1.

|
---|---
Figure 8: Samples generated from AFFJORD: MNIST and CelebA $32\times 32$.

In the case of MNIST, both models were trained for roughly 9 days, while in
the case of CIFAR-10 and CelebA ($32\times 32$), they were trained for
approximately 14 days. The results corresponding to the Real NVP and Glow
models, are taken from the original papers: [2] and [14].

As in the previous case, AFFJORD can generate samples by backintegrating.
However, due to the use of the augmented multiscale architecture, where for
each cycle we replace the augmented dimensions, these replaced augmented
dimensions must be saved in an array for the backward generative pass.
Examples of samples from AFFJORD are shown in Figure 8 for both MNIST and
CelebA ($32\times 32$) datasets. Additional generated samples from both
AFFJORD and FFJORD can be found in Appendix 0.D.

##  5 Limitations and Future Work

Number of Function Evaluations (NFE). As originally reported in [5], the
number of function evaluations is one of the main bottlenecks of Neural ODE-
based models.

---
Figure 9: Number of function evaluations of FFJORD and AFFJORD when the
’concatenate’ architecture is used.

Interestingly, if the concatenation architecture is used in AFFJORD, that is,
we only replace the concatenated time channel in FFJORD with the augmented
channel in AFFJORD, then the number of function evaluations decreases
significantly. This aspect is illustrated in Figure 9. However, the hypernet
architecture of AFFJORD is prone to the issue of the number of function
evaluations. Indeed, forward passes in AFFJORD-hypernet can be challenging, as
the learnt vector field becomes too stiff with an increasing number of
integration steps.

Addition of Self-Attention  [15] describe three modeling inefficiencies in
prior work on flow models. One such factor is the lack of expressive power of
convolutional layers used in normalizing flow models. Considering the
performance improvements demonstrated in [15], in the future we intend to test
the improvement in performance brought by the addition of self-attention in
AFFJORD.

Data Dependent Augmented Dimensions As described in Section 3.1, several
simplifications are made in the architecture of the general augmented neural
ODE flows, in order to ensure immediate bijectivity and reduce the
computational complexity. However, other possible architectures exist, where
for example $\boldsymbol{z}^{*}(0)$ is dependent on $\boldsymbol{z}(0)$, and
$g(\boldsymbol{z}^{*}(t))=-\boldsymbol{z}^{*}(t)$. This would ensure that the
augmented dimensions converge to zero, providing bijectivity. Since all
augmented dimensions would be different during training, this would imply that
the data is lifted to a higher plane, enabling richer transformations. Finding
an approximation of the loss in Equation 16 remains a challenge for the future
however.

Jacobian Regularization Theoretically speaking, all performance enhancing
modifications that can be applied to FFJORD are also applicable to AFFJORD.
Such a modification which reduces the training time of FFJORD is presented in
[20], where both the vector field and its Jacobian are regularized. Thus, an
interesting research direction in the future would be to test how the
performance of AFFJORD is affected by such amendments.

##  6 Conclusion

We have presented the generalization of the total derivative decomposition in
the continuous sense as well as the continuous generalization of the chain
rule, to which we refer as the cable rule. The cable rule is analogous to the
forward sensitivity of ODEs in the sense that, it gives the dynamics of the
Jacobian of the state with respect to the initial conditions, whereas forward
sensitivity gives the dynamics of the Jacobian of the state with respect to
the parameters of the flow. Motivated by this contribution, we propose a new
type of continuous normalizing flow, namely Augmented FFJORD (AFFJORD), which
outperforms the CNF state-of-art-approach, FFJORD, in the experiments we
conducted on the task of density estimation on both 2D toy data, and on high
dimensional datasets such as MNIST, CIFAR-10 and CelebA ($32\times 32$).

## References

  * [1] Rezende, Danilo Jimenez and Mohamed, Shakir: Variational Inference with Normalizing Flows. International Conference on Machine Learning, 2015.
  * [2] Dinh, Laurent and Sohl-Dickstein, Jascha and Bengio, Samy: Density estimation using Real NVP. International Conference on Learning Representations, 2017.
  * [3] Ivan Kobyzev and Simon J.D. Prince and Marcus A. Brubaker: Normalizing Flows: An Introduction and Review of Current Methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
  * [4] Chen, Ricky T. Q. and Rubanova, Yulia and Bettencourt, Jesse and Duvenaud, David: Neural Ordinary Differential Equations. Advances in neural information processing systems, 2018.
  * [5] Grathwohl, Will and Chen, Ricky T. Q. and Bettencourt, Jesse and Sutskever, Ilya and Duvenaud, David: FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models. International Conference on Learning Representations, 2019.
  * [6] Dupont, Emilien and Doucet, Arnaud and Teh, Yee Whye: Augmented Neural ODEs. Advances in neural information processing systems, 2019.
  * [7] Esteban G. Tabak and Eric Vanden-Eijnden: Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, March 2010.
  * [8] Tabak, E. G. and Turner, Cristina V.: A Family of Nonparametric Density Estimation Algorithms. Communications on Pure and Applied Mathematics, 2013.
  * [9] Dinh, Laurent and Krueger, David and Bengio, Yoshua: NICE: Non-linear Independent Components Estimation. Workshop paper, International Conference on Learning Representations, 2015.
  * [10] George Papamakarios and Eric Nalisnick and Danilo Jimenez Rezende and Shakir Mohamed and Balaji Lakshminarayanan: Normalizing Flows for Probabilistic Modeling and Inference. Journal of Machine Learning Research, 2021.
  * [11] M.F. Hutchinson : A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics - Simulation and Computation, 1990.
  * [12] Adams, Ryan P. and Pennington, Jeffrey and Johnson, Matthew J. and Smith, Jamie and Ovadia, Yaniv and Patton, Brian and Saunderson, James : Estimating the Spectral Density of Large Implicit Matrices. ArXiv Preprint, 2018.
  * [13] Pontrjagin, L.S. and Boltyanskii, V.G. and Gamkrelidze, R.V. and Mishchenko, E.F. and Brown, D.E. : The Mathematical Theory of Optimal Processes. International series of monographs in pure and applied mathematics, 1962.
  * [14] Kingma, Diederik P. and Dhariwal, Prafulla: Glow: Generative Flow with Invertible 1x1 Convolutions. Advances in neural information processing systems, 2018.
  * [15] Ho, Jonathan and Chen, Xi and Srinivas, Aravind and Duan, Yan and Abbeel, Pieter : Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Design. International Conference on Learning Representations, 2019.
  * [16] Magnus, W.: On the exponential solution of differential equations for a linear operator. Communications on Pure and Applied Mathematics, 1954.
  * [17] S. Blanes and F. Casas and J.A. Oteo and J. Ros: The Magnus expansion and some of its applications. Physics Reports, 2009.
  * [18] Zhang, Tianjun and Yao, Zhewei and Gholami, Amir and Gonzalez, Joseph E and Keutzer, Kurt and Mahoney, Michael W and Biros, George: ANODEV2: A Coupled Neural ODE Framework. Advances in Neural Information Processing Systems, 2019.
  * [19] Massaroli, Stefano and Poli, Michael and Park, Jinkyoo and Yamashita, Atsushi and Asama, Hajime: Dissecting Neural ODEs. Advances in Neural Information Processing Systems, 2020.
  * [20] Finlay, Chris and Jacobsen, Joern-Henrik and Nurbekyan, Levon and Oberman, Adam: How to Train Your Neural ODE. International Conference on Machine Learning, 2020.
  * [21] Haihua Zhao and Vincent A. Mousseau: Extended Forward Sensitivity Analysis for Uncertainty Quantification. Nuclear Technology, 2013.

##  Appendix 0.A Cable Rule: The Continuous Generalization of the Chain Rule

We will first assume that $z$ is one dimensional. If $z_{i}=f_{i}(z_{i-1})$
for $i\in\\{1,...,n\\}$ then by the chain rule we have

$\frac{dz_{n}}{dz_{0}}=\frac{\partial z_{n}}{\partial
z_{n-1}}\frac{dz_{n-1}}{dz_{0}}=\frac{\partial z_{n}}{\partial
z_{n-1}}\frac{\partial z_{n-1}}{\partial z_{n-2}}...\frac{\partial
z_{1}}{\partial z_{0}}=\frac{\partial f_{n}(z_{n-1})}{\partial
z_{n-1}}\frac{\partial f_{n-1}(z_{n-2})}{\partial z_{n-2}}...\frac{\partial
f_{1}(z_{0})}{\partial z_{0}}$ (19)

Now, if we assume that $z_{i}=z(t_{i})$ is transformed more gradually as in
$z_{i+1}=z_{i}+\epsilon f(z_{i})$, and that $t_{i+1}=t_{i}+\epsilon$, we get
that

$\frac{dz_{n}}{dz_{0}}=\frac{\partial z_{n}}{\partial z_{n-1}}\frac{\partial
z_{n-1}}{\partial z_{n-2}}...\frac{\partial z_{1}}{\partial
z_{0}}=(I+\epsilon\frac{\partial f(z_{n-1})}{\partial
z_{n-1}})(I+\epsilon\frac{\partial f(z_{n-2})}{\partial
z_{n-2}})...(I+\epsilon\frac{\partial f(z_{0})}{\partial z_{0}})$ (20)

We see that $z(t)=z(0)+\int_{0}^{t}f(z(\tau))d\tau$ is the limit of the
previous iterative definition $z_{i+1}=z_{i}+\epsilon
f(z_{i},\boldsymbol{\theta}(t_{i}))$, when $\epsilon\rightarrow 0$. For
simplicity, we have written $f(z_{i})=f(z_{i},\boldsymbol{\theta},t_{i})$. If
we decide to expand equation 20, we obtain

$\frac{dz_{n}}{dz_{0}}=I+\sum_{i=0}^{n-1}\frac{\partial f(z_{i})}{\partial
z_{i}}\epsilon+\sum_{i=0}^{n-1}\sum_{j<i}\frac{\partial f(z_{i})}{\partial
z_{i}}\frac{\partial f(z_{j})}{\partial z_{j}}\epsilon^{2}+...+\frac{\partial
f(z_{0})}{\partial z_{0}}...\frac{\partial f(z_{n})}{\partial
z_{n}}\epsilon^{n}$ (21) $=I+S_{1}^{n}+S_{2}^{n}+...+S_{n+1}^{n},$ (22)

where

$S_{2}^{n}=\sum_{i=0}^{n-1}\sum_{j<i}\frac{\partial f(z_{i})}{\partial
z_{i}}\frac{\partial f(z_{j})}{\partial z_{j}}\epsilon^{2},\
S_{3}^{n}=\sum_{i=0}^{n-1}\sum_{j<i}\sum_{k<j}\frac{\partial
f(z_{i})}{\partial z_{i}}\frac{\partial f(z_{j})}{\partial
z_{j}}\frac{\partial f(z_{k})}{\partial z_{k}}\epsilon^{3},...$ (23)

Figure 10: Rectangle $(i,j)$ in discretized $[0,T]\times[0,T]$

Figure 11: Symmetry of $g_{2}^{n}$

Figure 12: Rectangles participating in $S_{2}^{n}$

Figure 13: Rectangles in $\bar{S}_{2}^{n}$

We will focus on the sum $S_{2}^{n}$ for a moment. Let us define

$g_{2}^{n}(x,y)=\frac{\partial f(z_{t_{i}})}{\partial z_{t_{i}}}\frac{\partial
f(z_{t_{j}})}{\partial z_{t_{j}}},\text{ for
}x\in[t_{i},t_{i+1}],y\in[{t_{j}},{t_{j+1}}].$ (24)

It is clear that $g_{2}^{n}$ is the discretization of

$g_{2}(t,u)=\frac{\partial f(z_{t})}{\partial z_{t}}\frac{\partial
f(z_{u})}{\partial z_{u}},\text{ for }t\in[0,T],u\in[0,T].$ (25)

We can notice that $S_{2}^{n}=\sum_{i=0}^{n-1}\sum_{j<i}\
g_{2}^{n}(t_{i},t_{j})\epsilon^{2}$ and that
$g_{2}^{n}(t_{i},t_{j})=g_{2}^{n}(t_{j},t_{i})$ is symmetric, but in
$S_{2}^{n}$, only takes values in the rectangles under the diagonal as
illustrated in Figure 13. If we decide to expand the sum $S_{2}^{n}$ on the
rectangles above the diagonal as in Figure 13 and denote it as
$\bar{S}_{2}^{n}=\sum_{i=0}^{n-1}\sum_{j\neq i}\
g_{2}^{n}(t_{i},t_{j})\epsilon^{2}$, then due to the symmetry of
$g_{2}^{n}(t_{i},t_{j})=g_{2}^{n}(t_{j},t_{i})$, we deduce that
$S_{2}^{n}=\frac{1}{2}\bar{S}_{2}^{n}$. The only rectangles missing in
$\bar{S}_{2}^{n}$, regarding the discretization of $[0,T]\times[0,T]$, are the
ones corresponding to the cases when $i=j$, which are the rectangles in the
diagonal. It is important to emphasize here that $S_{2}^{n}$ is the sum of
$C_{n}^{2}$ terms, while $\bar{S}_{2}^{n}$ is made out of
$V_{n}^{2}=2!C_{n}^{2}$ terms. This is especially apparent when we notice the
discretization of $[0,T]\times[0,T]$ is composed of $n^{2}$ rectangles, while
the diagonal is composed of $n$ rectangles, hence $\bar{S}_{2}^{n}$ is the sum
of $n^{2}-n=V_{n}^{2}$ terms. The ratio of the collective mass of the
rectangles in the diagonal with respect to the entire $[0,T]\times[0,T]$ goes
to zero as $n\rightarrow\infty$, hence we conclude that

$\bar{S}_{2}^{n}=\sum_{i=0}^{n-1}\sum_{j\neq i}\frac{\partial
f(z_{i})}{\partial z_{i}}\frac{\partial f(z_{j})}{\partial
z_{j}}\epsilon^{2}\rightarrow\int_{[0,T]\times[0,T]}g_{2}(t,u)d(t,u)=(\int_{0}^{T}\frac{\partial
f(z_{t})}{\partial z_{t}}dt)^{2},$ (26)

thus,

$S_{2}^{n}=\frac{1}{2!}\bar{S}_{2}^{n}\rightarrow\frac{1}{2!}(\int_{0}^{T}\frac{\partial
f(z_{t})}{\partial z_{t}}dt)^{2}.$ (27)

Similarly for

$S_{3}^{n}=\sum_{i=0}^{n-1}\sum_{j<i}\sum_{k<j}\frac{\partial
f(z_{i})}{\partial z_{i}}\frac{\partial f(z_{j})}{\partial
z_{j}}\frac{\partial f(z_{k})}{\partial z_{k}}\epsilon^{3},$

we can define

$g_{3}^{n}(x,y,z)=\frac{\partial f(z_{t_{i}})}{\partial
z_{t_{i}}}\frac{\partial f(z_{t_{j}})}{\partial z_{t_{j}}}\frac{\partial
f(z_{t_{k}})}{\partial z_{t_{k}}},\text{ for
}x\in[t_{i},t_{i+1}],y\in[{t_{j}},{t_{j+1}}],z\in[{t_{k}},{t_{k+1}}],$ (28)

and

$g_{3}(t,u,v)=\frac{\partial f(z_{t})}{\partial z_{t}}\frac{\partial
f(z_{u})}{\partial z_{u}}\frac{\partial f(z_{v})}{\partial z_{v}},\text{ for
}t\in[0,T],u\in[0,T],v\in[0,T].$ (29)

In this case:

$g_{3}^{n}(x,y,z)=g_{3}^{n}(x,y,z)=g_{3}^{n}(x,z,y)=g_{3}^{n}(y,x,z)=$
$=g_{3}^{n}(y,z,x)=g_{3}^{n}(z,x,y)=g_{3}^{n}(z,y,x),$ (30)

where each equality corresponds to one of the $6=3!$ permutations of
$(x,y,z)$. As before we can expand the domain of $S_{3}^{n}$, by adding the
rectangles in the discretization of $[0,T]\times[0,T]\times[0,T]$ such that
$i$ might be smaller than $j$, as well as that $j$ could be smaller than $k$.
We denote this expanded sum as $\bar{S}_{3}^{n}$, and by the symmetry of
$g_{3}^{n}$, we notice that $S_{3}^{n}=\frac{1}{3!}\bar{S}_{3}^{n}$. This
implies that the rectangles participating in $\bar{S}_{3}^{n}$, now cover most
of $[0,T]\times[0,T]\times[0,T]$, where the only exceptions are the ones when
$i=j$ or $j=k$ (or both). The number of rectangles participating in
$\bar{S}_{3}^{n}$ is $V_{n}^{3}=3!C_{n}^{3}=n^{3}-3n^{2}+2n$, and since the
number of all rectangles in $[0,T]\times[0,T]\times[0,T]$ is $n^{3}$, this
implies that $3n^{2}-2n$ rectangles are missing in sum $\bar{S}_{3}^{n}$ from
the cases when $i=j$ or $j=k$ (or both). The ratio of the collective mass of
such rectangles with respect to the entire $[0,T]\times[0,T]\times[0,T]$ goes
to zero as $n\rightarrow\infty$, hence as before:

$\bar{S}_{3}^{n}=\sum_{i=0}^{n-1}\sum_{j\neq i}\sum_{i\neq k\neq
j}\frac{\partial f(z_{i})}{\partial z_{i}}\frac{\partial f(z_{j})}{\partial
z_{j}}\frac{\partial f(z_{k})}{\partial
z_{k}}\epsilon^{3}\rightarrow\int_{[0,T]\times[0,T]\times[0,T]}g_{3}(t,u,v)d(t,u,v)$
(31) $=(\int_{0}^{T}\frac{\partial f(z_{t})}{\partial z_{t}}dt)^{3},$ (32)

implying

$S_{3}^{n}=\frac{1}{3!}\bar{S}_{3}^{n}\rightarrow\frac{1}{3!}(\int_{0}^{T}\frac{\partial
f(z_{t})}{\partial z_{t}}dt)^{3}.$ (33)

In a similar fashion we can prove that
$S_{k}^{n}\rightarrow\frac{1}{k!}(\int_{0}^{T}\frac{\partial
f(z_{t})}{\partial z_{t}}dt)^{k}$. Thus we conclude that:

$\frac{dz(T)}{dz(0)}=\frac{1}{0!}I+\frac{1}{1!}\int_{0}^{T}\frac{\partial
f(z_{t})}{\partial z_{t}}dt+\frac{1}{2!}(\int_{0}^{T}\frac{\partial
f(z_{t})}{\partial z_{t}}dt)^{2}+\frac{1}{3!}(\int_{0}^{T}\frac{\partial
f(z_{t})}{\partial z_{t}}dt)^{3}+...$ (34)
$\frac{dz(T)}{dz(0)}=e^{\int_{0}^{T}\frac{\partial f(z_{t})}{\partial
z_{t}}dt}.$ (35)

Unfortunately, this result does not hold when the dimensionality of
$\boldsymbol{z}(t)$ is larger than one. Indeed, in this case, $\frac{\partial
f(\boldsymbol{z}_{t_{i}})}{\partial\boldsymbol{z}_{t_{i}}}$ is a matrix, hence
$g_{2}^{n}(x,y)=\frac{\partial
f(\boldsymbol{z}_{t_{i}})}{\partial\boldsymbol{z}_{t_{i}}}\frac{\partial
f(\boldsymbol{z}_{t_{j}})}{\partial\boldsymbol{z}_{t_{j}}}$ is not necessarily
symmetric, as the commutator $[\frac{\partial
f(\boldsymbol{z}_{t_{i}})}{\partial\boldsymbol{z}_{t_{i}}},\frac{\partial
f(\boldsymbol{z}_{t_{j}})}{\partial\boldsymbol{z}_{t_{j}}}]$ is not
necessarily zero. For this reason, inspired by the previous result we try a
different approach. Indeed, we can see from Equation 35 that
$\frac{dz(T)}{dz(0)}$ is the solution of the following ODE:

$\frac{d(\frac{dz(t)}{dz(0)})}{dt}=\frac{\partial f(z(t))}{\partial
z(t)}\frac{dz(t)}{dz(0)}$ (36)

as the initial condition $\frac{dz(t=0)}{dz(0)}=I$. Hence we wish to prove
that $\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}$ satisfies the same ODE in
higher dimensions as well.

We notice that we can write:

$\boldsymbol{z}(t)=\boldsymbol{z}(0)+\int_{0}^{t}f(\boldsymbol{z}(\tau))d\tau=g(\boldsymbol{z}(0),t),$
(37)

therefore,

$\frac{d(\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)})}{dt}=\frac{\partial^{2}g(\boldsymbol{z}(0),t)}{\partial
t\partial\boldsymbol{z}(0)}=\frac{\partial^{2}g(\boldsymbol{z}(0),t)}{\partial\boldsymbol{z}(0)\partial
t}=$
$=\frac{d(\frac{dg(\boldsymbol{z}(0),t)}{dt})}{d\boldsymbol{z}(0)}=\frac{df(\boldsymbol{z}(t))}{d\boldsymbol{z}(0)}=\frac{df(\boldsymbol{z}(t))}{d\boldsymbol{z}(t)}\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}.$
(38)

We pause for a moment, in order to highlight the similarity of expression 38
and the forward sensitivity:

$\frac{d(\frac{d\boldsymbol{z}(t)}{d\boldsymbol{\theta}})}{dt}=\frac{df(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{d\boldsymbol{\theta}}=\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}\frac{d\boldsymbol{z}(t)}{d\boldsymbol{\theta}}+\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}.$
(39)

Now from Expression 38, we infer that
$\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}$ is the solution of the
following linear ODE:

$\frac{d(\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)})}{dt}=\frac{\partial
f(\boldsymbol{z}(t))}{\partial\boldsymbol{z}(t)}\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}.$
(40)

If we write $\boldsymbol{Y}(t)=\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}$,
then the equation above becomes:

$\frac{d\boldsymbol{Y}(t)}{dt}=\frac{\partial
f(\boldsymbol{z}(t))}{\partial\boldsymbol{z}(t)}\boldsymbol{Y}(t).$ (41)

The general solution of first-order homogeneous linear ODEs is given in [16,
17], and in our case can be written as follows:

$\frac{d\boldsymbol{\boldsymbol{z}}(t)}{d\boldsymbol{\boldsymbol{z}}(0)}=e^{\boldsymbol{\Omega}(t)}\frac{d\boldsymbol{z}(0)}{d\boldsymbol{z}(0)}=e^{\boldsymbol{\Omega}(t)},\text{
for\ \ }\boldsymbol{\Omega}(t)=\sum_{k=1}^{\infty}\boldsymbol{\Omega}_{k}(t),$
(42)

where

$\displaystyle\boldsymbol{\Omega}_{1}(t)$
$\displaystyle=\int_{0}^{t}\boldsymbol{A}(t_{1})\,dt_{1},$ (43)
$\displaystyle\boldsymbol{\Omega}_{2}(t)$
$\displaystyle=\frac{1}{2}\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}\,[\boldsymbol{A}(t_{1}),\boldsymbol{A}(t_{2})],$
(44) $\displaystyle\boldsymbol{\Omega}_{3}(t)$ $\displaystyle=\ $ (45)
$=\frac{1}{6}\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}\int_{0}^{t_{2}}dt_{3}\,\Bigl{(}\big{[}\boldsymbol{A}(t_{1}),[\boldsymbol{A}(t_{2}),\boldsymbol{A}(t_{3})]\big{]}+\big{[}\boldsymbol{A}(t_{3}),[\boldsymbol{A}(t_{2}),\boldsymbol{A}(t_{1})]\big{]}\Bigr{)},$
(46)

and so on for $k>3$, where:

$\displaystyle\boldsymbol{A}(t)=\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{\boldsymbol{z}}(t)}.$
(47)

We notice that if $z(t)$ is one dimensional, then this result agrees with the
one in the previous approach. To conclude, we have proven the following
theorem:

###### Theorem 0.A.1

Let
$\frac{d\boldsymbol{z}(t)}{dt}=f(\boldsymbol{z}(t),\boldsymbol{\theta},t)$,
where $f$ is continuous in $t$ and Lipschitz continuous in
$\boldsymbol{z}(t)$. Then the following holds:

$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}=e^{\int_{0}^{T}\frac{\partial
f(\boldsymbol{z}_{t})}{\partial\boldsymbol{z}_{t}}dt+\boldsymbol{\Omega}_{2}(T)+....},$
(48)

where $\boldsymbol{\Omega}_{k>1}(t)$ are the terms of the Magnus series.

###### Proof

Since $f$ is continuous in $t$ and Lipschitz continuous in $\boldsymbol{z}$
then due to Picard–Lindelöf theorem, $\boldsymbol{z}(T)$ exists and is unique.
Furthermore, since $f$ is Lipschitz continuous in $\boldsymbol{z}(t)$, then
$\frac{\partial f(\boldsymbol{z}_{t})}{\partial\boldsymbol{z}_{t}}$ exists
almost everywhere. Based on the previous derivations we reach the desired
conclusion.

##  Appendix 0.B The Continuous Generalization of the Total Derivative
Decomposition

In Appendix 0.A, we assumed that
$\boldsymbol{z}(t)=\boldsymbol{z}(0)+\int_{0}^{t}f(\boldsymbol{z}(\tau),\boldsymbol{\theta},\tau)d\tau$,
which was the the limit of $\epsilon\rightarrow 0$ of the previous iterative
definition of $\boldsymbol{z}_{i+1}=\boldsymbol{z}_{i}+\epsilon
f(\boldsymbol{z}_{i},\boldsymbol{\theta},t_{i})$. Now, we assume that the set
of parameters in $f$ is different at each discrete time, that is
$\boldsymbol{z}_{i+1}=\boldsymbol{z}_{i}+\epsilon
f_{i}(\boldsymbol{z}_{i},\boldsymbol{\theta}_{i},t_{i})$, for independent sets
$\boldsymbol{\theta}_{i}$. First, we notice that

$\boldsymbol{z}_{n}=\boldsymbol{z}_{n}(\boldsymbol{z}(t_{0}),\boldsymbol{\theta}(t_{0}),\boldsymbol{\theta}(t_{1}),...,\boldsymbol{\theta}(t_{n-1}))=\boldsymbol{z}_{n}(\boldsymbol{z}(t_{1}),\boldsymbol{\theta}(t_{1}),\boldsymbol{\theta}(t_{2}),...,\boldsymbol{\theta}(t_{n-1}))=$
$...\boldsymbol{z}_{n}(\boldsymbol{z}(t_{k+1}),\boldsymbol{\theta}(t_{k+1}),\boldsymbol{\theta}(t_{k+2}),...,\boldsymbol{\theta}(t_{n-1}))=...$
$...=\boldsymbol{z}_{n}(\boldsymbol{z}(t_{n-1}),\boldsymbol{\theta}(t_{n-1}))=\boldsymbol{z}(t_{n}=T-\epsilon).$

This can be seen from Figure 14.

Figure 14: The dependencies in the discretisation of
$\frac{\boldsymbol{z}(t)}{t}=f(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)$.
Variable $\boldsymbol{z}_{t_{k+1}}$ completely absorbs the contribution of
$\boldsymbol{\theta}_{t_{k}}$ to $\boldsymbol{z}_{T}$, hence
$\boldsymbol{z}_{T}$ is a function of $\boldsymbol{z}_{t_{k+1}}$ and the
following sets of parameters
$\boldsymbol{\theta}_{t_{k+1}},\boldsymbol{\theta}_{t_{k+2}},...,\boldsymbol{\theta}_{t_{n}}$.

Thus

$\frac{d\boldsymbol{z}_{n}}{d\boldsymbol{\theta}(t_{k})}=\frac{d\boldsymbol{z}_{n}(\boldsymbol{z}(t_{k}+\epsilon),\boldsymbol{\theta}(t_{k}+\epsilon),\boldsymbol{\theta}(t_{k+2}),...,\boldsymbol{\theta}(t_{n}))}{d\boldsymbol{\theta}(t_{k})}$
(49)

$\frac{d\boldsymbol{z}_{n}}{d\boldsymbol{\theta}(t_{k})}=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}\frac{d\boldsymbol{z}(t_{k}+\epsilon)}{d\boldsymbol{\theta}(t_{k})}+0+0+...+0=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}\frac{\partial\boldsymbol{z}(t_{k}+\epsilon)}{\partial\boldsymbol{\theta}(t_{k})}.$
(50)

We now focus on analysing
$\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}$
and
$\frac{\partial\boldsymbol{z}(t_{k}+\epsilon)}{\partial\boldsymbol{\theta}(t_{k})}$.
First we notice that from:

$\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k})}=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}\frac{\partial\boldsymbol{z}(t_{k}+\epsilon)}{\partial\boldsymbol{z}(t_{k})}=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}(I+\epsilon\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{z}(t_{k})}+O(\epsilon^{2}))$
(51)

we get

$\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k})}-\epsilon\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{z}(t_{k})}+O(\epsilon^{2}).$
(52)

Regarding
$\frac{\partial\boldsymbol{z}(t_{k}+\epsilon)}{\partial\boldsymbol{\theta}(t_{k})}$
the following holds:

$\frac{\partial\boldsymbol{z}(t_{k}+\epsilon)}{\partial\boldsymbol{\theta}(t_{k})}=\frac{\partial(\boldsymbol{z}(t_{k})+f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))\epsilon+O(\epsilon^{2}))}{\partial\boldsymbol{\theta}(t_{k})}=0+\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{\theta}(t_{k})}\epsilon+O(\epsilon^{2}).$
(53)

After combining both Equation 52 and Equation 53 in expression 50, we deduce
that:

$\frac{d\boldsymbol{z}_{n}}{d\boldsymbol{\theta}(t_{k})}=$
$\big{(}\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k})}-\epsilon\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k}+\epsilon)}\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{z}(t_{k})}+O(\epsilon^{2})\big{)}\big{(}\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{\theta}(t_{k})}\epsilon+O(\epsilon^{2})\big{)},$
(54)

thus

$\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{\theta}(t_{k})}=\frac{d\boldsymbol{z}_{n}}{d\boldsymbol{\theta}(t_{k})}=\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k})}\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{\theta}(t_{k})}\epsilon+O(\epsilon^{2})$
(55)

Now assuming that $\boldsymbol{\theta}(t_{k})=g(t_{k},\boldsymbol{\theta})$,
we can write the total derivative of $\boldsymbol{z}_{n}$ with respect to
parameters $\boldsymbol{\theta}$ :

$\frac{d\boldsymbol{z}_{n}}{d\boldsymbol{\theta}}=\sum_{k=0}^{n}\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{\theta}(t_{k})}\frac{d\boldsymbol{\theta}(t_{k})}{d\boldsymbol{\theta}}=\sum_{k=0}^{n}\frac{\partial\boldsymbol{z}_{n}}{\partial\boldsymbol{z}(t_{k})}\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{\theta}(t_{k})}\frac{d\boldsymbol{\theta}(t_{k})}{d\boldsymbol{\theta}}\epsilon+O(\epsilon^{2})$
(56)

hence taking $\epsilon\rightarrow 0$, we conclude that

$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta}(t))}{\partial\boldsymbol{\theta}(t)}\frac{\partial\boldsymbol{\theta}(t)}{\partial\boldsymbol{\theta}}dt.$
(57)

We can see that we are integrating the infinitesimal contributions of
parameters $\boldsymbol{\theta}$, at each time $t$. In case that
$\boldsymbol{\theta}(t)=g(t,\boldsymbol{\theta})=\boldsymbol{\theta}$, then we
have

$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}\frac{\partial\boldsymbol{\theta}}{\partial\boldsymbol{\theta}}dt=\int_{0}^{T}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}Idt=$
$=\int_{0}^{T}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}dt=-\int_{T}^{0}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}dt.$
(58)

###### Theorem 0.B.1

Let $\frac{dz(t)}{dt}=f(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)$, where
$f$ is continuous in $t$ and Lipschitz continuous in $\boldsymbol{z}(t)$ and
$\boldsymbol{\theta}(t)$. Then the following holds:

$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial\boldsymbol{z}(T)}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)}{\partial\boldsymbol{\theta}(t)}\frac{\partial\boldsymbol{\theta}(t)}{\partial\boldsymbol{\theta}}dt.$
(59)

###### Proof

As before, since $f$ is continuous in $t$ and Lipschitz continuous in $z$ then
due to Picard–Lindelöf theorem, $z(T)$ exists and is unique. From Theorem
0.A.1, we can establish the existence of
$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(t)}$. Furthermore, since $f$ is
Lipschitz continuous in $\theta(t)$, then $\frac{\partial
f(\boldsymbol{\theta}(t))}{\partial\boldsymbol{\theta}(t)}$ exists almost
everywhere. Based on the previous derivations we reach the desired conclusion.

##  Appendix 0.C Generalization of Continuous Backpropagation into Piecewise
Continuous Backpropagation

We define $\boldsymbol{z}(t)$ as before, with the only difference being that
its derivative is discontinuous at a point (say $\frac{T}{2}$):

$\displaystyle\boldsymbol{z}(t)=\begin{cases}\boldsymbol{z}(0)+\int_{0}^{t}f(\boldsymbol{z}(\tau),\boldsymbol{\theta}(\tau))d\tau,\
t\in[0,\frac{T}{2}]\\\
\boldsymbol{z}(0)+\int_{0}^{\frac{T}{2}}f(\boldsymbol{z}(\tau),\boldsymbol{\theta}(\tau))d\tau+\int_{\frac{T}{2}}^{t}g(\boldsymbol{z}(\tau),\boldsymbol{\phi}(\tau))d\tau,\
t\in(\frac{T}{2},T]\end{cases}$ (60)

As in [4], we can get

$\displaystyle\frac{d(\frac{\partial
L}{\partial\boldsymbol{z}(t)})}{dt}=\begin{cases}-\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta}(t))}{\partial\boldsymbol{z}(t)},\
t\in[0,\frac{T}{2}]\\\ -\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{\phi}(t))}{\partial\boldsymbol{z}(t)},\
t\in(\frac{T}{2},T],\end{cases}$ (61)

thus

$\displaystyle\frac{\partial
L}{\partial\boldsymbol{z}(t)}=\begin{cases}\frac{\partial
L}{\partial\boldsymbol{z}(T)}-\int_{T}^{t}\frac{\partial
L}{\partial\boldsymbol{z}(\tau)}\frac{\partial
g(\boldsymbol{z}(\tau),\boldsymbol{\phi}(\tau))}{\partial\boldsymbol{z}(\tau)}d\tau,\
t\in(\frac{T}{2},T]\\\ \frac{\partial
L}{\partial\boldsymbol{z}(T)}-\int_{\frac{T}{2}}^{t}\frac{\partial
L}{\partial\boldsymbol{z}(\tau)}\frac{\partial
f(\boldsymbol{z}(\tau),\boldsymbol{\theta}(\tau))}{\partial\boldsymbol{z}(\tau)}d\tau-\int_{T}^{\frac{T}{2}}\frac{\partial
L}{\partial\boldsymbol{z}(\tau)}\frac{\partial
g(\boldsymbol{z}(\tau)\boldsymbol{\phi}(\tau))}{\partial\boldsymbol{z}(\tau)}d\tau,\
t\in(0,\frac{T}{2}].\end{cases}$ (62)

The approach developed in Appendix 0.B can be used to generalize continuous
backpropagation into piecewise continuous backpropagation. Indeed, identically
as before, we have the first order appoximation:

$\displaystyle\frac{dL_{\epsilon}}{d\boldsymbol{\theta}}=\sum_{k=0}^{n}\frac{\partial
L_{\epsilon}}{\partial\boldsymbol{z}(t_{k})}\frac{\partial
f(\boldsymbol{z}(t_{k}),\boldsymbol{\theta}(t_{k}))}{\partial\boldsymbol{\theta}(t_{k})}\frac{\partial\boldsymbol{\theta}(t_{k})}{\partial\boldsymbol{\theta}}\epsilon+\sum_{k=0}^{n}\frac{\partial
L_{\epsilon}}{\partial\boldsymbol{z}(t_{k})}\frac{\partial
g(\boldsymbol{z}(t_{k}),\boldsymbol{\phi}(t_{k}))}{\partial\boldsymbol{\phi}(t_{k})}\frac{\partial\boldsymbol{\phi}(t_{k})}{\partial\boldsymbol{\theta}}\epsilon$
(63)

Taking the limit we get:

$\frac{dL}{d\boldsymbol{\theta}}=\int_{0}^{\frac{T}{2}}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta}(t))}{\partial\boldsymbol{\theta}(t)}\frac{\partial\boldsymbol{\theta}(t)}{\partial\boldsymbol{\theta}}dt+\int_{\frac{T}{2}}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{\phi}(t))}{\partial\boldsymbol{\theta}(t)}\frac{\partial\boldsymbol{\phi}(t)}{\partial\boldsymbol{\theta}}dt$
(64)

and in the same manner we get:

$\frac{dL}{d\boldsymbol{\phi}}=\int_{0}^{\frac{T}{2}}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta}(t))}{\partial\boldsymbol{\theta}(t)}\frac{\partial\boldsymbol{\theta}(t)}{\partial\boldsymbol{\phi}}dt+\int_{\frac{T}{2}}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{\phi}(t))}{\partial\boldsymbol{\theta}(t)}\frac{\partial\boldsymbol{\phi}(t)}{\partial\boldsymbol{\phi}}dt,$
(65)

If we set $\boldsymbol{\theta}(t)=\boldsymbol{\theta}$ and
$\boldsymbol{\phi}(t)=\boldsymbol{\phi}$, then we get:

$\frac{dL}{d\boldsymbol{\theta}}=\int_{0}^{\frac{T}{2}}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta})}{\partial\boldsymbol{\theta}}dt,$
(66)

and

$\frac{dL}{d\boldsymbol{\phi}}=\int_{\frac{T}{2}}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{\phi})}{\partial\boldsymbol{\phi}}dt.$ (67)

##  Appendix 0.D Additional Generated Samples

Below in Figure 15, Figure 16 and Figure 17 , are presented additional
examples of generated samples of MNIST, CIFAR-10 and CelebA($32\times 32$) by
AFFJORD, respectively.

Figure 15: Generated samples of MNIST by AFFJORD Figure 16: Generated samples
of CIFAR-10 by AFFJORD Figure 17: Generated samples of CelebA($32\times 32$)
by AFFJORD

In addition in Figure 18, Figure 19 and Figure 20, are presented additional
examples of generated samples of MNIST, CIFAR-10 and CelebA($32\times 32$) by
FFJORD, respectively.

Figure 18: Generated samples of MNIST by FFJORD Figure 19: Generated samples
of CIFAR-10 by FFJORD Figure 20: Generated samples of CelebA ($32\times 32$)
by FFJORD

##  Appendix 0.E Deriving the Instantaneous Change of Variable via the Cable
Rule

The trace of a commutator
$[\boldsymbol{X},\boldsymbol{Y}]=\boldsymbol{X}\boldsymbol{Y}-\boldsymbol{Y}\boldsymbol{X}$
is always zero, hence we prove below that for all terms
$\boldsymbol{\Omega}_{k>1}(t)$ given in Equations 43 in Appendix 0.A, we have
$tr(\boldsymbol{\Omega}_{k}(t))=0$.

Indeed, since the trace and the integral are interchangeable we have:

$tr(\boldsymbol{\Omega}_{2}(t))=tr\frac{1}{2}\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}\,[\boldsymbol{A}(t_{1}),\boldsymbol{A}(t_{2})]$
$=\frac{1}{2}\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}\,tr[\boldsymbol{A}(t_{1}),\boldsymbol{A}(t_{2})]=0.$
(68)

$tr(\boldsymbol{\Omega}_{3}(t))=\frac{1}{6}\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}\int_{0}^{t_{2}}dt_{3}\,\Bigl{(}tr\big{[}\boldsymbol{A}(t_{1}),[\boldsymbol{A}(t_{2}),\boldsymbol{A}(t_{3})]\big{]}+$
$+tr\big{[}\boldsymbol{A}(t_{3}),[\boldsymbol{A}(t_{2}),\boldsymbol{A}(t_{1})]\big{]}\Bigr{)}$
(69) $=\frac{1}{6}\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}(0+0)dt_{2}=0,$ (70)

and so on for $k>3$.

The only exception is the case when $k=1$:

$\displaystyle
tr(\boldsymbol{\Omega}_{1}(t))=tr\int_{0}^{t}\boldsymbol{A}(t_{1})dt_{1}=\int_{0}^{t}tr\frac{\partial
f(\boldsymbol{z}(t_{1}),t_{1},\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{1})}dt_{1}.$
(71)

Finally, using Jacobi’s formula, we conclude that:

$\log{|\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}|}=\log{|e^{\boldsymbol{\Omega}(T)}|}=\log{e^{tr(\boldsymbol{\Omega}(T))}}={\int_{0}^{T}tr\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}}dt.$
(72)

##  Appendix 0.F Simplifications in Section 3.1

First we notice that if $z^{*}(0)$ does not depend on $z(0)$, then

$\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}=\begin{bmatrix}I,0\end{bmatrix}e^{\begin{bmatrix}\int_{0}^{T}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}dt&\int_{0}^{T}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t)}dt\\\
\int_{0}^{T}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}(t)}dt&\int_{0}^{T}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t)}dt\end{bmatrix}+\boldsymbol{\Omega}_{2}(T)+...}\
\ \begin{bmatrix}I\\\
\frac{d\boldsymbol{z}^{*}(0)}{d\boldsymbol{z}(0)}\end{bmatrix},$ (73)

becomes

$|\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}|=|\begin{bmatrix}I,0\end{bmatrix}e^{\begin{bmatrix}\int_{0}^{T}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}dt&\int_{0}^{T}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t)}dt\\\
\int_{0}^{T}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}(t)}dt&\int_{0}^{T}\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t)}dt\end{bmatrix}+\boldsymbol{\Omega}_{2}(T)+...}\
\ \begin{bmatrix}I\\\ 0\end{bmatrix}|.$ (74)

On the other hand, if the augmented dimensions do not depend on the main
dimensions then $\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}(t)}=0$.
This implies that

$\boldsymbol{A}(t)=\begin{bmatrix}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}&\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t)}\\\
\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}(t)}&\frac{\partial
g(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t)}\end{bmatrix}\rightarrow\begin{bmatrix}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}&\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t)}\\\
0&\frac{\partial
g(\boldsymbol{z^{*}}(t),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t)}\end{bmatrix}.$
(75)

Hence,

$\boldsymbol{\Omega}_{1}(t)=\int_{0}^{t}dt_{1}\begin{bmatrix}\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{1})}&\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t_{1})}\\\
0&\frac{\partial
g(\boldsymbol{z^{*}}(t_{1}),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t_{1})}\end{bmatrix}=\begin{bmatrix}\int_{0}^{t}dt_{1}\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{1})}&\boldsymbol{B}_{1}(t)\\\
0&\boldsymbol{D}_{1}(t)\end{bmatrix}$ (76)

$\ \ \ \ \ \ \boldsymbol{\Omega}_{2}(t)=$

$\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}\,\begin{bmatrix}\frac{\partial
f(\boldsymbol{z}(t_{2}),\boldsymbol{z}^{*}(t_{2}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{2})}&\frac{\partial
f(\boldsymbol{z}(t_{2}),\boldsymbol{z}^{*}(t_{2}),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t_{2})}\\\
0&\frac{\partial
g(\boldsymbol{z^{*}}(t_{2}),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t_{2})}\end{bmatrix}\begin{bmatrix}\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{1})}&\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t_{1})}\\\
0&\frac{\partial
g(\boldsymbol{z^{*}}(t_{1}),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t_{1})}\end{bmatrix}$
$-\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}\,\begin{bmatrix}\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{1})}&\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t_{1})}\\\
0&\frac{\partial
g(\boldsymbol{z^{*}}(t_{1}),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t_{1})}\end{bmatrix}\begin{bmatrix}\frac{\partial
f(\boldsymbol{z}(t_{2}),\boldsymbol{z}^{*}(t_{2}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{2})}&\frac{\partial
f(\boldsymbol{z}(t_{2}),\boldsymbol{z}^{*}(t_{2}),\boldsymbol{\theta})}{\partial\boldsymbol{z}^{*}(t_{2})}\\\
0&\frac{\partial
g(\boldsymbol{z^{*}}(t_{2}),\boldsymbol{\phi})}{\partial\boldsymbol{z}^{*}(t_{2})}\end{bmatrix}$
$=\begin{bmatrix}\int_{0}^{t}dt_{1}\int_{0}^{t_{1}}dt_{2}[\frac{\partial
f(\boldsymbol{z}(t_{1}),\boldsymbol{z}^{*}(t_{1}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{1})},\frac{\partial
f(\boldsymbol{z}(t_{2}),\boldsymbol{z}^{*}(t_{2}),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t_{2})}]&\boldsymbol{B}_{2}(t)\\\
0&\boldsymbol{D}_{2}(t)\end{bmatrix}=\begin{bmatrix}\boldsymbol{\Omega}_{2}^{[z]}(t)&\boldsymbol{B}_{2}(t)\\\
0&\boldsymbol{D}_{2}(t)\end{bmatrix}.$

In this way we can prove that

$\boldsymbol{\Omega}(t)=\sum_{k=1}^{\infty}\boldsymbol{\Omega}(t)_{k}=\begin{bmatrix}\sum_{k=1}^{\infty}\boldsymbol{\Omega}^{[z]}_{k}(t)&\sum_{k=1}^{\infty}\boldsymbol{B}_{k}(t)\\\
0&\sum_{k=1}^{\infty}\boldsymbol{D}_{k}(t)\end{bmatrix}=\begin{bmatrix}\boldsymbol{\Omega}^{[z]}(t)&\boldsymbol{B}(t)\\\
0&\boldsymbol{D}(t)\end{bmatrix}.$

Considering that the exponential of a matrix whose lower left block is zero,
will still have a zero lower left block, we have:

$\frac{dz(t)}{dz(0)}=e^{\boldsymbol{\Omega}(t)}=\begin{bmatrix}e^{\boldsymbol{\Omega}^{[z]}(t)}&\bar{\boldsymbol{B}}(t)\\\
0&\bar{\boldsymbol{D}}(t)\end{bmatrix}.$

Finally,

$|\frac{d\boldsymbol{z}(T)}{d\boldsymbol{z}(0)}|=|\begin{bmatrix}I,0\end{bmatrix}\begin{bmatrix}e^{\boldsymbol{\Omega}^{[z]}(t)}&\bar{\boldsymbol{B}}(t)\\\
0&\bar{\boldsymbol{D}}(t)\end{bmatrix}\ \ \begin{bmatrix}I\\\
0\end{bmatrix}|=|e^{\boldsymbol{\Omega}^{[z]}(t)}|=e^{\int_{0}^{T}tr\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{z}^{*}(t),\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}dt+0+...0+...}.$
(77)

##  Appendix 0.G Cable Rule Derived via Equation 5

Differentiating Equation 5 in Appendix 0.A, we get:

$\frac{d\big{(}\frac{dL}{d\boldsymbol{z}(t)}\big{)}}{dt}=-\frac{dL}{d\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}.$ (78)

Choosing $L$ to be $\boldsymbol{z}(0)$, we have

$\frac{d\big{(}\frac{d\boldsymbol{z}(0)}{d\boldsymbol{z}(t)}\big{)}}{dt}=-\frac{d\boldsymbol{z}(0)}{d\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}.$ (79)

We define $\boldsymbol{A}(t):=\frac{d\boldsymbol{z}(0)}{d\boldsymbol{z}(t)}$,
and
$\boldsymbol{B}(t):=\boldsymbol{A}(t)^{-1}=\frac{d\boldsymbol{z}(t)}{d\boldsymbol{z}(0)}$,
so that the equation above becomes

$\frac{d\boldsymbol{A}(t)}{dt}=-\boldsymbol{A}(t)\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)},$ (80)

thus

$-\boldsymbol{B}(t)\frac{d\boldsymbol{A}(t)}{dt}=\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}.$ (81)

Then from

$\boldsymbol{B}(t)\boldsymbol{A}(t)=I,$ (82)

we have

$-\boldsymbol{B}(t)\frac{d\boldsymbol{A}(t)}{dt}=\frac{d\boldsymbol{B}(t)}{dt}\boldsymbol{A}(t),$
(83)

thus, using this expression in Equation 81 we get

$\frac{d\boldsymbol{B}(t)}{dt}\boldsymbol{A}(t)=\frac{\partial
f(\boldsymbol{z}(t),t,\boldsymbol{\theta})}{\partial\boldsymbol{z}(t)}.$ (84)

Finally, we get the desired result by multiplying both sides from the right
with $\boldsymbol{B}(t)$ .

##  Appendix 0.H Equivalence Between Continuous Total Derivative Decomposition
and the Continuous Backpropagation

In Section 3.3. we derived the expression of continuous backpropagation from
the continuous total derivative decomposition. However, we can also derive the
the formula of the continuous total derivative decomposition (Equation 13)
from continuous backpropagation (Equation 14). Indeed, for a function
$f=f(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)$, where
$\boldsymbol{\theta}(t)=g(\boldsymbol{\theta},t)$, we can see $f$ as
$h(\boldsymbol{z}(t),\boldsymbol{\theta},t)=f(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)$.
Hence,

$\frac{\partial L}{\partial\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{dh(\boldsymbol{z}(t),\boldsymbol{\theta},t)}{d\boldsymbol{\theta}}dt=\int_{0}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{df(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)}{d\boldsymbol{\theta}}dt,$
(85)

therefore

$\frac{\partial L}{\partial\boldsymbol{\theta}}=\int_{0}^{T}\frac{\partial
L}{\partial\boldsymbol{z}(t)}\frac{\partial
f(\boldsymbol{z}(t),\boldsymbol{\theta}(t),t)}{\partial\boldsymbol{\theta}(t)}\frac{d\boldsymbol{\theta}(t)}{d\boldsymbol{\theta}}dt.$
(86)

Setting $L=\boldsymbol{z}(T)$, gives the desired result.

##  Appendix 0.I Additional Details About the Experiments

As mentioned in the main paper, we optimized the architecture of FFJORD first
and fine-tuned its hyper-parameters.

This architecture remains unchanged in the AFFJORD model for fair comparison,
and we only fine-tuned the augmented structure and dimension in addition.
Indeed, as it can be seen, the results we report for FFJORD (0.96 MNIST, 3.37
CIFAR) are better than those reported on the original paper (0.99 MNIST, 3.40
CIFAR). The aforementioned improvements were a result of reducing the number
of parameters during our tuning process, by reducing the number of CNF blocks
from 2 to 1. The total number of parameters for different architectures of
FFJORD and AFFJORD can be found on Table 2. It should be emphasized that in
the hypernet architecture, the parameters of AFFJORD are contained inside the
parameters in the main architecture of FFJORD, so a bigger model is not being
used, simply the flexibility of evolution of the main parameters in time is
being increased.

Table 2: Number of Total Parameters of FFJORD and AFFJORD (Concat and Hypernet
Architecture) for MNIST and CIFAR-10. The Multiscale Architecture Is Used in
All Cases.
Dataset | Concat | Hypernet
---|---|---
FFJORD | AFFJORD | FFJORD | AFFJORD
MNIST | 400323 | 417629 | 783019 | 4556169
CIFAR10 | 679441 | 820815 | 1332361 | 8976115
CelebA ($32\times 32$) | 679441 | 820815 | 1332361 | 8976115

More generally, our experiments show that augmented architectures generally
improve the performance of the standard non-augmented FFJORD counterparts,
independently from the baseline architecture used. Furthermore, the farther
the performance of FFJORD is from being optimal, the greater are the
improvements when using AFFJORD.


# Zero-field and Larmor spinor precessions
in a neutron polarimeter experiment

S. Sponara J. Kleppa G. Badureka Y. Hasegawaa,b<EMAIL_ADDRESS>aAtominstitut der Österreichischen Universitäten, Stadionallee 2, A-1020
Vienna, Austria
bPRESTO, Japan Science and Technology Agency (JST), Kawaguchi, Saitama
332-0012, Japan Corresponding author Phone +43-1-58801-14190; Fax
+43-1-58801-14199.

###### Abstract

We present a neutron polarimetric experiment where two kinds of spinor
precessions are observed: one is induced by different total energy of neutrons
(zero-field precession) and the other is induced by a stationary guide field
(Larmor precession). A characteristic of the former is the dependence of the
energy-difference, which is in practice tuned by the frequency of the
interacting oscillating magnetic field $\omega_{\textrm{\scriptsize{R}}}$. In
contrast the latter completely depends on the strength of the guide field,
namely Larmor frequency $\omega_{\textrm{\scriptsize{L}}}$. Our neutron-
polarimetric experiment exhibits individual tuning as well as specific
properties of each spinor precession, which assures the use of both spin
precessions for multi-entangled spinor manipulation.

###### keywords:

neutron , spin precession , Larmor precession , zero-field precession ,
polarimeter

, , , and

In the last years great effort has been devoted to the development of quantum
information and communication technology [1]. Among them creation and
manipulation of entanglement is a key issue and experimental realizations with
the use of photons, atoms, ions, nuclear-magnet, superconducting qubits etc.
have been reported [2]. There, the entanglement between spatially separated
systems are mainly used to exploit non-local property of quantum systems. In
contrast, investigations of two-qubit entanglement with neutrons, one of the
most useful quantum system to be utilized for studies of quantum mechanical
phenomena [3], have been performed with the use of their degrees of freedom,
namely spin and path, in a single particle: a violation of a Bell-like
inequality [4], contextual nature of quantum theory [5], and a full
tomographic analysis of Bell-states [6] were carried out.

For further development towards multi-entanglement, it is inevitable to
establish manipulation of other degrees of freedom. It is know that the total
energy of neutrons can be manipulated with the use of interactions between a
neutron’s magnetic moment and a time-dependent oscillating magnetic field [7,
8, 9]. Recently, we have accomplished an experiment to exhibit a coherent
energy manipulation of neutrons with the use of radio-frequency (RF)
oscillating magnetic field [10], enabling neutrons to be applied for
investigations of multi-entanglement. Thus, we find it significant to
demonstrate explicitly a phase manipulation for different energy eigenstates
which will be utilized in forthcoming experiments such as a demonstration of
triple-entanglment with a GHZ-state [11]. In the neutron scattering community,
a spin precession, alternatively a phase shift, due to energy difference is
known as zero-field precession [12, 13] and utilized for zero-field spin-echo
spectroscopy [14, 15, 16, 17].

When a neutron is exposed to a stationary magnetic field, the motion of its
spin vector is thoroughly described by the Bloch-equation exhibiting Larmor
precession:

$\frac{d\vec{S}}{dt}=\gamma\vec{S}\times\vec{B}$ (1)

where $\gamma$ is the gyromagnetic ratio given by $\mu/\lvert S\lvert$, $\mu$
and ${\vec{S}}$ are magnetic moment and the neutron’s spin, respectively. This
is the equation of motion of a classical magnetic dipole in a magnetic field,
which shows the precession of the spin vector ${\vec{S}}$ about the magnetic
field $\vec{B}$ with the Larmor frequency
$\omega_{\textrm{\scriptsize{L}}}=\left|{{{2\mu
B}\mathord{\left/{\vphantom{{2\mu B}\hbar}}\right.\kern-1.2pt}\hbar}}\right|$.
It is worth noting here that the Larmor precession angle (rotation angle) is
obtained with a frequency $\omega_{\textrm{\scriptsize{L}}}$ and the
propagation time T as $\omega_{\textrm{\scriptsize{L}}}\rm{T}$, solely depends
on the strength of the applied magnetic field. In practical experiments, the
Larmor precession is utilized by a variety of DC spin rotators.

In contrast, when a neutron interacts with a time-dependent, rather
oscillating, magnetic field, photon exchange occurs and the total energy of
neutrons is shifted. In particular when a so-called resonant spin-flipper is
activated, the neutron emits (or absorbs) a photon of energy
$\hbar\omega_{\textrm{\scriptsize{R}}}$, thus the total energy of the neutron
decreases (or increases) by $\hbar\omega_{\textrm{\scriptsize{R}}}$ during the
interaction [9]. As a consequence, an additional phase, a zero-field phase,
$\omega_{\textrm{\scriptsize{R}}}\rm{T}$ is accumulated during the propagation
afterwards. This phase shift, alternatively a spin precession, due to energy
shift emerges even without a guide field, thus it is referred to a zero-field
precession in literature [12, 13]. It should be emphasized that this zero-
field phase is independent of the applied guide field, namely the Larmor phase
$\omega_{\textrm{\scriptsize{L}}}\rm{T}$, and purely depends on the frequency
of the spin-flipper. In the following neutron polarimeter experiments, we
demonstrate how one can tune a zero-field phase
$\omega_{\textrm{\scriptsize{R}}}\rm{T}$ and a Larmor phase
$\omega_{\textrm{\scriptsize{L}}}\rm{T}$ independently.

Figure 1: Experimental setup for demonstrations of zero-field and Larmor spin
precessions. From a polarized incident beam the first $\pi/2$ spin-rotator
generates a superposition of $\lvert\pm z\rangle$. The first RF flipper
induces the energy shift $\hbar\omega_{R}$ for each flipped spin eigenstate
$\lvert\mp z\rangle$, which leads to the zero-field spin precession until the
second RF flipper where the energy difference is compensated. The guide field
$B_{0}$, permeated along the beam path, leads to the Larmor spin precession.
The second $\pi/2$ spin-rotator together with the analyzer enables to exhibit
the intensity modulation due to the spin precessions. The translation of the
DC and the RF flippers manipulate individually the two phases when the RF
flipper is turned on and off.

The experiments were carried out at the neutron polarimetry facility at the
250-kW TRIGA research reactor of the Atomic Institute of the Austrian
Universities in Vienna [18, 19, 20, 21]. A schematic view of the experimental
setup is shown in Fig.1. The incident neutron beam is monochromatized (with a
mean wave length of $\lambda=1.99\mbox{ \AA}$ which corresponds to a velocity
of 1990 m/s) by the use of a pyrolitic graphite crystal and polarized (average
degree of polarization over $98\%$) by reflection from a bent Co-Ti
supermirror array. The diameter of the beam is confined to about 4mm in
diameter by a Cd diaphragm. The initial polarization vector
$\vec{P}_{\textrm{i}}$ is perpendicular to the beam trajectory and defines the
$+\hat{\mathbf{z}}$-direction. Another suppermirror array is used to analyze
the final polarization $\vec{P}_{\textrm{f}}$. Depolarization of the neutron
beam is minimized by applying a guide field along the beam trajectory wherever
necessary: this guide field $B_{0}$ determines the Larmor frequency
$\omega_{\textrm{\scriptsize{L}}}$.

In a neutron polarimeter setup described above, after going through a $\pi/2$
spin-rotator, an incident polarization is interpreted as superposition of two
orthogonal spin states $\lvert\pm z\rangle$ and different phase shifts of
these two states under further spinor manipulations result in a polarization
change of the emerging neutron beam, i.e., we take a neutron polarization-
interference scheme. Between the two $\pi/2$ spin-rotators, three spin-
flippers are placed in the beam: the first two are RF spin-flippers followed
by a DC spin-flipper. In addition, one of the RF and the DC flippers are
mounted together on a single translator, allowing to tune the propagation time
of neutrons between the RF spin-flippers. The polarization vector after
passing through all spin-flippers is expected to lie in the xy-plane,
resulting from the interference between the states $\lvert\pm z\rangle$. A
spin analyzer together with the second $\pi/2$ spin-rotator is used to resolve
the different phase shifts accumulated through spin-flips.

Figure 2:  (a) Typical intensity oscillations with least square fits for the
Larmor spin precession. Curves are shown with the guide field $B_{0}=$10.90,
10.79, 10.69, 10.58G.(b) Dependence of the period on the strength of the guide
field depicted for seven values of the guide field. A clear linear dependence
agrees well with the theoretical predictions (see, Eq. (2)).

The first experiment shows a pure Larmor precession: both RF flippers are
turned off and only the DC flipper is in operation. In this case, the
superposed states $\lvert\pm z\rangle$ (and the flipped states $|\mp z\rangle$
later) simply propagate within the guide field $B_{0}$. These states
$|+z\rangle$ and $|-z\rangle$ are the eigenstates of the guide field: no
additional spinor rotation occurs. Then, each state merely obtains a Larmor
phase $\alpha$ due to the guide field, which is given by

$\displaystyle\alpha$ $\displaystyle=$
$\displaystyle{\omega}_{\textrm{\scriptsize{L}}}(\rm{T_{1}}-\rm{T_{2}})$ (2)
$\displaystyle=$
$\displaystyle{\omega}_{\textrm{\scriptsize{L}}}((\rm{T_{1}^{0}}+\Delta
T)-(\rm{T_{2}^{0}}-\Delta T))$ $\displaystyle=$
$\displaystyle\alpha_{0}+2{\omega}_{\textrm{\scriptsize{L}}}\Delta\rm{T},$

with
$\alpha_{0}=\omega_{\textrm{\scriptsize{L}}}(\rm{T_{1}^{0}}-\rm{T_{2}^{0}}+\rm{T_{3}^{0}})$.
$\rm{T_{1}}$, $\rm{T_{2}}$, $\rm{T_{j}^{0}}$ and $\Delta\rm{T}$ are the
propagation time before and after the DC flipper, the propagation time at the
initial position, and the time shift for $\rm{T_{1}}$, $\rm{T_{2}}$ by
shifting the position of the DC flipper. After the beam passes the $\pi/2$
spin-rotator and the analyzer, clear sinusoidal intensity oscillations of high
contrast were obtained. Typical experimental data with least square fits are
shown in Fig.2(a): the strength of the guide field was varied at
$B_{0}=$10.90, 10.79, 10.69, 10.58G. The shift of the oscillations is due to
$\omega_{\textrm{\scriptsize{L}}}\rm{T}$ (see, $\alpha_{0}$ in Eq.(2)). In
addition, the dependence of the period of the oscillations on the strength of
the guide field is plotted in Fig.2(b). A linear dependence, of which
inclination is in good agreement with the theoretically predicted value, is
seen, which confirms the precession angle given by Eq.(2).

Figure 3:  (a) Typical intensity oscillations with least square fits for the
zero-field spin precession. Curves are shown at the RF flipper frequencies
$\omega_{R}=$ 30.9, 30.6, 30.3, 30.0 kHz. (b) Dependence of the period on the
frequency depicted for seven values of the frequency. A clear linear
dependence agrees well with the theoretical predictions (see, Eq. (3)).

Next, characteristics of the zero-field precession are considered: both RF
flippers are turned on. In this case, the superposed states $\lvert\pm
z\rangle$ is once flipped to $\lvert\mp z\rangle$, and finally again flipped
back to $\lvert\pm z\rangle$. Then, the spin precession angle
$\alpha^{\prime}$ is expected to be given by the propagation time,
$\rm{T_{1}}$, $\rm{T_{2}}$, and $\rm{T_{3}}$ after each spin flipper by

$\displaystyle\alpha^{\prime}$ $\displaystyle=$
$\displaystyle{\omega}_{\textrm{\scriptsize{L}}}(\rm{T_{1}}-\rm{T_{2}}+\rm{T_{3}})+{\omega}_{\textrm{\scriptsize{R}}}\rm{T_{1}}$
(3) $\displaystyle=$
$\displaystyle{\omega}_{\textrm{\scriptsize{L}}}\\{(\rm{T_{1}^{0}}+\Delta
T)-\rm{T_{2}^{0}}+(\rm{T_{3}^{0}}-\Delta
T)\\}+{\omega}_{\textrm{\scriptsize{R}}}(\rm{T_{1}^{0}}+\Delta T)$
$\displaystyle=$
$\displaystyle\alpha^{\prime}_{0}+{\omega}_{\textrm{\scriptsize{R}}}\Delta\rm{T},$

with
$\alpha^{\prime}_{0}=\omega_{\textrm{\scriptsize{L}}}(\rm{T_{1}^{0}}-\rm{T_{2}^{0}}+\rm{T_{3}^{0}})+\omega_{\textrm{\scriptsize{R}}}\rm{T_{1}^{0}}$.
In this configuration, no Larmor precession is expected to be induced by the
change of $\Delta\rm{T}$, since a positive and a negative change of
$\rm{T_{1}}$ and $\rm{T_{3}}$ completely compensate each other by shifting the
position of the DC flipper and the second RF flipper. In order to prove the
frequency dependence of the precession, the frequency of the RF flippers were
varied with keeping the strength of the guide field by $B_{0}=$10.59G. Typical
intensity modulations are shown in Fig.3(a) at the RF flipper frequencies of
$\omega_{\textrm{\scriptsize{R}}}$=30.9, 30.6, 30.3, and 30.0kHz. Clear
sinusoidal oscillations of high contrast were again obtained. A slight
reduction of the amplitude of the obtained oscillations is solely due to the
detuning of the frequency $\omega_{\textrm{\scriptsize{R}}}$ in the vicinity
of the resonance for the flip-mode. And the shift of the oscillations arises
from $\omega_{\textrm{\scriptsize{R}}}\rm{T}$ (see, $\alpha^{\prime}_{0}$ in
eq.(3)). In addition, a dependence of the period on the frequency is plotted
in Fig.3(b). A linear dependence is seen and its inclination is in good
agreement with the theory, which confirms the precession angle given by
Eq.(3).

Figure 4:  (a) Typical intensity oscillations with least square fits for the
zero-field precession. Curves are shown with the guide field $B_{0}=$10.90,
10.79, 10.69, 10.58G. All curves are identical as expected. (b) The period is
plotted versus the strength of the guide field for seven values. It is
constant and independent of the strength of the guide field.

Note that the setup is constructed in a way that no spin rotation due to the
Larmor precession will occur, since the Larmor precessions before and after
the second RF flipper, i.e., $\Delta\rm{T_{1}}$ and $\Delta\rm{T_{3}}$, are
compensated each other with all flippers turned on (see Eq.(3)). This
independence of the Larmor precession can also be accessed in our experimental
setup. In particular, the frequency of the RF flippers were tuned at the
resonance $\omega_{0}=$31.8kHz and the strength of the guide field was varied
at $B_{0}=$10.90, 10.79, 10.69, 10.58G. Typical experimental data with least
square fits are shown Fig.4(a): identical sinusoidal oscillations are
obtained, confirming no spin rotation due to Larmor precessions. In addition,
the period is plotted versus the strength of the guide field in Fig.4(b).
Independent behavior of the period from the field strength is seen as expected
by Eq.(3).

The results are in good agreement with theoretically predicted behavior: the
frequency of the oscillating fields affects only on the zero-field precession
while the strength of the field does on the Larmor precession. The tactics,
that the RF and the DC flippers are displaced by the same amount, allows
individual control of both precessions. This is a great advantage for the
separate (phase) manipulation of the two-spaces, i.e., spin and energy spaces,
in future neutron optical experiments. We are now proceeding further neutron
interferometric and polarimetric experiments, where multi-entanglement in a
single neutron system will be investigated, in particular, with the use of
entanglement between degrees of freedoms such as path, spin, and energy.

In summary, we have presented experiments with a neutron polarimeter where two
spin-precessions are manipulated to show their characteristics. The zero-field
precession is independent of the strength of the guide field and purely
depends on the frequency of the RF spin-flippers, where an energy shift
occurs. In contrast, the Larmor precession solely depends on the strength of
the applied magnetic field. In addition, we have exhibited a method to
manipulate individually the two phases, resulting from each spin precession.
The method used here will be utilized in further neutron optical experiments
dealing with multi-entanglement in a single particle system.

This work has been partly supported by the Austrian Science Foundation, FWF
(P17803-N02). Y.H. would like to thank the Japan Science and Technology Agency
(JST) for financial support.

## References

  * [1] M.A. Nielsen and I. Chuang, _Quantum Computation and Quantum Information_ , (Cambridge Unviversity Press, Cambridge, 2000).
  * [2] S. Haroche and J. Raimond, _Exploring the Quantum_ , (Oxford Unviversity Press, Oxford, 2006).
  * [3] H. Rauch and S.A. Werner, _Neutron Interferometry_ , (Clarendon Press, Oxford, 2000).
  * [4] Y. Hasegawa, R. Loidl, G. Badurek, M. Baron, and H. Rauch, Nature 425, 45 (2003).
  * [5] Y. Hasegawa, R. Loidl, G. Badurek, M. Baron, and H. Rauch, Phys. Rev. Lett. 97, 230401 (2006).
  * [6] Y. Hasegawa, R. Loidl, G. Badurek, S. Filipp, J. Klepp, and H. Rauch, Phys. Rev. A 76, 052108 (2007).
  * [7] B.Alefeld, G. Badurek, and H. Rauch, Z. Phys B41 231 (1981).
  * [8] G. Badurek, H. Rauch, and J. Summhammer, Phys. Rev. Lett. 51, 1015 (1983).
  * [9] J. Summhammer, Phys. Rev. A 47, 56 (1993).
  * [10] S. Sponar, J. Klepp, R. Loidl, S. Filipp, G. Badurek, Y. Hasegawa, and H. Rauch, submitted to Phys. Rev. Lett.
  * [11] D.M. Greenberger, M.A. Horne, and A. Zeilinger, in _Bell‘s Theorem, Quantum Theory, and Concepts of the Universe_ , edited by M.Kafatos (Kluwer Academics, Dordrecht, The Netherlands, 1989).
  * [12] R. Golub, R. Gähler, and T. Keller, Am. J. Phys. 62 779 (1994).
  * [13] W.H. Kraan, S.V. Grigoriev, R. kreuger, F.M. Mulder and M.Th. Rekveldt, Physica B297, 23 (2001).
  * [14] R. Gähler, and R. Golub, Z. Phys. B65 269 (1987).
  * [15] R. Gähler, and R. Golub, Phys. Lett. A 123 43 (1987).
  * [16] D. Yamazaki, Nucl. instr. meth. A 488, 623 (2002).
  * [17] S.V. Grigoriev, W.H. Kraan, and M.Th. Rekveldt, Phys. Rev. A 69, 043615 (2004).
  * [18] Y. Hasegawa, S. Menhart, R. Meixner, and G. Badurek, Phys. Lett. A 234, 322 (1997)
  * [19] Y. Hasegawa, and G. Badurek, Phys. Rev. A 59, 4614 (1999)
  * [20] A.G. Wagh, G. Badurek, V.C. Rakhecha, R.J. Buchelt, A. Schricker, Phys. Lett. A 268, 209 (1997)
  * [21] J. Klepp, S. Sponar, Y. Hasegawa, E. Jericha,and G. Badurek, Phys. Lett. A 342 48 (2005).