We perform scalable approximate inference in continuous-depth Bayesian neural networks. In this model class, uncertainty about separate weights in each layer gives hidden units that follow a stochastic differential equation. We demonstrate gradient-based stochastic variational inference in this infinite-parameter setting, producing arbitrarily-flexible approximate posteriors. We also derive a novel gradient estimator that approaches zero variance as the approximate posterior over weights approaches the true posterior. This approach brings continuous-depth Bayesian neural nets to a competitive comparison against discrete-depth alternatives, while inheriting the memory-efficient training and tunable precision of Neural ODEs.
appendix.tex
\clearpage
\appendix

\thispagestyle{empty}


% For one-column format, uncomment the following:
\onecolumn \makesupplementtitle
% For two-column format, uncomment the following:
%\twocolumn[ \makesupplementtitle ]

% \onecolumn
% \appendix
% \section{Appendix}
\section{PROOFS}

\paragraph{Notation.} 
% Denote $\mathbb{T}=[0, T]\subseteq \mathbb{R}$ as the time horizon. 
Denote as $\phi$ the vector of variational parameters, $f_q$ as the approximate posterior on weights, $f_p$ as the prior on weights, $f_h$ as the dynamics of hidden units, and $\sigma$ as the diffusion function. Denote the Euclidean norm of a vector $u$ by $|u|$. For function $f$ denote its Jacobian as $\nabla f$.

% \subsection{Limit of the approximate posterior SDE}
% Let $x_0$ be a fixed initial state. Let $x_1, ..., x_N$ be states at times $\Delta t, 2\Delta t, \dots, N\Delta t = T$ with the Euler discretization:
% \eqn{
% x_{i+1} &= x_i + f_q(x_i) \Delta t + \sigma(x_i) (B(t + \Delta t) - B(t)),
% }
% where $B(\cdot)$ is the Brownian motion.
% This implies conditional on the previous state, the current state is normally distributed: 
% \eqn{
%     x_{i+1} | x_i &\sim \mathcal{N} ( x_i + f_q(x_i) \Delta t, \sigma(x_i)^2 \Delta t ) .
% }
% The following limit on the estimator with KL converges to $\E [ \int_0^T \frac{1}{2}\norm{u_t}^2 \dt ]$.
% \eq{
%     \lim_{N \to \infty} \sum_{i=0}^N \E_{x_i} \sbracks{ \KL \bracks{q(x_{i+1} | x_i) || p(x_{i+1} | x_i)}  }
%         =& \lim_{N \to \infty} \sum_{i=0}^N \E_{x_i} \sbracks{
%             \frac{ (f_q(x_i) - f_p(x_i))^2 }{ 2 \sigma(x_i)^2 } \Delta t
%         }\\
%          =& \E \sbracks{
%             \lim_{N \to \infty} \sum_{i=0}^N
%             \frac{ (f_q(x_i) - f_p(x_i))^2 }{ 2 \sigma(x_i)^2 } \Delta t
%         }\\
%         =& \E \sbracks{
%             \int_0^T \frac{(f_q(x_t) - f_p(x_t))^2}{ 2\sigma(x_t)^2} \dt 
%         }\\
%         =& \E \sbracks{ \int_0^T \frac{1}{2} |u_t|^2 \dt },
% }
% where (17) follows from the closed-form expression of KL between Gaussians (i.e. one step of the Euler discretization of the prior and approximate posterior from the same state, (18) follows from the dominated convergence theorem given that the diffusion processes for both the prior and approximate posterior are equivalent and hence the KL is finite, and (19) follows from the convergence of Reimann sums. Finally (20) is as defined in equation 10 of \cite{li2020scalable}.

\subsection{Derivation of an Alternative Monte Carlo Estimator}\label{app:mcestimator}
The goal of this section is to derive a Monte Carlo estimator of the KL-divergence on path space that is similar to the \textit{fully Monte Carlo} estimator described in~\cite{roeder2017sticking}. This will serve as the basis for the subsequent heuristic derivation of the continuous-time sticking-the-landing trick.

Let $w_0$ be a fixed initial state. Let $w_1, ..., w_N$ be states at times $\Delta t, 2\Delta t, \dots, N\Delta t = T$ generated by the Euler discretization:
\eq{
\label{eq:euler_discretize_posterior}
w_{i+1} &= w_i + f_q(w_i) \Delta t + \sigma(w_i) (B_{t + \Delta t} - B_{t} ) \\
        &= w_i + f_q(w_i) \Delta t + \sigma(w_i) \Delta t^{1/2} \epsilon_{i+1}, \quad \epsilon_{i+1} \sim \mathcal{N}(0, 1).
}
where $\{B_t\}_{t\ge 0}$ is the Brownian motion.
This implies that conditional on the previous state, the current state is normally distributed: 
\eqn{
    w_{i+1} | w_i &\sim \mathcal{N} ( w_i + f_q(w_i) \Delta t, \sigma(w_i)^2 \Delta t ) .
}
Thus, the log-densities can be evaluated as
\eq{
    \label{eq:logq_step}
    \log q(w_{i+1} | w_i) &=
        -\frac{1}{2} \log ( 2\pi \sigma(w_i)^2 \Delta t ) -\frac{1}{2} \frac{
            \bracks{w_{i+1} - (w_i + f_q(w_i) \Delta t )}^2
        }{
            \sigma(w_i)^2 \Delta t
    }, \quad i = 0, \dots N - 1.
}
On the other hand, if at any time, the next state was generated from the current state based on the prior process, we would have the following log-densities:
\eq{
    \label{eq:logp_step}
    \log p(w_{i+1} | w_i) &=
        -\frac{1}{2} \log ( 2\pi \sigma(w_i)^2 \Delta t ) -\frac{1}{2} \frac{
            \bracks{w_{i+1} - (w_i + f_p(w_i) \Delta t )}^2
        }{
            \sigma(w_i)^2 \Delta t
    }, \quad i = 0, \dots N - 1.
}
Now, we substitute the form of $w_{i+1}$ based on~\eqref{eq:euler_discretize_posterior} into~\eqref{eq:logq_step} and~\eqref{eq:logp_step} and obtain
\eqn{
\log q(w_{i+1} | w_i) =& -\frac{1}{2} \log ( 2\pi \sigma(w_i)^2 \Delta t ) -\frac{1}{2} \epsilon_{i+1}^2, \\
\log p(w_{i+1} | w_i) =&
    -\frac{1}{2} \log ( 2\pi \sigma(w_i)^2 \Delta t )
    \\&
    -\frac{1}{2} \Bigg(
        \frac{ (f_q(w_i) - f_p(w_i))^2 }{ \sigma(w_i)^2 } \Delta t + \frac{ 2 ( f_q(w_i) - f_p(w_i) ) \epsilon_{i+1} }{ \sigma(w_{i}) } \Delta t^{1/2}
        + \epsilon_{i+1}^2
    \Bigg).
}
The KL divergence on the path space could then be regarded as a sum of infinitely many KL-divergences between Gaussians:
\eq{
     &\lim_{N \to \infty} \sum_{i=0}^N \E_{w_i} \sbracks{ \KL \bracks{q(w_{i+1} | w_i) || p(w_{i+1} | w_i)}  } \\
    =&
        \lim_{N \to \infty} \sum_{i=0}^N \E_{w_i} \sbracks{
        \E_{w_{i + 1} \sim q(w_{i+1} | w_i)} \sbracks{
        \log \frac{q(w_{i+1} | w_i)}{p(w_{i+1} | w_i)}
        }
    } \\
    =&
        \lim_{N \to \infty} \sum_{i=0}^N \E_{w_i} \sbracks{
            \E_{\epsilon_{i+1}} \sbracks{
                \frac{ (f_q(w_i) - f_p(w_i))^2 }{ 2 \sigma(w_i)^2 } \Delta t + 
                \frac{ (f_q(w_i) - f_p(w_i)) }{ \sigma(w_i) } \Delta t^{1/2} \epsilon_{i+1}
            }
        } \\
    =& \;
        \E \sbracks{
            \frac{1}{2} \int_0^T | u_t |^2 \dt + \int_0^T u_t \dB_t
        }.
    \label{eq:fmc}
}

\subsection{Sticking-the-landing in Continuous Time}
\label{app:stlheuristic}
For a non-sequential latent variable model, the sticking-the-landing (STL) trick removes from the fully Monte Carlo ELBO estimator a score function term of the form $\partial \log q(w, \phi) / \partial \phi$, where $w$ is sampled using the reparameterization trick and may depend on $\phi$. 
The score function term has $0$ expectation, but may affect the variance of the gradient estimator for the inference distribution's parameters.

Here, we exploit this intuition and apply it to each step before taking the limit. More precisely, we apply the STL trick to estimate the gradient of $\KL ( q(w_{i+1} | w_i) || p(w_{i+1} | w_i) )$ for $i=1, 2, \dots ,N$, and thereafter take the limit as the mesh size of the discretization goes to $0$. For each individual term, the score function term to be removed is 
\eqn{
    \frac{\partial }{\partial \phi} \log q(w_{i+1} | w_i, \phi) =&
        -\frac{1}{ 2 \sigma^2(w_i) \Delta t } \frac{\partial }{\partial \phi}
        \sbracks{
            \bracks{
                w_{i+1} - (w_i + f_q( w_i, \phi) \Delta t)
            }^2
        } \\
    =&
        \frac{\partial }{\partial \phi} \sbracks{
            \frac{ f_q(w_i, \phi) }{ \sigma(w_i) }
        } \epsilon_{i+1} \Delta t^{1/2}.
}
Now, we sum up all of these terms and take the limit as $\Delta t \to 0$. This gives us
\eqn{
    &\lim_{N \to \infty} \sum_{i=0}^N \E_{w_i} \sbracks{
        \E_{ w_{i+1} \sim q(w_{i+1} | w_i) } \sbracks{
            \frac{\partial }{\partial \phi} \log q(w_{i+1} | w_i)
        }
    }\\
    =&
        \lim_{N \to \infty} \sum_{i=0}^N \E_{w_i} \sbracks{
            \E_{ \epsilon_{i+1} } \sbracks{
                \frac{\partial }{\partial \phi} \sbracks{
                    \frac{ f_q(w_i, \phi) }{ \sigma(w_i) }
                } \epsilon_{i+1} \Delta t^{1/2}
            }
        } \\
    =&
        \; \E\sbracks{
            \int_0^T \frac{\partial }{\partial \phi} \sbracks{ 
                \frac{ f_q(w_t, \phi) }{ \sigma(w_t) }
            } \dB_t
        } \\
    =&
        \; \E\sbracks{
            \int_0^T
            \frac{\partial }{\partial \phi} \sbracks{u_t} \dB_t
        }.
}

% \eqn{
%     h_T = h_0 + \int_0^T f(h_t, t)dt + \int_0^T g(h_t, t) dB_t
% }

Removing this term from the fully Monte Carlo estimator in~\eqref{eq:fmc} gives rise to the following estimator of a surrogate objective that facilitates implementation:
\eqn{
     \widehat{\text{ELBO}} =&
        \log p ( \mathcal{D} \mid  w ) 
        - \int_{t_0}^{t_1} \frac{1}{2} \norm{ u(w_t, t, \phi) }_2^2 \dt 
        \\&
        - \int_{t_0}^{t_1} u(w_t, t, \texttt{stop\_gradient}(\phi)) \dB_t, \quad w(\cdot) \sim q_\phi().
}
% \XL{Unify notation on drawing a sample from the approximate posterior.}

% \eq{
%     &\lim_{N \to \infty} \sum_{i=0}^N \E_{x_i} \sbracks{ \KL \bracks{q(x_{i+1} | x_i) || p(x_{i+1} | x_i)}  } \\
%     =&
%         \lim_{N \to \infty} \sum_{i=0}^N \E_{x_i} \sbracks{
%         \E_{x_{i + 1} \sim q(x_{i+1} | x_i)} \sbracks{
%         \log \frac{q(x_{i+1} | x_i)}{p(x_{i+1} | x_i)}
%         }
%     } \\
%     =&
%         \lim_{N \to \infty} \sum_{i=0}^N \E_{x_i} \sbracks{
%             \E_{\epsilon_{i+1}} \sbracks{
%                 \frac{ (f_q(x_i) - f_p(x_i))^2 }{ 2 \sigma(x_i)^2 } \Delta t
%         \\
%         & + \frac{ ( f_q(x_i) - f_p(x_i) ) \epsilon_{i+1} }{ \sigma(x_{i}) } \Delta t^{1/2}
%          \epsilon_{i+1}^2
%             }
%         } \\
%     =&
%         \E \sbracks{
%             \frac{1}{2} \int_0^T | u_t |^2 \dt + \int_0^T u_t \dB_t
%         }.
% }

\newpage
\section{EXPERIMENTAL SETTINGS}
\begin{table*}[!htbp]
\centering
% \ra{1.1}
\setlength{\tabcolsep}{0.8em}
\resizebox{1\textwidth}{!}{
\begin{tabular}{@{} l l l l l l @{}}\toprule
  & \multicolumn{1}{c}{\textbf{\textsc{}}} &
  & \multicolumn{1}{c}{\textbf{\textsc{Experiments}}} &
  & \multicolumn{1}{c}{\textbf{\textsc{}}} \\
% \cmidrule(lr){1-2} 
\cmidrule(l){3-5} 
Model & Hyper-parameter & 1D Regression
& MNIST~\cite{deng2012mnist} & CIFAR-10~\cite{krizhevsky2014cifar} \\
\cmidrule(r){1-2} \cmidrule(l){3-5} 
ResNet32 
 & Learning Rate & -- & 1e-3 & 7e-4 \\
 & Batch Size & -- & 128 & 128 \\
 & Activation & -- & tanh & tanh \\
 & Epochs & -- & 100 & 500 \\
\cmidrule(r){1-2} \cmidrule(l){3-5} 
ODEnet
 & Augment dim. & 2 & 2 & 2 \\
 & \# blocks & 1 & 1 & 2-2-2 \\
 & Diffusion $\sigma$ & 0 & 0 & 0 \\
 & KL coef. & 0 & 0 & 0 \\
 & Learning Rate & 1e-3 & 1e-3 & 7e-4 \\
 & \# Solver Steps & 10 & 20 & 20 \\
 & Batch Size & 40 & 128 & 128 \\
 & Activation & swish & tanh & tanh \\
 & Epochs & 800 & 100 & 500 \\
\cmidrule(r){1-2} \cmidrule(l){3-5} 
HyperODEnet 
 & <ODEnet> & -- & <ODEnet> & <ODEnet> \\
 & KL coef. & -- & 1e-3 & 1e-3 \\
 & Drift $f_w$ dim. & -- & 1-64-1 & 1-128-1 \\
\cmidrule(r){1-2} \cmidrule(l){3-5} 
MFVI ResNet32 
 & <ResNet32> & -- & <ResNet32> & <ResNet32> \\
 & KL coef. & -- & 1e-3 & 1e-3 \\
\cmidrule(r){1-2} \cmidrule(lr){3-5}
MFVI ODEnet 
 & <ODEnet> & -- & <ODEnet> & <ODEnet> \\
 & KL coef. & -- & 1e-3 & 1e-3 \\
\cmidrule(r){1-2} \cmidrule(lr){3-5}
MFVI HyperODEnet 
 & <MFVI ODEnet> & -- & <MFVI ODEnet> & <MFVI ODEnet> \\
 & Drift $f_w$ dim. & -- & 1-64-1 & 1-128-1 \\
\cmidrule(r){1-2} \cmidrule(lr){3-5}
SDE BNN 
 & <ODEnet> & <ODEnet> & <ODEnet> & <ODEnet> \\
 & Learning Rate & 1e-3 & 1e-3 & 7e-4 \\
 & \# blocks & 1 & 1 & 2-2-2 \\
 & Drift $f_x$ dim. & 32 & 32 & 64 \\
 & Drift $f_w$ dim. & 32 & 1-64-1 & 2-128-2 \\
 & Diffusion $\sigma$ & 0.2 & 0.1 & 0.1 \\
 & \# Posterior Samples & 20 & 1 & 1 \\
\cmidrule(r){1-2} \cmidrule(lr){3-5}
SDE BNN (+ STL) & <SDE BNN> & <SDE BNN> & <SDE BNN> & <SDE BNN> \\
% SDE BNN $w0$ inferred & 99.04 $\pm$ 0.03 & 0.73 $\pm$ 0.04 & &  ---  &  ---  \\
% SDE BNN $w0$ inferred (+ STL) & 99.05 $\pm$ 0.00 & 0.79 $\pm$ 0.04 & &  ---  &  ---  \\
\bottomrule
\end{tabular}
}
\caption{These are the hyper-parameters for each method of evaluation pertaining to results in the toy and classification tasks of Table \ref{tab1:img_results}. Each model was run on a single Nvidia RTX6000 GPU on our compute clusters. SDE and learning optimization parameters were tuned according to a validation set sampled randomly from 10\% of the training set. No schedules of any kind on the hyper-parameters were used in training. Settings with high overlap with another model are indicated using \texttt{<model>} with additional parameters overridden as necessary. Each block is separated either by a downsampling or upsampling convolutional layer (i.e., the \texttt{-}'s).}
\label{tab:hparam settings}
\end{table*}

\newpage
\section{ADDITIONAL RESULTS}\label{app:additional_figures}

\subsection{Augmentation in Differential Equation Models}

\begin{figure}[h]
    \centering
    \begin{subfigure}[b]{0.32\linewidth}
         \centering
         \includegraphics[trim=8cm 0 0 0,clip,width=\linewidth]{plots/fig1_augdim0/predictive.pdf}
         \caption{Non-augmented dimension}
    \end{subfigure}
    \begin{subfigure}[b]{0.32\linewidth}
         \centering
         \includegraphics[trim=8cm 0 0 0,clip,width=\linewidth]{plots/fig1_augdim1/diff0.3/predictive.pdf}
         \caption{from 2nd augmented dimension}
    \end{subfigure}
    \begin{subfigure}[b]{0.32\linewidth}
         \centering
         \includegraphics[trim=8cm 0 0 0,clip,width=\linewidth]{plots/fig1_augdim-1/sdebnn_sample.pdf}
         \caption{from last augmented dimension}
     \end{subfigure}
     \vspace{-0.6em}
    \caption{Example flows sampled from learned SDE dynamics. All continuous-depth models were trained by augmenting the state by 2 dimensions, refer to Figure~\ref{fig:dynamics} for main results. \textit{Left:} The SDE-BNN learns meaningful parameterizations on the non-extraneous dimensions of the input state vector. In the case of a true function being monotonic, the augmented dimensions simply help the main output.
    %Here, the diffusion was increased from 0.2 in the toy regression task to 0.3 to show a proportionate increase in variance when sampling the dynamics. 
    \textit{Middle:} The model learns to ignore dimensions that are not necessary to train on, especially on simpler tasks as in the toy setting. Samples in augmented dimensions can overlap for different input values in the given domain $(-5, 5)$. \textit{Right:} Similarly, the last output dimension was also associated with augmentation and was not a well learned representation of the data, ignoring the initial inputs entirely (all values are 0).}
    \label{fig:posteriorpredictive}
\end{figure}

\subsection{Classification Results}
\begin{table*}[!htbp]
\centering
% \ra{1.1}
\setlength{\tabcolsep}{0.8em}
\resizebox{1\textwidth}{!}{
\begin{tabular}{@{} l c c c c c c @{}}\toprule
  & \multicolumn{2}{c}{\textbf{\textsc{MNIST}}} &
  & \multicolumn{2}{c}{\textbf{\textsc{CIFAR-10}}} \\
\cmidrule(lr){2-3} \cmidrule(l){5-6} 
Model & \tcen{Accuracy (\%) } & \tcen{ECE ($\times 10^{-2}$)}
& & \tcen{Accuracy (\%)} & \tcen{ECE ($\times 10^{-2}$)} \\
\cmidrule(r){1-1}\cmidrule(lr){2-3} \cmidrule(l){5-6} 
SDE BNN & 99.30 $\pm$ 0.09 & 0.63 $\pm$ 0.10 & &  88.08 $\pm$ 1.25 & 7.53 $\pm$ 0.44 \\
SDE BNN (+ STL) & 99.10 $\pm$ 0.09 & 0.78 $\pm$ 0.12 & & 87.95 $\pm$ 1.32 & 7.94 $\pm$ 0.59 \\
SDE BNN $w0$ inferred & 99.04 $\pm$ 0.03 & 0.73 $\pm$ 0.04 & & 88.04 $\pm$ 0.30 &  6.56 $\pm$ 0.62 \\
SDE BNN $w0$ inferred (+ STL) & 99.05 $\pm$ 0.00 & 0.79 $\pm$ 0.04 & &  87.37 $\pm$ 0.63  &  6.44 $\pm$ 0.10\\
\bottomrule
\end{tabular}
}
\caption{Classification accuracy and expected calibration error on MNIST ($100^{th}$ epoch) and CIFAR-10 ($300^{th}$ epoch) for additional baseline to the ones in Table~\ref{tab1:img_results}. Values are compared at the 100th epoch for MNIST and 300th for CIFAR-10. Here $w0$, the initial drift of the posterior SDE, is inferred using a Gaussian prior rather than being a fixed value. The best prior variance was selected in a preliminary sweep between values in the range [0.1, 0.44]. The performance is slightly worse than the point estimate but displays better calibration as a trade-off. It can be noted that calibration may appear better earlier on in training, as in prior to converging and reaching non-uniform confidence, but model predictions are not necessarily correct.}
\label{tab3:add_img_results}
\end{table*}
 
\newpage
\subsection{Sticking the Landing Results}
\begin{table*}[!htbp] %[h]
\centering
% \ra{1.1}
\setlength{\tabcolsep}{1em}
\resizebox{1\textwidth}{!}{
\begin{tabular}{lccc}
\hline
Method &
\begin{tabular}[c]{@{}c@{}}Accuracy (\%)\end{tabular} &
\begin{tabular}[c]{@{}c@{}}Negative \\ Log-likelihood ($\times 10^{-4}$) \end{tabular} &
\begin{tabular}[c]{@{}c@{}}ELBO \end{tabular} \\ 
\hline
\begin{tabular}[c]{@{}l@{}}SDE BNN\end{tabular} & 95.91 $\pm$ 0.2 & 1.17 $\pm$ 0.309 & 1.40 $\pm$ 0.2 \\ \hline
SDE BNN (+STL) & 96.89 $\pm$ 0.2 & 0.309 $\pm$ 0.15 & 1.183 $\pm$ 0.2  \\
\end{tabular}
}
\caption{Training with STL estimator on CIFAR-10 shows training time improvements in accuracy, negative log likelihood, and ELBO objective in addition to reducing variance. This improvement to the standard gradient estimator can be especially useful in settings where the approximate posterior is sufficiently flexible (i.e. the drift neural net is very large relative to the state size).}
\label{tab4:stl_improvement}
\end{table*}


\subsection{Calibration Results}
\label{app:calibration}
\begin{figure}[H]
    \centering
    \begin{subfigure}[b]{0.45\linewidth}
         \centering
         \includegraphics[width=\linewidth]{cifar10c/model vs perf/with_error-fig1_error_bar.pdf}
         \caption{Model vs Classification Error}
         \label{fig:cifarcmodelvserror}
    \end{subfigure}
    \begin{subfigure}[b]{0.45\linewidth}
         \centering
         \includegraphics[width=\linewidth]{cifar10c/model vs perf/with_error-fig2_ece_bar.pdf}
         \caption{Model vs Expected Calibration Error}
         \label{fig:cifarcmodelvsece}
    \end{subfigure}
    \begin{subfigure}[b]{0.45\linewidth}
         \centering
         \includegraphics[width=\linewidth]{cifar10c/model vs perf/with_error-fig3_brier_bar.pdf}
         \caption{Model vs Brier Score}
         \label{fig:cifarcmodelvsbrier}
     \end{subfigure}
    \begin{subfigure}[b]{0.45\linewidth}
         \centering
         \includegraphics[width=\linewidth]{cifar10c/severity vs metric/error_bar-fig6_severity_brier_bar.pdf}
         \caption{Corruption severity vs Brier Score}
         \label{fig:cifarcseverityvsbrier}
    \end{subfigure}
    \caption{Figures~\ref{fig:cifarcmodelvserror}-\ref{fig:cifarcmodelvsbrier} show that the SDE BNN and SDE BNN + STL models outperform their non-continuous depth ResNet counterparts on all three robustness metrics when evaluated on the corrupt CIFAR-10C benchmarks. Figure~\ref{fig:cifarcseverityvsbrier} indicates that the accuracy of predictions is relatively consistent across all severity levels with the SDE-BNN and SDE-BNN + STL models having relatively better calibrated predictions.}
    \label{fig:otherbnnevals}
\end{figure}

\newpage
\subsection{Comparisons with Other Bayesian Models}
\begin{table*}[!htbp] %[h]
\centering
% \ra{1.1}
\setlength{\tabcolsep}{0.2em}
\resizebox{1\textwidth}{!}{
\begin{tabular}{lcccc}
\hline
Method &
\begin{tabular}[c]{@{}c@{}}Posterior over\\Stochastic Process\end{tabular} &
\begin{tabular}[c]{@{}c@{}}Flexible \\Approximate Posterior\end{tabular} &
\begin{tabular}[c]{@{}c@{}}Adaptive \\Computation \end{tabular} &
References \\ \hline
\begin{tabular}[c]{@{}l@{}}Bayes by Backprop\end{tabular} & \xmark & \xmark & \xmark &  \citet{blundell2015weight}  \\ \hline
MCMC for BNNs & \xmark & \cmark & \xmark & \citep{neal2012bayesian, wenzel2020good, izmailov2021bayesian}   \\
\begin{tabular}[c]{@{}l@{}}Bayesian Hypernets\end{tabular} & \xmark & \cmark & \xmark & \citet{krueger2018bayesian} \\ \hline
BBVI for SDEs & \cmark & \xmark & \xmark & \citet{ryder2018blackbox}  \\ \hline
Bayesian Neural ODEs & \xmark & \xmark & \cmark & 
\begin{tabular}[c]{@{}c@{}}\citet{yildiz2019ode}\\\citet{dandekar2020bayesian}\end{tabular} \\ \hline
SDE-BNN & \cmark & \cmark & \cmark  & current work
\end{tabular}
}
\caption{Properties of various Bayesian supervised learning approaches.}
\label{tab5:comparemethods}
% \vspace{-1.5em} % Add back if in main text not in appendix!
\end{table*}
% \improvement{Per Neurips style guide, caption of table should appear at top and not bottom!!!}

\subsection{Other Bayesian Methods}
\begin{figure}[!htbp] %[H]
% \vspace{-2.2em}
\vspace{-1em}
    \centering
    \subfloat{{\includegraphics[width=0.31\linewidth]{bnns/gp_toy_square.pdf} }}%
    % \qquad
    \vspace{-0.1cm}
    \subfloat{{\includegraphics[width=0.31\linewidth]{bnns/ensemble_shaded_square.pdf} }}%
    % \qquad
    \vspace{-0.1cm}
    \subfloat{{\includegraphics[width=0.31\linewidth]{bnns/mfvi_toy.pdf} }}%
    \vspace{-0.5em}
    \caption{Approximate posteriors from other common Bayesian statistical models. \textit{Left:} Gaussian Process. \text{Center:} Deep Ensemble K=8. \text{Right:} MFVI. Different variances and extrapolations are learned from the SDE-BNN across other Bayesian model parameterizations, which can result in more or less reasonable uncertainty bounds depending on interpretation.}%
    % \vspace{-1.5em}
    \label{fig:morebaselines}%
\end{figure}


% \newpage
\subsection{Robustness to solver error at test time}
\label{app:solvererror}
\begin{figure}[H]
    \centering
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth]{cifar10/sdebnn/train/val-calib1210-20.pdf}
    \vspace{-2em}
    \caption{160 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth]{cifar10/sdebnn/train/val-calib1210-22.pdf}
    \vspace{-2em}
    \caption{176 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth]{cifar10/sdebnn/train/val-calib1210-24.pdf}
    \vspace{-2em}
    \caption{192 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth]{cifar10/sdebnn/train/val-calib1210-26.pdf}
    \vspace{-2em}
    \caption{208 steps}
    \end{subfigure}
    \caption{CIFAR10 image classification with a SDE-BNN. Better calibration can be obtained by increasing solver step sizes during inference without substantially changing the training error.}
    \label{fig:cifar10_calib_spectrum}
\end{figure}

\begin{figure}[H]
    \centering
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth, trim=11.5px 10px 12px 19px, clip]{cifar10/sdebnn-ou/calib/test-calib0-18.pdf}
    \vspace{-2em}
    % \label{fig:ou_sdebnn_cifar_test18}
    \caption{154 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth, trim=11.5px 10px 12px 19px, clip]{cifar10/sdebnn-ou/calib/test-calib0-20.pdf}
    \vspace{-2em}
    % \label{fig:ousdebnncifartest20}
    \caption{160 steps}
    \label{fig:stepstrained}
    \end{subfigure}
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth, trim=11.5px 10px 12px 19px, clip]{cifar10/sdebnn-ou/calib/test-calib0-22.pdf}
    \vspace{-2em}
    % \label{fig:ou_sdebnn_cifar_test22}
    \caption{176 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.23\linewidth}
    \includegraphics[width=\linewidth, trim=11.5px 10px 12px 19px, clip]{cifar10/sdebnn-ou/calib/test-calib0-24.pdf}
    \vspace{-2em}
    % \label{fig:ou_sdebnn_cifar_test24}
    \caption{192 steps}
    \end{subfigure}
    \caption{CIFAR10 image classification with a SDE-BNN. Generalization improves marginally compared to a trained model during inference in \ref{fig:stepstrained}, as tuning solver step size does not yield significant differences in calibration outcomes.}
    \label{fig:cifar10_calib_spectrum_inference}
\end{figure}
\vspace{-0.5cm
}
\subsection{SDE solver and adjoint settings}
\label{app:adjointadaptivesolver}

These were run with a SDE-BNN for MNIST image classification, to compare the performance and run-time cost across different solver settings. Comparably, backpropagation through the solver averaged 162.58 sec / epoch while the adjoint method averaged 135.90 sec / epoch in terms of wall clock time.
% \vspace{-0.2cm}
\begin{figure}[H] %[!h] %[H]
    \centering
    \begin{subfigure}[b]{0.32\linewidth}
    \includegraphics[width=\linewidth]{torchsde/adjoint/100 and time/plot_backprop_vs_adjoint_ep100_val_xent_time1_back1.pdf}
    \vspace{-2em}
    % \caption{160 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.32\linewidth}
    \includegraphics[width=\linewidth]{torchsde/adjoint/100 and time/plot_backprop_vs_adjoint_ep100_train_loss_time1_back1.pdf}
    \vspace{-2em}
    % \caption{176 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.32\linewidth}
    \includegraphics[width=\linewidth]{torchsde/adjoint/100 and time/plot_backprop_vs_adjoint_ep100_train_kl_time1_back1.pdf}
    \vspace{-2em}
    % \caption{192 steps}
    \end{subfigure}
    \caption{Backpropagation through the SDE solver yields similar optimization dynamics but is less time efficient than the adjoint method.}
    \label{fig:mnist_adjoint}
\end{figure}
\vspace{-0.8cm}
\begin{figure}[H] %[!h]
    \centering
    \begin{subfigure}[b]{0.31\linewidth}
    \includegraphics[trim={0 0 2.15cm 0},clip,width=\linewidth]{torchsde/adaptive/adaptive nfes/plot_adaptive_vs_adjoint_ep100_train_acc_nfe1_adjFalse.pdf}
    \vspace{-2em}
    % \caption{160 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.31\linewidth}
    \includegraphics[trim={0 0 2.15cm 0},clip,width=\linewidth]{torchsde/adaptive/adaptive nfes/plot_adaptive_vs_adjoint_ep100_train_loss_nfe1_adjFalse.pdf}
    \vspace{-2em}
    % \caption{176 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.36\linewidth}
    \includegraphics[width=\linewidth]{torchsde/adaptive/adaptive nfes/plot_adaptive_vs_adjoint_ep100_val_xent_nfe1_adjFalse.pdf}
    \vspace{-2em}
    % \caption{192 steps}
    \end{subfigure}
    \caption{Trade-off between solver speed and convergence during training. Adaptive refers to training with the stochastic adjoint in both forward and reverse modes here.}
    \label{fig:mnist_adaptive}
\end{figure}
\vspace{-0.75cm}
\begin{figure}[H]
    \centering
    \begin{subfigure}[b]{0.31\linewidth}
    \includegraphics[trim={0 0 2.15cm 0},clip,width=\linewidth]{torchsde/adaptive/100 epochs/plot_adaptive_vs_adjoint_ep100_train_acc_nfe1.pdf}
    \vspace{-2em}
    % \caption{160 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.31\linewidth}
    \includegraphics[trim={0 0 2.15cm 0},clip,width=\linewidth]{torchsde/adaptive/100 epochs/plot_adaptive_vs_adjoint_ep100_train_loss_nfe1.pdf}
    \vspace{-2em}
    % \caption{176 steps}
    \end{subfigure}
    \begin{subfigure}[b]{0.36\linewidth}
    \includegraphics[width=\linewidth]{torchsde/adaptive/100 epochs/plot_adaptive_vs_adjoint_ep100_val_xent_nfe1.pdf}
    \vspace{-2em}
    % \caption{192 steps}
    \end{subfigure}
    \caption{Trade-off between solver speed and precision during training. Adaptive-order optimization trajectories were comparable to fixed-order solvers and were thus not applied to the classification tasks since computational resources were not under constraint.}
    \label{fig:mnist_adaptive_adjoint}
\end{figure}
main.tex
\documentclass[twoside]{article}

% \usepackage{aistats2022}
% If your paper is accepted, change the options for the package
% aistats2022 as follows:
%
\usepackage[accepted]{aistats2022}
%
% This option will print headings for the title of your paper and
% headings for the authors names, plus a copyright note at the end of
% the first column of the first page.

% If you set papersize explicitly, activate the following three lines:
%\special{papersize = 8.5in, 11in}
%\setlength{\pdfpageheight}{11in}
%\setlength{\pdfpagewidth}{8.5in}

% If you use natbib package, activate the following three lines:
\usepackage[round]{natbib}
\renewcommand{\bibname}{References}
\renewcommand{\bibsection}{\subsubsection*{\bibname}}

% If you use BibTeX in apalike style, activate the following line:
%\bibliographystyle{apalike}

\input{mathdef.tex}
\input{math_commands.tex}

\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc}    % use 8-bit T1 fonts
\usepackage{hyperref}       % hyperlinks
\usepackage{url}            % simple URL typesetting
\usepackage{booktabs}       % professional-quality tables
\usepackage{amsfonts}       % blackboard math symbols
\usepackage{nicefrac}       % compact symbols for 1/2, etc.
\usepackage{microtype}      % microtypography
\usepackage[pdftex,dvipsnames]{xcolor}  % Coloured text etc.
\usepackage{wrapfig,lipsum}
\usepackage{pifont}
\definecolor{myred}{RGB}{215,48,39}
\definecolor{mygreen}{RGB}{26,152,80}
\newcommand{\cmark}{\textcolor{mygreen}{\ding{51}}}
\newcommand{\xmark}{\textcolor{myred}{\ding{55}}}
\usepackage{xspace}
\usepackage{enumitem}
\newcommand*{\eg}{{\it e.g.}\@\xspace}
\newcommand*{\ie}{{\it i.e.}\@\xspace}
\newcommand\numberthis{\addtocounter{equation}{1}\tag{\theequation}}
\renewcommand{\ra}[1]{\renewcommand{\arraystretch}{#1}}
\newcommand{\D}{\mathcal{D}}

\usepackage{mathtools}

\usepackage{floatrow}
\floatsetup[table]{capposition=top}
% Recommended, but optional, packages for figures and better typesetting:
\usepackage{microtype}
\usepackage{graphicx}
% \usepackage{subfigure}
\usepackage{booktabs} % for professional tables

\usepackage{wrapfig}
\usepackage{caption}
\usepackage{subcaption}
\usepackage{float}

\usepackage{xargs}                      % Use more than one optional 
\usepackage[colorinlistoftodos]{todonotes}
\newcommandx{\improvement}[2][1=]{\todo[linecolor=Plum,backgroundcolor=Plum!25,bordercolor=Plum,#1]{#2}}

\definecolor{mydarkblue}{rgb}{0,0.08,0.45}
\hypersetup{
    colorlinks=true,
    linkcolor=mydarkblue,
    citecolor=mydarkblue,
    filecolor=mydarkblue,
    urlcolor=mydarkblue
}

\newif\ifcomments

\commentstrue
% \commentsfalse

% \newif
\ifcomments
\newcommand{\ricky}[1]{\textcolor{blue}{[RC: #1]}}
\else
\newcommand{\ricky}[1]{}
\fi

\ifcomments
\newcommand{\winnie}[1]{\textcolor{blue}{[WX: #1]}}
\else
\newcommand{\winnie}[1]{}
\fi

\ifcomments
\newcommand{\XL}[1]{\textcolor{cyan}{(XL: #1)}}
\else
\newcommand{\XL}[1]{}
\fi

\begin{document}

% If your paper is accepted and the title of your paper is very long,
% the style will print as headings an error message. Use the following
% command to supply a shorter title of your paper so that it can be
% used as headings.
%
\runningtitle{Infinitely Deep Bayesian Neural Networks with SDEs}

% If your paper is accepted and the number of authors is large, the
% style will print as headings an error message. Use the following
% command to supply a shorter version of the authors names so that
% they can be used as headings (for example, use only the surnames)
%
%\runningauthor{Surname 1, Surname 2, Surname 3, ...., Surname n}

\twocolumn[

\aistatstitle{Infinitely Deep Bayesian Neural Networks \\ with Stochastic Differential Equations}

\aistatsauthor{ Winnie Xu \And Ricky T.Q. Chen \And  Xuechen Li \And David Duvenaud }

\aistatsaddress{ University of Toronto \And  University of Toronto \And Stanford Unviersity \And University of Toronto } ]

\begin{abstract}
We perform scalable approximate inference in continuous-depth Bayesian neural networks.
In this model class, uncertainty about separate weights in each layer gives hidden units that follow a stochastic differential equation.
We demonstrate gradient-based stochastic variational inference in this infinite-parameter setting, producing arbitrarily-flexible approximate posteriors.
We also derive a novel gradient estimator that approaches zero variance as the approximate posterior over weights approaches the true posterior.
This approach brings continuous-depth Bayesian neural nets to a competitive comparison against discrete-depth alternatives, while inheriting the memory-efficient training and tunable precision of Neural ODEs.
\end{abstract}

\section{INTRODUCTION}\label{sec:intro}
Taking the limit of neural networks to be the composition of infinitely many residual layers provides a way to implicitly define its output as the solution to an ODE~\citep{haber2017stable,E2017}.
This continuous-depth parameterization decouples the specification of the model from its computation.
While the paradigm adds complexity, it has several benefits:
(1) Computational cost can be traded for precision in a fine-grained manner by specifying error tolerances for adaptive computation, and
(2) memory costs for training can be significantly reduced by running the dynamics backwards in time to reconstruct activations of intermediate states needed for backpropagation.

On the other hand, the Bayesian treatment for neural networks modifies the typical training pipeline where instead of performing point estimates, a distribution over parameters is inferred.
Although this approach adds complexity, it automatically accounts for model uncertainty. In turn, model averaging can be done to combat overfitting and improve calibration, especially on out-of-distribution data~\citep{zhang2018noisy, osawa2019practical}. 

% \begin{wrapfigure}[21]{r}{6.5cm}

\begin{figure}[t]
    \centering
    % \vspace{-2em}
    \phantom{bl} ODE-Net $\qquad \ \ \ \ \ \ \quad\quad$ SDE-BNN\\
    % \vspace{-0.2em}
    % \begin{subfigure}[b]{\linewidth} %trim={1.8cm 1.5cm 10cm 0.55cm},clip,
    \begin{subfigure}[b]{\textwidth}
    % \begin{subfigure}[b]{0.45\textwidth}
    \includegraphics[trim={0.3cm 0.5cm 0.2cm 0.5cm},clip,width=\linewidth]{toy/fig1_aug0_predictive.pdf}
    \end{subfigure}
% \vspace{-1.3em}
\caption{Hidden unit trajectories in an ODE-Net and an SDE-BNN. \textit{Left}: A continuous-depth residual network has deterministic transformations of its hidden units from depths $t=0$ to $t=1$. \textit{Right}: Uncertainty in the weights of a Bayesian continuous-depth residual network implies uncertainty in its hidden unit activation trajectories. Shaded regions show densities over samples from the learned posterior dynamics. \textit{Both}: Each distinct color corresponds to a different initial state corresponding to different data inputs.}%ercentiles are displayed for each initial state.}

\label{fig:dynamics}
\end{figure}

% \end{wrapfigure} 

How can we combine the benefits of continuous-depth models with those of Bayesian neural networks?
The simplest approach is a ``Bayesian neural ODE''~\citep{yildiz2019ode,dandekar2020bayesian}, which integrates out the finitely-many parameters of a standard neural ODE for prediction.

\begin{figure}[t]
    \centering
    % \vspace{-1.5em}
    % \begin{subfigure}[b]{0.9\linewidth}
    \begin{subfigure}[b]{\linewidth}
         \centering
         $\qquad \qquad$ Prior $\ \ \qquad\qquad$ Approximate posterior\\
         \includegraphics[clip, trim=1em 1em 1em 1em, width=\linewidth]{plots/toy_sdebnn.pdf}
         \label{fig:toy_prior_posterior}
     \end{subfigure}\\
    %  \vspace{-1em}
    \caption{Predictive prior and posterior of the SDE-BNN on a non-monotonic toy dataset.
    %\emph{Left}: Prior marginals. \emph{Right}: Posterior marginals. 
    %\emph{Both}: 
    Blue areas indicate density percentiles, and distinct colored lines show model samples.
    %$Left$: Our Bayesian SDE network gives an approximate posterior distribution on monotonic dataset ($x^3$). Percentiles $\pm 3$ standard deviations about the mean indicate the distribution over activations as states change through hidden units. $Right$:
    %Approximate posterior distribution for non-monotonic ($cos(x)$) problem given augmented inputs. $Both$: Percentiles indicate $\pm 3$ standard deviations above the mean. 
    }
    % \vspace{-0.3em}
    % \vspace{-1em}
    \label{fig:toypriorposterior}
\end{figure}

This approach is straightforward to implement, and can inherit the advantages of both Bayesian and continuous-depth neural nets.
However, empirically, standard Gaussian approximate posteriors are a relatively poor match for neural ODEs, not to mention the drawbacks of also being used in the prior.
Additionally, it does not exploit the special synergy available between continuous-time models and approximate inference.

In this paper, we show that an alternative construction of Bayesian continuous-depth neural networks has additional practical benefits.
Specifically, we consider the limit of infinite-depth Bayesian neural networks with separate unknown weights at each layer, a model class that we refer to as SDE-BNNs. 
% This construction, explored by \citet{peluchetti2020doubly,peluchetti2020infinitely},  corresponds to defining an SDE over the weights.
% A Bayesian treatment of this model class and its variational objective were characterized theoretically by \citet{tzen2019neural}.
We develop a unique network architecture that enhances model expressivity through time-correlated weights and scales linearly, instead of quadratically, with the parameter dimensionality. Combined with our novel formulation of a zero-variance gradient estimator, we show that approximate inference can be realized through the maximization of our modified variational lower bound, effectively scaling up the gradient-based variational inference scheme described by~\citet{li2020scalable} (preliminary forms of which appeared in earlier works~\citep{archambeau2008variational,opper2019variational,tzen2019neural}).

With this approach,  the state of the output layer is computed by a black-box adaptive SDE solver. %, and the model is  trained to maximize a variational lower bound.
Figure~\ref{fig:dynamics} contrasts our neural SDE with the % standard 
neural ODE parameterization.
This approach maintains the adaptive computation and constant-memory cost of training Bayesian neural ODEs and adds two unique benefits:
\begin{itemize}[topsep=0pt,leftmargin=*]
    \setlength\itemsep{0.2em}
    \item The variational posterior can be made arbitrarily expressive by simply enlarging the neural network that parameterizes the dynamics of the approximate posterior. Under mild conditions, this approach can approximate the true posterior arbitrarily closely.%family of approximate posteriors contains the true posterior.
    \item The variational objective admits a variance-reduced gradient estimator that is a natural extension of the ``sticking the landing'' trick~\citep{roeder2017sticking}. Combined with arbitrarily expressive approximate posteriors, it is consistent and has vanishing variance as the approximate posterior approaches the true.
\end{itemize}

Notably, our low-variance gradient estimator can also be applied to variational inference in SDEs more generally, such as for time-series modeling, but such applications are beyond the scope of this paper.
% \vspace{-1em}
\section{BACKGROUND} \label{sec:background}
% \vspace{-1em}
\paragraph{Bayesian Neural Networks}
%Neural networks define expressive families of functions. 
Given a dataset, there are often many functions that fit the data well, which a given neural network can express with different parameter values.
Instead of making point estimates of the parameters, the Bayesian paradigm frames learning as posterior inference. Predictions are obtained through integrating over many possible parameter settings. 
Formally, given a dataset $\D= \{(x_i, y_i) \}_{i=1}^N$ and prior distribution over model weights $p(w)$, we want to compute a posterior $p(w | \D) \propto p(\D | w) p(w)$. We can optimize an approximate posterior distribution $q(w)$ that minimizes the Kullback-Leibler (KL) divergence, i.e. maximizing the Evidence Lower Bound (ELBO):
\eq{
\label{eq:elbo}
\! \mathcal{L}_{\text{ELBO}}(\phi) =& 
\E_{q(w)} \sbracks{ \log p(\D| w) } - \KL\bracks{ q(w) || p(w) }.
}
% \vspace{-0.3em}
%However, exact inference is intractable, and the true posterior must be approximated.
Estimating gradients of this objective using simple Monte Carlo is known as stochastic variational inference (SVI)~\citep{hoffman2013stochastic, rezende2014stochastic}).
%turns posterior inference into a stochastic optimization problem.

One of the main technical challenges of SVI is choosing a parametric family of approximate posteriors that is tractable to sample from and evaluate, while being flexible enough to approximate the true posterior well. 
Most scalable inference techniques use Gaussian approximate posteriors with restricted covariance structure between network parameters
%Many works have investigated the construction of variational posteriors
~\citep{graves2011practical,blundell2015weight,zhang2018noisy,mishkin2018slang}. 
Others construct complex approximate posteriors with normalizing flows~\citep{krueger2018bayesian,louizos2017multiplicative} or through distillation~\citep{balan2015bayesian,wang2018adversarial}.

% \vspace{-4em}
\paragraph{Neural Ordinary Differential Equations}

\begin{figure}
% \vspace{-2.5em}
% \vspace{-4em} % for submission
% \begin{wrapfigure}[15]{R}{0.45\textwidth}
    \centering
    % \vspace{-1.5em}
    % \vspace{-4em} % submissio
    % \vspace{-2em} % submissio
    \begin{subfigure}[b]{0.05\linewidth}
        \rotatebox{90}{$z(t)$} \vspace{5em}
    \end{subfigure}
    \begin{subfigure}[b]{0.37\linewidth}
    % \begin{subfigure}[b]{0.35\linewidth}
    \includegraphics[trim={1.8cm 0.5cm 0 0},clip,width=\linewidth]{multimodal/no axis/fig_295_lightblue.pdf}
    \hspace{-2em}
    \vspace{0.5em}
    % {multimodal/no axis/fig_295.pdf}
    \phantom{blablab} $t$
    \end{subfigure}
    \begin{subfigure}[b]{0.05\linewidth}
    \rotatebox{90}{$z(t_1)$} \vspace{4.8em}
    \end{subfigure}
    \begin{subfigure}[b]{0.46\linewidth}
    \includegraphics[trim={1.5cm 1.5cm 0 0},clip,width=0.95\linewidth]
    {multimodal/no axis/joint_xlim295_t0.0_t1.0.pdf}\\
    % \\
    \hspace{-2em}
    \vspace{0.5em}
    \phantom{blabl} $z(t_0)$
    \end{subfigure}
    % \vspace{-1.2em}
    % \vspace{-2em}
\caption{Neural SDEs can learn arbitrarily expressive approximate posteriors.
\emph{Left:} Samples from an approximate posterior, trained with an OU prior and conditioned on two observations with Cauchy likelihoods.
\emph{Right:} Joint distribution and marginals of the approximate posterior process $z$ at times $t_0$ and $t_1$.}
\label{fig:multimodalcauchy}
\end{figure}

% \vspace{1em}
% \end{wrapfigure} 

% \vspace{3.5em}
Neural ordinary differential equations~\citep{chen2018neural} define ODEs using neural networks:
\eq{
\dee h_t = f_\theta(h_t, t) \dt, \quad h_0 \in \R^d,
}
where $f:\R^d \times \R \to \R^d$ is a Lipschitz function defined by a neural network with parameters $\theta$.
%The state $h$ of this ODE can itself be viewed as the hidden units of a continuous-depth neural network, where $t$ is analogous to the depth of the current layer.
%$indexes the depth of the network.
%
Starting at an initial value $h_0 = x$ given by a data example and integrating these dynamics forward for a finite time can be seen as passing the input through an infinitely-deep residual network.
For learning scalar-valued functions, adding extra dimensions to $h$ and a linear final layer induces similar universal approximations to 
% if extra dimensions are added to $h$, and the network is capped with a linear layer at the end, then these networks have similar universal approximation properties as 
standard neural networks~\citep{dupont2019augmented,zhang2019approximation} trained by standard stochastic gradient descent methods.
Using adaptive ODE solvers can trade evaluation speed for precision. The adjoint sensitivity method saves memory during training through reconstructing the trajectory of the hidden units $h$ by running the dynamics backwards in backpropagation.
% \vspace{-0.7em}
\subsection{Latent Stochastic Differential Equations}\label{ref:latentsde}
% \vspace{-0.5em}
Informally, an SDE can be viewed as an ODE with infinitesimal noise added throughout time. Formally:
% \vspace{-0.5em}
\begin{align}
\dee w_t = f_\theta(w_t, t)\dt+ g_\theta(w_t, t) \dB_t,
\label{eq:prior}
\end{align}
% \vspace{-0.5em}
where $w_0 \in \R^d$ is the initial state, 
%$\{w_t\}$ is the process of interest, 
$f_\theta:\R^d \times \R \to \R^d$ and $g_\theta:\R^d \times \R \to \R^{d \times m}$ are functions Lipschitz in both arguments, dubbed the \textit{drift} and \textit{diffusion}, respectively, and $\{B_t\}$ is an $m$-dimensional Brownian motion.

Some works have considered training SDEs with dynamics parameterized by neural networks~\citep{li2020scalable,tzen2019neural,peluchetti2020infinitely,innes2019zygote,kong2020sde,liu2019neural}. 
Note that directly optimizing the drift and diffusion to maximize the average log-likelihood of an observation $\log p(y_t | w_t)$ would result in the diffusion approaching $0$, conditional on the ODE fitting the training data well with the diffusion somewhat unconstrained.

Instead of directly optimizing the parameters of an SDE to match the data, a better approach is to use an SDE to define a prior over trajectories of $w$, and optimize the marginal likelihood of the data, integrating over all trajectories of $w$ weighted by the prior.
%condition on the data (by multiplying by the data likelihood), and optimize the 
%
%introduce a stochastic process over the states $\{w_t\}$ defined by an %SDE with the form:
%\eq{
%\dee {w_t} &= f_\theta({w_t}, t) \dt + g({w_t}, t) \dB_t,
%\label{eq:prior}
%}
Luckily, we can specify an approximate posterior over trajectories using a second SDE.
%conditioned on the data perform posterior inference for the state process using another process that induces the approximate posterior on path space. 
We define the approximate posterior by
% \vspace{-1.5em}
\begin{align}
\dw_t &= f_\phi(w_t, t) \dt + g_\theta(w_t, t) \dB_t. \label{eq:posterior}
\end{align}
% \vspace{-2em}
When the dynamics of the approximate posterior $f_\phi$ is parameterized by a neural network, this family of  posteriors is extremely expressive.
Figure~\ref{fig:multimodalcauchy} shows that such a variational family can easily approximate non-Gaussian and multi-modal posteriors on path space.

If both the SDE defined by~\eqref{eq:prior} and~\eqref{eq:posterior} share the same diffusion function, then the KL between the two induced measures on path space has the following form~\citep{li2020scalable,tzen2019neural}:
\vspace{-1.7em}
\begin{align}
\KL \bracks{ \mu_q || \mu_p } = 
    \mathbb{E}_{q_\phi(w)}\left[
        \int_0^1 \!\!\! \tfrac{1}{2} \norm{ u(t, \phi) }_2^2 \dt \right] \label{eq:kl} \textnormal{where} \\ %+ \!\!\! \int_0^1 \!\!\! u(t, \phi) \dB_t
    %  \textnormal{where} 
     \qquad u(t, \phi) = g_\theta(w_t, t)^{-1} \left[f_\theta(w_t, t) - f_\phi(w_t, t) \right] \label{eq:u}
\end{align}
% \vspace{-2em}

where $\mu_q$ and $\mu_p$ are path space probability measures induced respectively by~\eqref{eq:posterior} and~\eqref{eq:prior}, and the expectation is taken under the approximate posterior, denoted $q_\phi(w)$.
Intuitively, this KL divergence resembles the summative difference over time horizon [0, 1] between the prior drift $f_\theta$ and $f_\phi$, scaled by the diffusion.
%
This divergence can be estimated up to a constant with simple Monte Carlo, sampling trajectories from the dynamics given by the approximate posterior.
%The term being integrated over time is analogous to the KL between marginals of the prior and posterior processes, and is simply the difference in drifts scaled by the inverse diffusion.
% \vspace{-1em}
\paragraph{SDEs as expressive approximate posteriors}
To ensure that the KL divergence between the prior and approximate posterior on path space is finite, the same diffusion function $g_\theta(w_t, t)$ must be used for the approximate posterior and prior.
Surprisingly, this does not limit the expressivity of the approximate posterior. 
\citet{boue1998variational} show that there is a one-to-one correspondence between the space of path measures and drift functions that result in the same path space KL divergence. 
This implies that any path space measure close to the true posterior can be instantiated by SDEs with appropriate drifts. It follows that an approximate posterior parameterized by a sufficiently expressive family of function approximators can be made arbitrarily close to the true posterior.
The Girsanov reparameterization of the variational formula, derived from Boue \citep[Section 4]{tzen2019neural}, proves that the ELBO is tight when the drift is optimal. This means that there exists a ground truth drift function that can make the ELBO tight, the approximation of which can be achieved with a high capacity neural network.

% Similarly, \citet[Section 4]{tzen2019neural}, using the fact derived from Boue, characterize the Girsanov reparameterization of the variational formula, where the evidence lower bound is tight when the drift is optimal.
%This implies that if $u_t$ is the infimum over a increasingly larger space of processes then the the KL on path space becomes increasingly smaller.
% \vspace{-1em}
\section{INFINITELY DEEP BNNs}\label{sec:infinite bnns}
% \vspace{-1em}
Standard discrete-depth residual networks can be defined as a composition of layers of the form:
\begin{align}
    h_{t+\epsilon} = h_t + \epsilon f(h_t, w_t),  \quad t = 1 \dots T,
\end{align}
where $t$ is the layer index, $h_t \in \mathbb{R}^{D_h}$ denotes a vector of hidden unit activations at layer $t$, the input $h_0 = x$, and $w_t \in \mathbb{R}^{D_w}$ represents the parameters for layer $t$.
In the discrete setting, $\epsilon = 1$, $\in\mathbb{R}$.

We can construct a continuous-depth variant of residual networks by setting $\epsilon = \nicefrac{1}{T}$ and taking the limit as $T \rightarrow \infty$.
This yields a differential equation that describes the hidden unit evolution as a function of depth $t$.
Since standard residual networks are parameterized with different layerwise ``weights'', we denote them $w_t$.
To specify different weights at each layer with a finite number of parameters, we introduce a hypernetwork $f_w$ that specifies the change in weights as a function of depth and the current weights~\citep{ha2016hypernetworks}.
The evolution of the hidden unit activations and weights can then be combined into a single differential equation: 
\begin{align}
\frac{\dee}{\dt}
\begin{bmatrix}
h_t \\
w_t
\end{bmatrix} = 
\begin{bmatrix}
f_h(t, h_t, w_t) \\
f_w(t, w_t)
\end{bmatrix}
\end{align}
with some learned initial weight value $w_{t_0}$. Using time-varying weights is similar to augmenting the state~\citep{dupont2019augmented,anodev2}. See Appendix Figure~\ref{fig:posteriorpredictive} on the effects of augmentation.
We perform Bayesian inference on the weight process $w_t$, assigning a suitable prior stochastic process and performing variational inference in this infinitesimal limit.

Like all Bayesian neural networks with observation likelihoods, our framework models uncertainty both about parameters and about individual observations:
The likelihood $p(y | h_1)$ captures the observational noise, while the SDE encodes weight uncertainty. 
% \vspace{-1em}
\paragraph{Prior process on weights}
Typical priors for Bayesian neural networks are independent Gaussians across all weights and layers.
Taking the infinitesimal limit of such a prior gives a white noise process prior on the weights $w(\cdot)$.
However, initializing this noise while maintaining finite variance at scale is difficult~\citep{peluchetti2020doubly,peluchetti2020infinitely}.

Instead, we use the Ornstein–Uhlenbeck (OU) process as the prior on weights.
The process is characterized by an SDE with drift and diffusion:
\begin{align}
    f_p(w_t, t) = - w_t, \quad g(w_t, t) = \sigma I_d,
\end{align}
respectively, where $\sigma$ is a hyperparameter.
We choose this prior for its simplicity and bounded marginal variance at a constant in the large time limit.
% We choose this prior due to its simplicity and because its marginal variance approaches a constant in the large time limit, remaining bounded.
% \vspace{-1em}
\paragraph{Approximate posterior over weights}
We parameterize the approximate posterior on weights implicitly using another SDE with the following drift function:
\begin{align}
    f_q(w_t, t, \phi) &= \text{NN}_\phi(w_t, t, \phi) - f_p(w_t, t).
\end{align}
%
This drift $f_q$ is parameterized by a small neural network (NN) with parameters $\phi$.
% To ensure optimization starts at a stable set of dynamics, we also subtract the prior drift so that when the NN is initialized to output zero, the approximate posterior equals the prior.
With this drift, the approximate posterior process will generally have non-Gaussian, non-factorized marginals; its expressive capacity can be increased by making the neural net larger.

\paragraph{Evaluating the network}
% \vspace{-1em}
Given an input, we marginalize over weight and hidden unit trajectories.
This can be done with simple Monte Carlo, sampling a weight path $\{w_t\}$ from the posterior process and evaluating the network activations $\{h_t\}$ given the sampled weights and input.
Both steps require solving a differential equation.
Luckily, both can be solved simultaneously with the augmented state SDE:
\begin{align}\label{augsde}
    \dee \begin{bmatrix}
    w_t \\
    h_t
    \end{bmatrix} = 
    \begin{bmatrix}
    f_w(w_t, t, \phi) \\
    f_h(h_t, t,  w_t)
    \end{bmatrix}  \dt + 
    \begin{bmatrix}
    g_w(w_t, t)\\
    \mathbf{0}
    \end{bmatrix} \dB_t, 
\end{align}
where $h_0 = x$, the input. 
The learnable parameters are the initial weight values at time zero $w_0$ (either point estimated or inferred) and those of drift function $\phi$.
% \vspace{-0.5em}
\paragraph{Output likelihood}
The final state of the hidden units $h_1$ is used to parameterize the likelihood of the target output $y$:
 $\log p(y | x, w) = \log p(y | h_1)$.
For instance, $p(y | h_1)$ could be a Cauchy likelihood for regression, or categorical likelihood for classification.
% \vspace{-0.5em}
\paragraph{Training objective}
To fit the network to data, we maximize the lower bound on marginal likelihood given by the infinite-dimensional ELBO:
% \begin{align}
% &\mathcal{L}_{\text{ELBO}_\infty}(\phi)  
% = \Exp { \log p(\mathcal{D} | w) } - \KL(\phi) \\
% =& \E \bigg[
%     \log p(\mathcal{D} | w) \!-\!\!
%     \int_0^1 \!\! \tfrac{1}{2}
%         \norm{ \text{NN}(w_t, t, \phi)  g_w(w_t, t)^{-1} }_2^2 \dt \bigg] \nonumber
%     \nonumber \label{eq:sde_elbo}
% \end{align}
\begin{equation}
\mathcal{L}_{\text{ELBO}_\infty}(\phi)  
%= \mathbb{E}_{q_\phi(w)} \left \log p(\mathcal{D} | w) \right] - \KL(\phi)
= \mathbb{E}_{q_\phi(w)} \left[
    \log p(\mathcal{D} | w) \!-\!\!
    \int_0^1 \!\! \tfrac{1}{2}
        \norm{u(w_t, t, \phi)}_2^2 \dt \right]. \nonumber
    \nonumber \label{eq:sde_elbo}
\end{equation}
% \begin{equation*}
% \begin{multlined}
% \begin{aligned}
% \mathcal{L}_{\text{ELBO}_\infty}(\phi)  
% = \Exp { \log p(\mathcal{D} | w) } - \KL(\phi)\\
% = \E \bigg[
%     \log p(\mathcal{D} | w) \!-\!\!
%     \int_0^1 \!\! \tfrac{1}{2}
%         \norm{ \text{NN}(w_t, t, \phi)  g_w(w_t, t)^{-1} }_2^2 \dt \bigg] \nonumber
%     \nonumber \label{eq:sde_elbo}
% \end{aligned}
% \end{multlined}
% \end{equation*}
% where $u(w_t, t, \phi) = (f_q(w_t, t, \phi) - f_p(w_t, t)) / \sigma$.
The sampled weights, the hidden activations, and the training objective are all computed simultaneously with a single call to an adaptive SDE solver.
Gradients of the sampled loss can also be efficiently computed using adaptive solvers, following \citet{li2020scalable}.

% \begin{wrapfigure}{r}{6.25cm}
%     \centering
%     \begin{subfigure}[b]{\linewidth}
%         %  \centering
%         %  $\qquad \quad\quad$ Prior $\qquad$ Approximate posterior\\
%         %  \includegraphics[clip, trim=1em 1em 1em 1em, width=\linewidth]{plots/toy_sdebnn.pdf}
%         %  \label{fig:toy_prior_posterior}
%      \end{subfigure}\\
%     %  % \vspace{-2em}
%     % \caption{Predictive prior and posterior of the SDE-BNN on a non-monotonic toy dataset.
%     % Blue areas indicate density percentiles, and distinct colored lines show model samples.
%     % }
%     % \vspace{-0.3em}
%     \label{fig:toypriorposterior}
% % \end{figure}
% \end{wrapfigure}
% \vspace{-1em}
\section{VARIANCE-REDUCED GRADIENTS}\label{sec:stl}
% \vspace{-1em}
\citet{roeder2017sticking} showed that when optimizing expectations using the reparameterization gradient, a gradient estimator with lower variance can be constructed by removing a score function term that has zero expectation, and that the variance of this gradient estimator approaches zero as the approximate posterior approaches the true posterior.
We refer to this trick as ``sticking the landing'' (STL). 
We generalize this to our SDE setting by replacing the original estimator of the path space KL with the following STL estimator:
\eq{
\widehat{\text{KL}}_{\text{STL}} \!=\! 
    \int_0^1
        \!\! \tfrac{1}{2} \norm{ u(w_t, t, \phi) }_2^2 \dt \! + \!\!
    \int_0^1
        \!\! u(w_t, t, \bot(\phi)) \dB_t, 
        % \qquad 
        % w(\cdot) \sim q_\phi(w)
\label{eq:stl}
}
where $w(\cdot) \sim q_\phi(w)$ and $u$ is defined in \eqref{eq:u}, the path $\{w_t\}_{t \in [0, T]}$ is sampled from the approximate posterior process, and $\bot(\cdot)$ is the stop gradient function that renders the input a constant with respect to which gradient propagation is stopped. Note that \ref{eq:stl} is the fully Monte Carlo version referred to in \ref{fig:toy_stl} from which our STL variant is derived.
% See \citep[Appendix 9.5]{li2020scalable} for a derivation.

The second term in~\eqref{eq:stl} is a martingale and has expectation zero. Therefore, in prior works~\citep{li2020scalable,tzen2019neural,tzen2019theoretical}, Monte Carlo estimation was only performed for the first term, but we find that this approach does not necessarily reduce the variance of the \emph{gradient}~(Figure~\ref{fig:toy_stl}).

\begin{figure}[t]
% \begin{wrapfigure}[16]{L}{6.3cm}
    \centering
    % \vspace{-0.5em}
    \begin{subfigure}[b]{\linewidth}
    % \begin{subfigure}[b]{0.9\linewidth}
    \centering
    \includegraphics[width=\linewidth, trim=2px 2px 0 2px, clip]{plots/grad_var_avg.pdf}
    \end{subfigure}
    % \vspace{-1.5em}
    \caption{Comparison of the variance in three gradient estimators. On this toy problem, our new gradient estimator reduces variance by a factor of roughly 4.}
    \label{fig:toy_stl}
% \end{wrapfigure}
\end{figure}
%
Because our approximate posterior can be made arbitrarily expressive, we conjecture that our approach can achieve arbitrarily low gradient variance towards the end of training if the $f_w$ parameterization is expressive enough. See Appendix~is \ref{app:stlheuristic} for a heuristic derivation.

We show the variance of different gradient estimators in Figure~\ref{fig:toy_stl}, averaged across the parameters $\theta$, in a 1D regression setting. We compare STL against a ``Full Monte Carlo'' estimate which includes the second additional term without gradient stopping, as well as the estimator that was previously used by~\citet{li2020scalable} which ignores the second term.
Figure~\ref{fig:toy_stl} shows that STL obtains lower variance than alternatives, when matching an exponentiated Brownian motion.
Table~\ref{tab4:stl_improvement} shows training performance improvements.
% \vspace{-1em}
\section{EXPERIMENTS}\label{sec:experiments}
% \vspace{-1em}
We investigate the effectiveness of our proposed approximate inference method for training continuous-depth neural nets, referred to as SDE-BNN, in terms of classification accuracy, calibration, perturbation robustness, and speed-precision trade-offs. 
Our code is publicly available
\href{https://github.com/xwinxu/bayeSDE}{here} and experimental settings in Table~\ref{tab:hparam settings}.

\newcommand{\tcen}[1]{\multicolumn{1}{c}{#1}}
\newcommand{\tnan}{---}
\begin{table*}\centering
\ra{1.1}
\setlength{\tabcolsep}{0.8em}
\resizebox{0.75\textwidth}{!}{
\begin{tabular}{@{} l c c c c c c @{}}\toprule
  & \multicolumn{2}{c}{\textbf{\textsc{MNIST}}} &
  & \multicolumn{2}{c}{\textbf{\textsc{CIFAR-10}}} \\
\cmidrule(lr){2-3} \cmidrule(l){5-6} 
Model & \tcen{Accuracy (\%) } & \tcen{ECE ($\times 10^{-2}$)}
& & \tcen{Accuracy (\%)} & \tcen{ECE ($\times 10^{-2}$)} \\
\cmidrule(r){1-1}\cmidrule(lr){2-3} \cmidrule(l){5-6} 
ResNet32 & 99.46 $\pm$ 0.00 & 2.88 $\pm$ 0.94 & & 87.35 $\pm$ 0.00 & 8.47 $\pm$ 0.39 \\
ODEnet & 98.90 $\pm$ 0.04 & 1.11 $\pm$ 0.10 & & 88.30 $\pm$ 0.29 & 8.71 $\pm$ 0.21 \\
HyperODEnet & 99.04 $\pm$ 0.00 & 1.04 $\pm$ 0.09 & & 87.92 $\pm$ 0.46 & 15.86 $\pm$ 1.25 \\
\cmidrule(r){1-1}\cmidrule(lr){2-3} \cmidrule(l){5-6} 
MFVI ResNet32 & 99.44 $\pm$ 0.00 & 2.76 $\pm$ 1.28 & & 86.97 $\pm$ 0.00 & 3.04 $\pm$ 0.94 \\
MFVI$^\dagger$ & ---  & --- & & 86.48 \phantom{$\pm$ 0.00} & 1.95 \phantom{$\pm$ 0.00} \\
Deep Ensemble$^\dagger$ & --- & --- & & 89.22 \phantom{$\pm$ 0.00} & 2.79 \phantom{$\pm$ 0.00} \\
HMC (``gold standard'')$^\dagger$ & 98.31 \phantom{$\pm$ 0.00}  & 1.79 \phantom{$\pm$ 0.00} & & 90.70 \phantom{$\pm$ 0.00} & 5.94 \phantom{$\pm$ 0.00} \\ 
\cmidrule(r){1-1}\cmidrule(lr){2-3} \cmidrule(l){5-6} 
MFVI ODEnet & 98.81 $\pm$ 0.00 & 2.63 $\pm$ 0.31 & & 81.59 $\pm$ 0.01 & 3.62 $\pm$ 0.40 \\ 
MFVI HyperODEnet & 98.77 $\pm$ 0.01 & 2.82 $\pm$ 1.34 & & 80.62 $\pm$ 0.00 & 4.29 $\pm$ 1.10 \\
SDE BNN & 99.30 $\pm$ 0.09 & 0.63 $\pm$ 0.10 & & 89.84 $\pm$ 0.94 & 7.19 $\pm$ 0.37 \\
SDE BNN (+ STL) & 99.10 $\pm$ 0.09 & 0.78 $\pm$ 0.12 & & 89.10 $\pm$ 0.45 & 7.97 $\pm$ 0.51 \\
% SDE BNN $w0$ inferred & 99.04 $\pm$ 0.03 & 0.73 $\pm$ 0.04 & &  ---  &  ---  \\
% SDE BNN $w0$ inferred (+ STL) & 99.05 $\pm$ 0.00 & 0.79 $\pm$ 0.04 & &  ---  &  ---  \\
\bottomrule
\end{tabular}
}
\caption{Classification accuracy and expected calibration error (ECE) on MNIST and CIFAR-10. 
We separate models into point estimates, discrete-time models, and continuous-time models.
Our SDE-BNN outperforms other continuous-time Bayesian neural nets (BNNs) and perform competitively against discrete-time BNNs.
% Note that ResNet32 was run without Batch Normalization (BN) for fairer comparison with all mean field and SDE models. 
% No batch normalization was used as it is unclear how it can be applied to continuous-depth models.
% Our approach (SDE BNN and SDE BNN (+STL) ) is either competitive or outperforms baselines (ODENet), while obtaining better generalization as seen in lower expected calibration error (ECE). 
%This suggests the utility of trading off ease of training (i.e. maximum likelihood) for learning richly parameterized, yet cheap, approximate posteriors towards better calibration, robustness, and generalization.
$^\dagger$Results by \citet{izmailov2021bayesian} where a modified residual network architecture was used; only one seed was reported.
}
\label{tab1:img_results}
\end{table*}

We consider toy regression and image classification tasks on MNIST and CIFAR-10.
%to validate the usage of an infinite-dimensional prior and our variational inference objective.
We also investigate out-of-distribution generalization.
%the continuous-depth parameterizations.
Notably, our approach does not require \emph{post hoc} recalibration methods such as training with temperature scaling \citep{guo2017calibration} or isotonic regression \citep{zadrozny2002transforming}.
% \vspace{-2em}
% \vspace{-0.2em}
\paragraph{Backpropagation through solvers vs. adjoint}
We experimented with fixed- and adaptive-step SDE solvers, and the stochastic adjoint of~\citet{li2020scalable}.
Figure~\ref{fig:bpadjoint} shows similar convergence for both approaches.  Appendix~\ref{app:additional_figures} shows that both had similar numbers of dynamics function evaluations and wall-clock time.

% Investigation into alternative priors is beyond the current scope.

\begin{figure}
% \begin{wrapfigure}[21]{L}{6cm}
% \vspace{-0.5em}
% \vspace{-0.8em}
    \centering
    \vspace{0em}
    \begin{subfigure}[b]{0.9\linewidth}
         \centering
         \includegraphics[width=\linewidth]{torchsde/plot_backprop_vs_adjoint_ep100_val_xent_train_loss.pdf}
        %  \label{fig:mnist}
     \end{subfigure}\\
    %  \vspace{-1.4em}
    % \vspace{0.4em}
    \caption{Benchmarking two gradient computation methods: (1) Back-propagation through the SDE solver, and (2) the memory-efficient stochastic adjoint of \citet{li2020scalable}.
    Both methods have similar optimization dynamics, final performance, and wall-clock time, but the adjoint approach is more memory-efficient.
    %but the adjoint method has better speed, averaging 135.91 vs 162.59 sec/epoch. 
    Detailed comparisons of wall-clock time and evaluation step results in Appendix~\ref{app:adjointadaptivesolver}.}% Average across 3 seeds.} % avg backprop: 162.58540952497057 | avg adjoint: 135.90531070722474 sec / epoch | relative: 1.1963138797071857
    % \vspace{-1.8em}
    \label{fig:bpadjoint}
\end{figure}
% \end{wrapfigure}

The overhead for estimating error in our adaptive solvers was substantial; therefore, for final model evaluation, we trained with fixed-step solvers, where the number of steps is chosen to be large enough to match the convergence speed of our adaptive-step solvers.
%\vspace{-1em}

\paragraph{Baselines}
For a fixed-depth network baseline, we compare to standard residual networks.
We then test variational inference on the weights of these models.

We also perform ablation studies to compare with standard variational inference approaches over continuous-depth networks.
Specifically, we compare to a mean field variational inference (MFVI) ODEnet where stochastic variational inference is performed over depth-invariant weights.
This baseline is a fully-factorized Gaussian approximate posterior, \ie mean-field approximation, and been used for Neural ODEs by \citet{look2019differential,dandekar2020bayesian}.

We further compare our model to a MFVI HyperODEnet, where a learned drift is applied to $w$, but mean-field inference is instead performed over the parameters of the hypernetwork. 
Alternatively, one can interpret this as another MFVI ODEnet with a larger state and a more complex drift function but with similar computational complexity to SDE-BNN.
This setting contrasts our approach of doing Bayesian inference over the entire continuous-depth network as a stochastic process. 
\vspace{-1em}
\paragraph{Parameterizing the drift function} 
\vspace{-1em}
%The (hyper)network for the drift function of the variational posterior $f_w$ required special handling.
We parameterized the drift of the variational posterior $f_w$ with a simple multilayer perceptron.
% To make training easier, we added the prior drift to the output of this multilayer perceptron. This results in a zero initialization of the last layer, which corresponds to an initialization of the approximate posterior SDE to the prior SDE.
To ensure optimization starts at a stable set of dynamics, we subtract the prior drift so that the approximate posterior equals the prior when the final layer is initialized to output zero.

% \vspace{-1em}
\paragraph{Hyperparameters}
We swept learning rates in the range [1e-4, 1e-3], selecting the optimal based on the validation set. We train with the default Adam optimizer~\citep{Kingma2015AdamAM}.
In image classification experiments, all convolutional layers of the drift network are time-conditional and use the \texttt{tanh} non-linearity. The diffusion coefficient $\sigma$ was selected from validation performance over \{0.1, 0.2, 0.5\}.

% \vspace{-0.5em}
% \vspace{-1em}
\subsection{1D Regression}
% \vspace{-1em}
We first verify the capabilities of the SDE-BNN on a 1D regression problem. 
Conditioned on a sample from the diffusion process, each sample from a one-dimensional SDE-BNN is a bijective mapping from the inputs to the outputs.
This implies that every function sampled from a 1D SDE-BNN is monotonic.
To be able to sample non-monotonic functions, we augment the state with 2 extra dimensions initialized to zero, as in~\citet{dupont2019augmented}.
Figure~\ref{fig:toypriorposterior} shows that our model learns a reasonably flexible approximate posterior on a synthetic non-monotonic 1D dataset.
We emphasize that the samples from our model are smooth w.r.t. depth because the hidden states $h$ do not receive additive instantaneous noise, only the weights $w$ do.
% Only on the weights $w$ do we apply instantaneous noise.
% \vspace{-1em}
\subsection{Image Classification}
% \vspace{-1em}
\begin{figure}[htb]
    \centering
    % \vspace{-1.55}
    \begin{subfigure}[b]{\textwidth}
    % \begin{subfigure}[b]{0.52\textwidth}
        \centering
        \includegraphics[width=0.475\linewidth,trim={1.15cm 0.45cm 0.53cm 0.7cm},clip]{cifar10/sdebnn/pareto/loss_inftime_CIFAR10_paretotest_mean.pdf}%
        \hfill
        \includegraphics[width=0.475\linewidth,trim={1.17cm 0.42cm 0.3cm 0.7cm},clip]{mnist/sdebnn/pareto/ece_inftime_MNIST_paretoval_mean.pdf}
        \caption{CIFAR-10. \textit{Left:} Negative log likelihood. \textit{Right:} ECE. Adjusting SDE-BNN solver tolerance at test time trades off computational speed for predictive performance. Grey line is solver's training tolerance. Averaged across 3 seeds.}
        \label{fig:pareto}
    \end{subfigure}
    \hspace{0.3em}
    % \vskip\baselineskip
    \begin{subfigure}[b]{\textwidth}
    % \begin{subfigure}[b]{0.45\textwidth}
        \centering
        \includegraphics[width=0.49\linewidth]{cifar10/odenet/test/test-calib0-20.pdf}%
        % \hfill
        \hspace{0.1em}
        \includegraphics[width=0.49\linewidth]{cifar10/sdebnn-stl-ou/calib/test-calib1715-20.pdf}
        % \vspace{-1.8em}
        \caption{Calibration on the CIFAR-10 test set for a neural ODE (left) and a SDE-BNN (right). The SDE-BNN displays better calibration and generalization.}
        \label{fig:cifarcalibration}
    \end{subfigure}
    \caption{Performance of SDE-BNN on standard CIFAR-10 classification task.}
    \label{fig:imgclassplots}
    % \vspace{-0.7em} % submission
\end{figure}
% \vspace{-3em}
% Next, we evaluate the performance of SDE-BNN on standard MNIST and CIFAR-10 classification tasks.
Instantaneous changes to the hidden state ($f_h$) are parameterized using a convolutional neural network, including one strided convolution for downsampling and a transposed convolution layer for upsampling. 
We then set the $w$ to be the filters and biases of all the convolutional layers.
The approximate posterior drift dynamics ($f_w$) is a multilayer perceptron with hidden layer widths of 2, 128, and 2. The small hidden width of the bottleneck layers was chosen to reduce the number of variational parameters and promote linear scaling with respect to the dimension of $w$.
On MNIST, we used one such SDE-BNN block, while on CIFAR-10, we used a multi-scale variant where multiple SDE-BNN blocks were stacked with the invertible downsampling from \citet{dinh2016density} in between.

We report classification results in Table~\ref{tab1:img_results}.
Our SDE-BNN generally outperforms the baselines. While the continuous-depth Neural ODE (ODEnet) models can achieve similar classification performance on a standard residual network, it consistently has poorer calibration. 

The SDE-BNN matches and outperforms the accuracy of standard residual networks on MNIST and CIFAR-10, respectively, while obtaining lower expected calibration errors (ECE).
From ablation studies, we found that it was harder to achieve similar performance with either of the mean field variants of an ODEnet as they had a poorer trade-off between performance and calibration.

Figure~\ref{fig:pareto} shows the ability of SDE-BNNs to trade off computation time for precision. 
Figure~\ref{fig:cifar10_calib_spectrum_inference} in Appendix~\ref{app:calibration}  indicates that calibration is insensitive to solver tolerances close to the value used during training.

% \vspace{-1em}
\subsubsection{Calibration}
% \vspace{-0.5em}
%A well calibrated model is able to forecast its accuracy with a corresponding level of confidence. 
Table~\ref{tab1:img_results} quantifies our model's calibration with expected calibration error (ECE;~\citet{guo2017calibration}). 
The SDE-BNN appears better calibrated than the Neural ODE \citep{chen2018neural} and mean field ResNet baselines.
Figure~\ref{fig:cifarcalibration} shows better calibration than neural ODEs with similar accuracy.
Appendix Figure~\ref{fig:cifar10_calib_spectrum} shows the insensitivity of these results to solver step size.
\begin{figure}
% \begin{wrapfigure}[18]{L}{0.48\linewidth}
    \centering
    % \vspace{-2em}
    \begin{subfigure}[b]{0.8\linewidth}
         \centering
         \includegraphics[width=\linewidth]{cifar10c/severity vs metric/error_bar-fig4_severity_error_bar.pdf}
        %  \caption{Classification error.}
         \label{fig:corruption_accs}
     \end{subfigure}\\
     % \vspace{-3em}
     \begin{subfigure}[b]{0.8\linewidth}
         \centering
         \includegraphics[width=\linewidth]{cifar10c/severity vs metric/error_bar-fig5_severity_ece_bar.pdf}
        %  \caption{Calibration Error.}
         \label{fig:corruption_eces}
     \end{subfigure}
     % \vspace{-2.2em}
    %  \vspace{-2.5em}
    \caption{CIFAR10-C. Robustness to distributional shifts on CIFAR-10. SDE-based neural nets show better accuracy and calibration than non-Bayesian and mean-field methods. Black bars show standard deviation over 3 seeds.}
    \label{fig:severitycorruption}
\end{figure}
% \end{wrapfigure}
% \vspace{-5em}
% \vspace{-1em}
\subsubsection{Robustness to Input Corruption}
% \vspace{-1em}
We report the robustness of SDE-BNNs by evaluating on all 19 non-adversarial corruptions across 5 severity levels in  CIFAR10-C~\cite{hendrycks2019robustness}.
These corruptions mimic real-world perturbations such as noise, blur, and weather.
To evaluate the classification robustness of SDE-BNN, we compare the mean corruption error (mCE), an average error for each intensity level summed across all 19 perturbations, to the top-1 error rate on the corresponding clean CIFAR-10.

Figure~\ref{fig:severitycorruption} shows error on the corrupted test set relative to uncorrupted data, demonstrating a steady increase in mCE across increasing perturbation levels along with the overall error measurement summarized in Table~\ref{tab1:img_results}. On both CIFAR-10 and CIFAR10-C, the SDE-BNN and SDE-BNN + STL models achieve lower overall test error and better calibration than the baselines.

Compared to standard baselines (ResNet32 and Mean Field (MF) ResNet32), SDE-BNN achieves around 4.4\% lower absolute corruption error (CE), the total classification error for all corruption tasks across all 5 severity levels \citep{hendrycks2019robustness}, in comparison to the clean errors. The effectiveness of learned uncertainty on out-of-domain inputs indicates that SDE-BNN is more robust to observation noise despite not being trained on such diverse forms of corruptions. 

% \vspace{-1em}
% \todo{add vspace if necessary}
\section{SCOPE AND LIMITATIONS}\label{sec:scopelimitations}
% \vspace{-0.3em}
% \vspace{-1em}

\paragraph{Computational speed}
The cost of evaluating our model grows in $\mathcal{O}(DT)$, where $D$ is the number of weights, and $T$ the number of iterations taken by the solver.
This may seem advantageous compared to the $\mathcal{O}(D^3)$ cost for non-factorized Gaussian approximate posteriors, but the number of steps required is difficult to characterize.
%
Although our approach allows adjustment of the computational cost at test time, it is harder to control the cost of evaluation during training time, making our method relatively slow to train.
However, it should be straightforward to regularize these models to be faster to solve, as in \citet{kelly2020easynode}.
Relatedly, \citet{dusenberry2020efficient} recently demonstrated an $\mathcal{O}(DK)$ cost approximate posterior in standard BNNs.
% \vspace{-2em}

\paragraph{Batch norm}
% There are two places in our model at which batch normalization (BN) could be applied:  (1) in the weight hypernetwork, and (2) in the dynamics of the hidden unit activations.
We did not incorporate batch normalization~\citep{ioffe2015batch} in any of our neural network components.
Introducing any normalization (e.g. batchnorm, layernorm, etc.) compromises the Lipschitz property required for SDEs to have a unique solution. 
Since BN introduces dependence between samples within a batch, it is also unclear how to incorporate BN while maintaining the consistency properties of Bayesian inference.
\citet{zhang2019fixup, Chang2020PrincipledWI} proposed initializations that yield the same performance without needing batch normalization.
% \vspace{-1em}

\paragraph{Low-variance gradients for other domains}
Our extended STL gradient estimator~\citep{roeder2017sticking} to the infinite-dimensional variational objective could be applied to other settings for faster convergence, e.g. time series applications \citet{li2020scalable} investigated.
% \vspace{-1em}

\section{RELATED WORK}\label{sec:related work}
% \vspace{-1em}
\paragraph{Initial theoretical investigations}
The earliest theoretical treatment of infinitely-deep Bayesian neural networks was made by \citet[Chapter 2]{neal2012bayesian}, but no practical training or evaluation method was proposed.
%, who also characterized infinitely-wide Bayesian neural networks as Gaussian processes.
\citet{DuvRipAdaGha2014} also investigated the theoretical properties of kernel-based constructions of infinitely-deep Bayesian neural networks.
% \vspace{-1em}

\paragraph{Diffusion limits of discrete-time models}
We expect existing discrete-depth constructions to converge to diffusion limits in the infinitesimal limit if a system is updated with appropriately scaled Gaussian noise at each timestep. 
\citet{peluchetti2020infinitely,peluchetti2020doubly} show this holds for the output of residual networks with shallow residual blocks whose weight initializations are appropriately scaled. 
While our construction of SDE-BNN given by~\eqref{augsde} seems similar \cite{peluchetti2020infinitely}, there are two key differences: (i) We strictly enforce hidden states to follow a diffusion throughout training by directly learning a neural SDE, whereas~\citet{peluchetti2020infinitely} only ensures SDE-driven dynamics at initialization. (ii) We adopt a more general neural net architecture for the residual blocks than the shallow ones considered in~\citep{peluchetti2020infinitely}. Their work mainly discusses the convergence of shallow ResNets to SDEs, in order to analyze training stability for regular ResNets and verify that a scaled gradient formulation leads to faster convergence at the first epoch.
The consequence of (i) is that operations on diffusions (e.g., computing path-space KL) remain applicable even after our model has been trained.
While (ii) appears to be a minor difference, it actually uncovers a fundamental distinction in our analysis: Since we start out with an SDE, and only discretize for numerical computations, our model is able to incorporate any type of Lipschitz smooth residual block. Additionally, no training algorithm was specified for learning SDE models.
The analysis by \citet{peluchetti2020infinitely} relies on Taylor expanding the residual block function, which is not easy in the presence of complex residual block architectures and would require modifications to the initialization.
%when such alternate architectures are employed.
\citet{tzen2019neural} show that particle trajectories of the approximate posterior in discrete deep latent Gaussian models converge to a diffusion, and that the ELBO may be written with KL of measures on path space.
This construction has been explored in various forms in the past~\citep{opper2019variational, archambeau2008variational}.
%, with a more general formulation encapsulating non-constant diffusions and practical implementation applied to time series data by~\citet{li2020scalable}. 
% \vspace{-1em}
%Our work applies specifically this to Bayesian learning of neural nets.
%Our work introduces a new variant of SDE-BNN, shows that the scalable gradient-based stochastic variational inference of \cite{li2020scalable} applies in this setting, 
%Our demonstrates the special flexibility of approximate posteriors in this setting, and derives a new gradient estimator for this approach that takes advantage of this flexibility to achieve lower variance.

% \paragraph{Theory of Neural SDEs}
% \citet{peluchetti2020doubly,peluchetti2020infinitely} derive correspondences between infinitely deep stochastic neural networks and SDEs, but do not treat them as Bayesian neural networks.
% Specifically, \citet{peluchetti2020doubly} derived the drift and diffusion functions corresponding to a particular randomly initialized convolutional residual network, and explored the properties of infinite-wide variants of this related model class.
% In contrast, we use a Bayesian approach, tracking the uncertainty arising from limited data.
% \citet{tzen2019theoretical} formulated neural SDEs from a latent variable model, and their connections to deep neural networks was made precise in \citet{tzen2019neural}.
% Specifically, \citet{tzen2019neural} proposed the use of the Radon-Nikodym derivative as 
% an infinite-dimensional analog to the variational evidence lower bound (ELBO) objective. 
% However their approach relied on forward-mode differentiation, so did not scale to large models.

\paragraph{Neural SDEs with other training objectives} Models making use of SDEs have appeared in the past, though many make use of somewhat \textit{ad-hoc} combinations of methods involving both discrete and continuous components.
\citet{kong2020sde} proposed fitting a neural SDE by using a heuristic training objective based on encouraging the diffusion to be large away from the training data and a fixed Euler-Maruyama (E-M) discretization.
\citet{innes2019zygote} trained neural SDEs by backpropagating through the operations of the solver, however their training objective simply matched the first two moments of the training data, implying that it could not consistently estimate diffusion functions.
This approach is also relatively memory-intensive.
\citet{liu2019neural} and \citet{oganesyan2020stochasticity} add noise to the solver operations in a neural ODE, although the diffusion must be tuned as a hyperparameter.
\citet{hegde2018deep} proposed a form of neural SDE using Gaussian processes to parameterize the drift and diffusion functions for a fixed E-M discretization.  However, the diffusion functions are based on an \emph{ad-hoc} construction from a Gaussian process posterior conditioned on inducing points.
\citet{ryder2018blackbox} used a Gaussian process variational posterior, effectively a continuous-time analog of a mean field approximation that may not always be expressive enough to model the true posterior.
\citet{kidger2021neural} learn neural SDEs by jointly learning a discriminator~\citep{kidger2020neural} and formalize the problem as learning generative adversarial networks. However, this would involve many more hyperparameters and require extensive tuning compared to our variational inference approach.
% \vspace{-0.5em}
\paragraph{ODEnets with finite-dimensional stochasticity} Some methods based on building variational autoencoders with a neural ODE share similar training objectives, since the ELBO appears frequently in posterior inference. The Latent ODE model~\citep{rubanova2019latent} only performs inference on the distribution at an initial time of a continuous hidden state. %\citet{rubanova2019latent} and 
\citet{de2019gru} introduced stochastic jumps at data locations, and do not perform continuous-time inference. While performing amortized inference for time series modeling, \citet{yildiz2019ode} also infer the weights of an ODE drift function.
\citet{dandekar2020bayesian} have a similar setting but for supervised learning.
% \vspace{-0.5em}
\paragraph{Approximate posteriors defined as neural nets}
\citet{krueger2018bayesian} and \citet{louizos2017multiplicative} use normalizing flows to construct an unfactorized, non-Gaussian approximate posterior in BNNs.
However, normalizing flows have poor scaling with dimension and point estimates were used for most of the weights in the neural network.
Table~\ref{tab5:comparemethods} in Appendix~\ref{tab5:comparemethods} compares qualities of our approach to existing methods for stochastic variational inference in BNNs.
% % \vspace{-1em}
% \paragraph{Models that Denoise a Diffusion.} \citet{song2020score} build a generative model using an encoder-decoder pair of SDEs defined implicitly by the gradients of an energy function. They then use the Fokker-Planck equation to directly update their model without differentiating through sample paths.
% \improvement{This is not precise.}

\section{CONCLUSION}
% \vspace{-1em}
We developed a practical method for approximate inference in continuous-depth Bayesian neural networks. Our approach exploits a special synergy between continuous-depth models and variational inference for SDEs, providing additional benefits over standard approaches.
In particular, our method allows arbitrarily-expressive, non-factorized approximate posteriors implicitly defined through neural SDEs.
We also developed an unbiased gradient estimator for SDE variational inference whose variance approaches zero as the approximate posterior approaches the true posterior.
This combination gives our family of Bayesian continuous-depth neural networks a special property, which is that the gradients' bias and variance can be made arbitrarily small during training.
%
Where standard applications of MFVI on continuous-depth models perform poorly, our approach brings continuous-depth Bayesian neural networks to a comparable performance with standard Bayesian neural networks.
Furthermore, we demonstrated the ability of this continuous-depth model class to use adaptive SDE solvers.
This allows a memory-efficient training, and a fine-grained trade-off between precision and speed.

% \section{Broader Impact}
% We expect the main impact of this work to be towards the improvement of models with better calibration and predictive uncertainties that would be conducive to better signals regardless of the domain of application in which they are used. The societal impacts aside from the general of deep learning approaches and ethical concerns of neural networks are not immediately foreseeable.

\subsubsection*{Acknowledgements}
We thank Jesse Bettencourt, Radford M. Neal, and Patrick Kidger for helpful technical discussions and revisions on earlier drafts of this work. We also thank James Bradbury for his support while implementing differential equation solvers in JAX.

% \newpage
\bibliographystyle{apalike}
\bibliography{references}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% \newpage
% \input{checklist}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% \newpage
% \appendix

% \section{Appendix}
% Optionally include extra information (complete proofs, additional experiments and plots) in the appendix.
% This section will often be part of the supplemental material.

% \clearpage
% \appendix

% \thispagestyle{empty}
% \onecolumn \makesupplementtitle

\input{appendix.tex}


\end{document}

math_commands.tex
%%%%% NEW MATH DEFINITIONS %%%%%

\usepackage{amsmath,amsfonts,bm}

% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}

% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}


% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}

\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}

\def\eps{{\epsilon}}


% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}

% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}

% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}

% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}

% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}

% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}

% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}

% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}

% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}


% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}

% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}

% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}

% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}

% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}

\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution

\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}

\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.

\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}

\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak

mathdef.tex
% definitions and conditions
% \newtheorem{defi}{Definition}
% \newtheorem*{defi*}{Definition}
% \numberwithin{defi}{section}
% \newtheorem{fact}[defi]{Fact}
% \newtheorem{algo}[defi]{Algorithm}
% \newtheorem{cond}[defi]{Condition}
% \newtheorem{prop}[defi]{Proposition}
% \newtheorem{exam}[defi]{\textbf{Example}}

% proposition
% \newtheorem{prop}{Proposition}
% \numberwithin{prop}{section}

% lemma, theorem, corollary
\newtheorem{lemm}{Lemma}
\newtheorem{theo}[lemm]{Theorem}
\newtheorem{coro}[lemm]{Corollary}

% \theoremstyle{remark}
% \newtheorem*{remark}{\textbf{Remark}}

% \theoremstyle{definition}
\newtheorem{exam}{\textbf{Example}}
% \numberwithin{exam}{section}

%\newtheorem{claim}[defi]{Claim}
% \newcommand{\theoproof}[1]{\noindent{\bf (Proof of Theorem #1) }}
% \newcommand{\propproof}[1]{\noindent{\bf (Proof of Proposition #1) }}
% \newcommand{\lemmproof}[1]{\noindent{\bf (Proof of Lemma #1) }\\}
% \newcommand{\lemmproofn}[1]{\noindent{\bf (Proof of Lemma #1) }}
% \newcommand{\coroproof}[1]{\noindent{\bf (Proof of Corollary #1) }\\}
% \newcommand{\proof}{\noindent\textbf{(Proof)} }
\newcommand{\proofend}{\hfill$\Box$\vspace{2mm}}
% \newcommand{\QED}{\rule{7pt}{7pt}}

\newcommand{\sgn}{\mathrm{sgn}}
\newcommand{\define}{\stackrel{\mathrm{def}}{=}}
% \newcommand{\argmin}{\mathop{\mathrm{argmin\,}}}
% \newcommand{\argmax}{\mathop{\mathrm{argmax\,}}}
\newcommand{\rank}[1]{\mathrm{rank}\left(#1\right)}
\newcommand{\tr}[1]{\mathrm{tr}(#1)}
\newcommand{\diag}[1]{\mathrm{diag}\left(#1\right)}
\newcommand{\sinc}{\mathrm{sinc}\;}
\newcommand{\unorm}[1]{\|#1\|}
\newcommand{\unorms}[1]{\unorm{#1}^2}
\newcommand{\inner}[2]{\langle #1,#2\rangle}
\newcommand{\iid}{\stackrel{\mathrm{i.i.d.}}{\sim}}
\newcommand{\indep}{\mathop{\perp\!\!\!\perp}}

\newcommand{\pnu}{p}
\newcommand{\pde}{p'}
\newcommand{\Pnu}{P}
\newcommand{\Pde}{P'}
\newcommand{\pmix}{q_\alpha}
\newcommand{\numparams}{n}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%55

\newcommand{\mathbbC}{\mathbb{C}}
\newcommand{\mathbbN}{\mathbb{N}}
% \newcommand{\R}{\mathbb{R}}
\newcommand{\mathbbS}{\mathbb{S}}

\newcommand{\boldzero}{{\boldsymbol{0}}}
\newcommand{\boldone}{{\boldsymbol{1}}}

\newcommand{\boldA}{{\boldsymbol{A}}}
\newcommand{\boldB}{{\boldsymbol{B}}}
\newcommand{\boldC}{{\boldsymbol{C}}}
\newcommand{\boldD}{{\boldsymbol{D}}}
\newcommand{\boldE}{{\boldsymbol{E}}}
\newcommand{\boldF}{{\boldsymbol{F}}}
\newcommand{\boldG}{{\boldsymbol{G}}}
\newcommand{\boldH}{{\boldsymbol{H}}}
\newcommand{\boldI}{{\boldsymbol{I}}}
\newcommand{\boldJ}{{\boldsymbol{J}}}
\newcommand{\boldK}{{\boldsymbol{K}}}
\newcommand{\boldL}{{\boldsymbol{L}}}
\newcommand{\boldM}{{\boldsymbol{M}}}
\newcommand{\boldN}{{\boldsymbol{N}}}
\newcommand{\boldO}{{\boldsymbol{O}}}
\newcommand{\boldP}{{\boldsymbol{P}}}
\newcommand{\boldQ}{{\boldsymbol{Q}}}
\newcommand{\boldR}{{\boldsymbol{R}}}
\newcommand{\boldS}{{\boldsymbol{S}}}
\newcommand{\boldT}{{\boldsymbol{T}}}
\newcommand{\boldU}{{\boldsymbol{U}}}
\newcommand{\boldV}{{\boldsymbol{V}}}
\newcommand{\boldW}{{\boldsymbol{W}}}
\newcommand{\boldX}{{\boldsymbol{X}}}
\newcommand{\boldY}{{\boldsymbol{Y}}}
\newcommand{\boldZ}{{\boldsymbol{Z}}}
\newcommand{\bW}{{\boldsymbol{W}}}
\newcommand{\bX}{{\boldsymbol{X}}}
\newcommand{\bY}{{\boldsymbol{Y}}}

\newcommand{\bolda}{{\boldsymbol{a}}}
\newcommand{\boldb}{{\boldsymbol{b}}}
\newcommand{\boldc}{{\boldsymbol{c}}}
\newcommand{\boldd}{{\boldsymbol{d}}}
\newcommand{\bolde}{{\boldsymbol{e}}}
\newcommand{\boldf}{{\boldsymbol{f}}}
\newcommand{\boldg}{{\boldsymbol{g}}}
\newcommand{\boldh}{{\boldsymbol{h}}}
\newcommand{\boldi}{{\boldsymbol{i}}}
\newcommand{\boldj}{{\boldsymbol{j}}}
\newcommand{\boldk}{{\boldsymbol{k}}}
\newcommand{\boldl}{{\boldsymbol{l}}}
\newcommand{\boldm}{{\boldsymbol{m}}}
\newcommand{\boldn}{{\boldsymbol{n}}}
\newcommand{\boldo}{{\boldsymbol{o}}}
\newcommand{\boldp}{{\boldsymbol{p}}}
\newcommand{\boldq}{{\boldsymbol{q}}}
\newcommand{\boldr}{{\boldsymbol{r}}}
\newcommand{\bolds}{{\boldsymbol{s}}}
\newcommand{\boldt}{{\boldsymbol{t}}}
\newcommand{\boldu}{{\boldsymbol{u}}}
\newcommand{\boldv}{{\boldsymbol{v}}}
\newcommand{\bw}{{\boldsymbol{w}}}
\newcommand{\bx}{{\boldsymbol{x}}}
\newcommand{\by}{{\boldsymbol{y}}}
\newcommand{\bz}{{\boldsymbol{z}}}


\newcommand{\boldztr}{\boldz^{\mathrm{tr}}}
\newcommand{\boldzte}{\boldz^{\mathrm{te}}}


\newcommand{\boldalpha}{{\boldsymbol{\alpha}}}
\newcommand{\boldbeta}{{\boldsymbol{\beta}}}
\newcommand{\boldgamma}{{\boldsymbol{\gamma}}}
\newcommand{\bolddelta}{{\boldsymbol{\delta}}}
\newcommand{\boldepsilon}{{\boldsymbol{\epsilon}}}
\newcommand{\boldzeta}{{\boldsymbol{\zeta}}}
\newcommand{\boldeta}{{\boldsymbol{\eta}}}
\newcommand{\boldtheta}{{\boldsymbol{\theta}}}
\newcommand{\boldiota}{{\boldsymbol{\iota}}}
\newcommand{\boldkappa}{{\boldsymbol{\kappa}}}
\newcommand{\boldlambda}{{\boldsymbol{\lambda}}}
\newcommand{\boldmu}{{\boldsymbol{\mu}}}
\newcommand{\boldnu}{{\boldsymbol{\nu}}}
\newcommand{\boldxi}{{\boldsymbol{\xi}}}
\newcommand{\boldpi}{{\boldsymbol{\pi}}}
\newcommand{\boldrho}{{\boldsymbol{\rho}}}
\newcommand{\boldsigma}{{\boldsymbol{\sigma}}}
\newcommand{\boldtau}{{\boldsymbol{\tau}}}
\newcommand{\boldupsilon}{{\boldsymbol{\upsilon}}}
\newcommand{\boldphi}{{\boldsymbol{\phi}}}
\newcommand{\boldchi}{{\boldsymbol{\chi}}}
\newcommand{\boldpsi}{{\boldsymbol{\psi}}}
\newcommand{\boldomega}{{\boldsymbol{\omega}}}
\newcommand{\boldGamma}{{\boldsymbol{\Gamma}}}
\newcommand{\boldDelta}{{\boldsymbol{\Delta}}}
\newcommand{\boldTheta}{{\boldsymbol{\Theta}}}
\newcommand{\boldLambda}{{\boldsymbol{\Lambda}}}
\newcommand{\boldXi}{{\boldsymbol{\Xi}}}
\newcommand{\boldPi}{{\boldsymbol{\Pi}}}
\newcommand{\boldSigma}{{\boldsymbol{\Sigma}}}
\newcommand{\boldUpsilon}{{\boldsymbol{\Upsilon}}}
\newcommand{\boldPhi}{{\boldsymbol{\Phi}}}
\newcommand{\boldPsi}{{\boldsymbol{\Psi}}}
\newcommand{\boldOmega}{{\boldsymbol{\Omega}}}
\newcommand{\boldvarepsilon}{{\boldsymbol{\varepsilon}}}
\newcommand{\boldvartheta}{{\boldsymbol{\vartheta}}}
\newcommand{\boldvarpi}{{\boldsymbol{\varpi}}}
\newcommand{\boldvarrho}{{\boldsymbol{\varrho}}}
\newcommand{\boldvarsigma}{{\boldsymbol{\varsigma}}}
\newcommand{\boldvarphi}{{\boldsymbol{\varphi}}}
 
\newcommand{\calA}{{\mathcal{A}}}
\newcommand{\calB}{{\mathcal{B}}}
\newcommand{\calC}{{\mathcal{C}}}
\newcommand{\calD}{{\mathcal{D}}}
\newcommand{\calE}{{\mathcal{E}}}
\newcommand{\calF}{{\mathcal{F}}}
\newcommand{\calG}{{\mathcal{G}}}
\newcommand{\calH}{{\mathcal{H}}}
\newcommand{\calI}{{\mathcal{I}}}
\newcommand{\calJ}{{\mathcal{J}}}
\newcommand{\calK}{{\mathcal{K}}}
\newcommand{\calL}{{\mathcal{L}}}
\newcommand{\calM}{{\mathcal{M}}}
\newcommand{\calN}{{\mathcal{N}}}
\newcommand{\calO}{{\mathcal{O}}}
\newcommand{\calP}{{\mathcal{P}}}
\newcommand{\calQ}{{\mathcal{Q}}}
\newcommand{\calR}{{\mathcal{R}}}
\newcommand{\calS}{{\mathcal{S}}}
\newcommand{\calT}{{\mathcal{T}}}
\newcommand{\calU}{{\mathcal{U}}}
\newcommand{\calV}{{\mathcal{V}}}
\newcommand{\calW}{{\mathcal{W}}}
\newcommand{\calX}{{\mathcal{X}}}
\newcommand{\calY}{{\mathcal{Y}}}
\newcommand{\calZ}{{\mathcal{Z}}}

\newcommand{\boldcalZ}{{\bm \calZ}}
\newcommand{\boldcalY}{{\bm \calY}}
\newcommand{\boldcalB}{{\bm \calB}}

\newcommand{\calXtr}{{\mathcal{X}^{\mathrm{tr}}}}
\newcommand{\calYtr}{{\mathcal{Y}^{\mathrm{tr}}}}
\newcommand{\calXte}{{\mathcal{X}^{\mathrm{te}}}}
\newcommand{\calYte}{{\mathcal{Y}^{\mathrm{te}}}}

\newcommand{\thetah}{{\widehat{\theta}}}
\newcommand{\thetat}{{\widetilde{\theta}}}
\newcommand{\boldthetah}{{\widehat{\boldtheta}}}
\newcommand{\boldthetat}{{\widetilde{\boldtheta}}}
\newcommand{\boldzetah}{{\widehat{\boldzeta}}}
\newcommand{\calAh}{{\widehat{\calA}}}
\newcommand{\boldEh}{{\widehat{\boldE}}}
\newcommand{\boldGh}{{\widehat{\boldG}}}
\newcommand{\boldUh}{{\widehat{\boldU}}}
\newcommand{\Eh}{{\widehat{E}}}
\newcommand{\jh}{{\widehat{j}}}
\newcommand{\boldbetah}{{\widehat{\boldbeta}}}
\newcommand{\boldbetat}{{\widetilde{\boldbeta}}}

\newcommand{\mathrmd}{{\mathrm{d}}}
\newcommand{\longrightarrowD}{\stackrel{\mathrm{D}}{\longrightarrow}}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\newcommand{\domainx}{\calX}
\newcommand{\domainy}{\calY}
\newcommand{\domainz}{\calZ}
\newcommand{\inputdim}{d}
\newcommand{\reducedim}{m}
\newcommand{\reducedimh}{\widehat{\reducedim}}
\newcommand{\ntask}{m}
\newcommand{\nmix}{c}
\newcommand{\inputdimx}{\inputdim_{\mathrm{x}}}
\newcommand{\inputdimy}{\inputdim_{\mathrm{y}}}
\newcommand{\inputdimz}{\inputdim_{\mathrm{z}}}

\newcommand{\ratiosymbol}{r}
\newcommand{\ratio}{\ratiosymbol}
\newcommand{\ratioh}{\widehat{\ratiosymbol}}
\newcommand{\ratiomodel}{\ratiosymbol_{\alpha}}
\newcommand{\relratio}{\ratio_{\alpha}}

\newcommand{\densitysymbol}{p}
\newcommand{\cumulativesymbol}{P}
\newcommand{\density}{\densitysymbol^\ast}
\newcommand{\densitynu}{\densitysymbol_{\mathrm{te}}^\ast}
\newcommand{\densityde}{\densitysymbol_{\mathrm{de}}^\ast}
\newcommand{\densitydet}{\widetilde{\densitysymbol}_{\mathrm{de}}^\ast}
\newcommand{\cumulativede}{\cumulativesymbol_{\mathrm{de}}^\ast}
\newcommand{\cumulativenu}{\cumulativesymbol_{\mathrm{te}}^\ast}
\newcommand{\cumulativedeh}{\widehat{\cumulativesymbol}_{\mathrm{de}}}
\newcommand{\densityh}{\widehat{\densitysymbol}}
\newcommand{\densityt}{\widetilde{\densitysymbol}}
\newcommand{\densitynuh}{\widehat{\densitysymbol}_{\mathrm{te}}}
\newcommand{\densitydeh}{\widehat{\densitysymbol}_{\mathrm{de}}}
%\newcommand{\densitymodelnu}{\densitysymbol_{\mathrm{te}}}
%\newcommand{\densitymodelde}{\densitysymbol_{\mathrm{de}}}
\newcommand{\densitymodel}{\densitysymbol}
\newcommand{\prior}{\pi}
\newcommand{\posterior}{\densitysymbol}
\newcommand{\evidence}{\densitysymbol}

\newcommand{\xy}{\mathrm{xy}}
\newcommand{\zy}{\mathrm{zy}}
\newcommand{\x}{\boldsymbol{x}}
\newcommand{\y}{\boldsymbol{y}}


\newcommand{\ptr}{p_{\mathrm{tr}}}
\newcommand{\pte}{p_{\mathrm{te}}}
\newcommand{\pxy}{p_{\mathrm{x,y}}}
\newcommand{\pzy}{p_{\mathrm{z,y}}}
\newcommand{\px}{p_{\mathrm{x}}}
\newcommand{\pz}{p_{\mathrm{z}}}
\newcommand{\py}{p_{\mathrm{y}}}
\newcommand{\boldxtr}{\boldx^{\mathrm{tr}}}
\newcommand{\boldxte}{\boldx^{\mathrm{te}}}
\newcommand{\xtr}{x^{\mathrm{tr}}}
\newcommand{\xte}{x^{\mathrm{te}}}
\newcommand{\ytr}{y^{\mathrm{tr}}}
\newcommand{\yte}{y^{\mathrm{te}}}
\newcommand{\epsilontr}{\epsilon^{\mathrm{tr}}}
\newcommand{\epsilonte}{\epsilon^{\mathrm{te}}}
\newcommand{\ntr}{n}
\newcommand{\nte}{n_{\mathrm{te}}}
\newcommand{\nteu}{n_{\mathrm{te}}'}
\newcommand{\ntep}{n_{\mathrm{te}}}
\newcommand{\densitytr}{\densitysymbol_{\mathrm{tr}}^\ast}
\newcommand{\densityte}{\densitysymbol_{\mathrm{te}}^\ast}
\newcommand{\Xtr}{X^{\mathrm{tr}}}
\newcommand{\Xte}{X^{\mathrm{te}}}
\newcommand{\boldXtr}{\boldX^{\mathrm{tr}}}
\newcommand{\boldXte}{\boldX^{\mathrm{te}}}
\newcommand{\boldytr}{\boldy^{\mathrm{tr}}}
\newcommand{\boldyte}{\boldy^{\mathrm{te}}}
\newcommand{\Gen}{G}
\newcommand{\Genh}{\widehat{\Gen}}

\newcommand{\boldxnu}{\boldx^{\mathrm{te}}}
\newcommand{\boldxde}{\boldx^{\mathrm{tr}}}
\newcommand{\boldxt}{\widetilde{\boldx}}

\newcommand{\boldunu}{\boldu^{\mathrm{te}}}
\newcommand{\boldude}{\boldu^{\mathrm{tr}}}
\newcommand{\boldvnu}{\boldv^{\mathrm{te}}}
\newcommand{\boldvde}{\boldv^{\mathrm{tr}}}
\newcommand{\boldunuh}{\widehat{\boldu}^{\mathrm{te}}}
\newcommand{\boldudeh}{\widehat{\boldu}^{\mathrm{tr}}}
\newcommand{\unu}{u^{\mathrm{te}}}
\newcommand{\ude}{u^{\mathrm{tr}}}

\newcommand{\nde}{n_{\mathrm{tr}}}
\newcommand{\nnu}{n_{\mathrm{te}}}
\newcommand{\nsample}{n}

\newcommand{\nparam}{b}
\newcommand{\nparamm}{t}
\newcommand{\nclass}{c}
\newcommand{\ndet}{\widetilde{n}_{\mathrm{tr}}}

\newcommand{\data}{D}
\newcommand{\datanu}{D^{\mathrm{te}}}
\newcommand{\datade}{D^{\mathrm{tr}}}

\newcommand{\boldratiosymbol}{\boldr}
\newcommand{\boldratio}{\boldratiosymbol^\ast}
\newcommand{\boldratiomodel}{{\boldratiosymbol}}
\newcommand{\boldratiode}{\boldratiosymbol_{\mathrm{tr}}^\ast}
\newcommand{\boldrationu}{\boldratiosymbol_{\mathrm{te}}^\ast}
\newcommand{\boldratioh}{\widehat{\boldratiosymbol}}
\newcommand{\boldratiodeh}{\boldratioh_{\mathrm{tr}}}
\newcommand{\boldrationuh}{\boldratioh_{\mathrm{te}}}
\newcommand{\boldratiodemodel}{\boldratiomodel_{\mathrm{tr}}}
\newcommand{\boldrationumodel}{\boldratiomodel_{\mathrm{te}}}
\newcommand{\boldPhide}{\boldPhi_{\mathrm{tr}}}
\newcommand{\boldPhinu}{\boldPhi_{\mathrm{te}}}
\newcommand{\boldPside}{\boldPsi_{\mathrm{tr}}}
\newcommand{\boldPsinu}{\boldPsi_{\mathrm{te}}}

\newcommand{\Bregman}{\mathrm{BR}}
\newcommand{\fdivergence}{\mathrm{ASC}}

% \newcommand{\fh}{\widehat{f}}
\newcommand{\function}{f^\ast}
\newcommand{\functionmodel}{f}
\newcommand{\functionh}{\widehat{f}}

\newcommand{\xde}{x^{\mathrm{tr}}}
\newcommand{\xdet}{\widetilde{x}^{\mathrm{tr}}}
\newcommand{\xnu}{x^{\mathrm{te}}}
\newcommand{\zde}{z^{\mathrm{tr}}}
\newcommand{\boldxdet}{\widetilde{\boldx}^{\mathrm{tr}}}
\newcommand{\boldxnut}{\widetilde{\boldx}^{\mathrm{te}}}
\newcommand{\etade}{\eta^{\mathrm{tr}}}
\newcommand{\etanu}{\eta^{\mathrm{te}}}
\newcommand{\boldepsilonde}{\boldepsilon^{\mathrm{tr}}}
\newcommand{\boldepsilondet}{\widetilde{\boldepsilon}^{\mathrm{tr}}}
\newcommand{\epsilonde}{\epsilon^{\mathrm{tr}}}
\newcommand{\epsilonnu}{\epsilon^{\mathrm{te}}}
\newcommand{\epsilondeh}{\widehat{\epsilon}^{\mathrm{tr}}}
\newcommand{\epsilondet}{\widetilde{\epsilon}^{\mathrm{tr}}}
\newcommand{\boldrde}{\boldr^{\mathrm{tr}}}
\newcommand{\boldrnu}{\boldr^{\mathrm{te}}}
\newcommand{\Ade}{A^{\mathrm{tr}}}
\newcommand{\Anu}{A^{\mathrm{te}}}
\newcommand{\Xde}{X^{\mathrm{tr}}}
\newcommand{\Xdet}{\widetilde{X}^{\mathrm{tr}}}
\newcommand{\Xnu}{X^{\mathrm{te}}}
\newcommand{\boldXde}{\boldX^{\mathrm{tr}}}
\newcommand{\boldXdet}{\widetilde{\boldX}^{\mathrm{tr}}}
\newcommand{\boldXnu}{\boldX^{\mathrm{te}}}
\newcommand{\boldGde}{\boldG^{\mathrm{tr}}}
\newcommand{\boldGnu}{\boldG^{\mathrm{te}}}
\newcommand{\boldAde}{\boldA^{\mathrm{tr}}}
\newcommand{\boldAnu}{\boldA^{\mathrm{te}}}
\newcommand{\boldKdede}{\boldK_{\mathrm{tr},\mathrm{tr}}}
\newcommand{\boldKdenu}{\boldK_{\mathrm{tr},\mathrm{te}}}
\newcommand{\boldKde}{\boldK_{\mathrm{tr}}}
\newcommand{\boldKnu}{\boldK_{\mathrm{te}}}
\newcommand{\boldyde}{\boldy^{\mathrm{tr}}}
\newcommand{\boldydet}{\widetilde{\boldy}^{\mathrm{tr}}}
\newcommand{\Kdede}{K_{\mathrm{tr},\mathrm{tr}}}
\newcommand{\Kdenu}{K_{\mathrm{tr},\mathrm{te}}}
\newcommand{\Kde}{K_{\mathrm{tr}}}
\newcommand{\yde}{y^{\mathrm{tr}}}
\newcommand{\ydet}{\widetilde{y}^{\mathrm{tr}}}
\newcommand{\ynu}{y^{\mathrm{te}}}
\newcommand{\yh}{\widehat{y}}
\newcommand{\ynuh}{\widehat{y}^{\mathrm{te}}}
\newcommand{\calXde}{\calX^{\mathrm{tr}}}
\newcommand{\calYde}{\calY^{\mathrm{tr}}}
\newcommand{\calZde}{\calZ^{\mathrm{tr}}}
\newcommand{\calXnu}{\calX^{\mathrm{te}}}
\newcommand{\calYnu}{\calY^{\mathrm{te}}}
\newcommand{\calZnu}{\calZ^{\mathrm{te}}}
\newcommand{\calZt}{\widetilde{\calZ}}

\newcommand{\Pb}{{P^{\ast}}'}
\newcommand{\Pa}{P^\ast}
\newcommand{\pb}{{p^{\ast}}'}
\newcommand{\pa}{p^\ast}
\newcommand{\nb}{n'}
\newcommand{\na}{n}
\newcommand{\calXb}{\calX'}
\newcommand{\calXa}{\calX}
\newcommand{\calXbt}{\widetilde{\calX}'}
\newcommand{\calXat}{\widetilde{\calX}}
\newcommand{\boldxb}{\boldx'}
\newcommand{\boldxa}{\boldx}
\newcommand{\boldHhde}{{\widehat{\boldH}}}
\newcommand{\boldhhde}{{\widehat{\boldh}}'}
\newcommand{\boldhhnu}{{\widehat{\boldh}}}
\newcommand{\boldGhde}{{\widehat{\boldG}}}
\newcommand{\boldghde}{{\widehat{\boldg}}'}
\newcommand{\boldghnu}{{\widehat{\boldg}}}
\newcommand{\nn}{\overline{n}}

\newcommand{\boldTh}{{\widehat{\boldT}}}
\newcommand{\boldgh}{{\widehat{\boldg}}}
\newcommand{\boldHh}{{\widehat{\boldH}}}
\newcommand{\boldhh}{{\widehat{\boldh}}}
\newcommand{\boldBh}{{\widehat{\boldB}}}
\newcommand{\Hh}{{\widehat{H}}}
\newcommand{\hh}{{\widehat{h}}}
\newcommand{\boldHt}{{\widetilde{\boldH}}}
\newcommand{\boldht}{{\widetilde{\boldh}}}
\newcommand{\boldDt}{{\widetilde{\boldD}}}
\newcommand{\boldWt}{{\widetilde{\boldW}}}
\newcommand{\boldXt}{{\widetilde{\boldX}}}
\newcommand{\Wt}{{\widetilde{W}}}
\newcommand{\Dt}{{\widetilde{D}}}
\newcommand{\Qt}{{\widetilde{Q}}}
\newcommand{\Qh}{{\widehat{Q}}}
\newcommand{\Ht}{{\widetilde{H}}}
\newcommand{\Xt}{{\widetilde{X}}}
\newcommand{\Jh}{\widehat{J}}
\newcommand{\Jt}{\widetilde{J}}
\newcommand{\Vh}{\widehat{V}}
\newcommand{\Fh}{\widehat{F}}
\newcommand{\Gh}{\widehat{G}}
% \newcommand{\KL}{\mathrm{KL}}
\newcommand{\KLh}{\widehat{\KL}}
\newcommand{\KLt}{\widetilde{\KL}}
\newcommand{\boldxih}{\widehat{\boldxi}}
\newcommand{\boldalphah}{\widehat{\boldalpha}}
\newcommand{\boldalphat}{\widetilde{\boldalpha}}
\newcommand{\alphah}{\widehat{\alpha}}
\newcommand{\sigmah}{{\widehat{\sigma}}}
\newcommand{\lambdah}{{\widehat{\lambda}}}

\newcommand{\boldSb}{\boldS^{\mathrm{b}}}
\newcommand{\boldSw}{\boldS^{\mathrm{w}}}
\newcommand{\boldWb}{\boldW^{\mathrm{b}}}
\newcommand{\boldWw}{\boldW^{\mathrm{w}}}
\newcommand{\Wb}{W^{\mathrm{b}}}
\newcommand{\Ww}{W^{\mathrm{w}}}
\newcommand{\boldSlb}{\boldS^{\mathrm{lb}}}
\newcommand{\boldSlw}{\boldS^{\mathrm{lw}}}
\newcommand{\boldWlb}{\boldW^{\mathrm{lb}}}
\newcommand{\boldWlw}{\boldW^{\mathrm{lw}}}
\newcommand{\Wlb}{W^{\mathrm{lb}}}
\newcommand{\Wlw}{W^{\mathrm{lw}}}
\newcommand{\weight}{w}
\newcommand{\weighth}{\widehat{w}}


%\newcommand{\resample}{\omega}
\newcommand{\resample}{\zeta}
\newcommand{\CEVOW}{\textrm{CV}_\textrm{W}}
\newcommand{\PCEVOW}{\textrm{PCV}_\textrm{W}}
\newcommand{\CEVOO}{\textrm{CV}_\textrm{O}}
\newcommand{\PCEVOO}{\textrm{PCV}_\textrm{O}}
\newcommand{\FEVOW}{\textrm{FV}_\textrm{W}}
\newcommand{\PFEVOW}{\textrm{PFV}_\textrm{W}}
\newcommand{\FEBVOW}{\textrm{FBV}_\textrm{OW}}
\newcommand{\PFEBVOW}{\textrm{PFBV}_\textrm{OW}}
\newcommand{\ALICE}{\textrm{ALICE}}
\newcommand{\PALICE}{\textrm{PALICE}}

\newcommand{\model}{M}
\newcommand{\GhatAL}{\Genh^{(AL)}}
\newcommand{\GhatMS}{\Genh^{(MS)}}
\newcommand{\GhatuCEAL}{\Genh^{(EAL)}}
\newcommand{\GhatCEAL}{\Genh^{(nEAL)}}

\newcommand{\hatgn}{\widehat{g}_n}
\newcommand{\gnstar}{g_n^*}
\newcommand{\alzn}{\alpha_0^n}
\newcommand{\supp}{\mathrm{supp}}
\newcommand{\gzn}{\ratio}
\newcommand{\Gconv}{\mathcal{G}^{(\mathrm{c})}}

% custom notation
\newcommand*{\A}{\mathcal{A}}
\newcommand*{\B}{\mathcal{B}}
\newcommand*{\F}{\mathcal{F}}
\newcommand*{\X}{\mathcal{X}}
\newcommand*{\V}{\mathbb{V}}
\newcommand*{\N}{\mathcal{N}}
\newcommand*{\TV}{\mathrm{TV}}
\newcommand*{\LL}{\left}
\newcommand*{\RR}{\right}
\newcommand*{\tPhi}{\tilde{\Phi}}
\newcommand\widebar[1]{\mathop{\overline{#1}}}

% \newcommand*{\Cov}{\mathbb{Cov}}
\newcommand*{\Exp}[1]{\mathbb{E} \left[#1\right]}
\newcommand*{\Covar}{\mathbb{Cov}\left( \right)}
% \newcommand*{\Var}[1]{\mathbb{V}\left(#1\right)}

\newcommand*{\bracks}[1]{\left(#1\right)}  % usual brackets
\newcommand*{\abracks}[1]{\left\langle#1\right\rangle}  % angle brackets
\newcommand*{\sbracks}[1]{\left[#1\right]}  % square brackets

% differentials
\newcommand{\dee}{\mathop{\mathrm{d}\!}}
\newcommand{\dt}{\,\dee t}
\newcommand{\de}{\,\dee e}
\newcommand{\ds}{\,\dee s}
\newcommand{\dx}{\,\dee x}
\newcommand{\dX}{\,\dee X}
\newcommand{\dy}{\,\dee y}
\newcommand{\dY}{\,\dee Y}
\newcommand{\dz}{\,\dee z}
\newcommand{\dZ}{\,\dee Z}
\newcommand{\dv}{\,\dee v}
\newcommand{\du}{\,\dee u}
\newcommand{\dw}{\,\dee w}
\newcommand{\dr}{\,\dee r}
\newcommand{\dB}{\,\dee B} % Brownian motion
\newcommand{\dW}{\,\dee W} % Wiener process
\newcommand{\dtau}{\,\dee \tau}
\newcommand{\dmu}{\,\dee \mu}
\newcommand{\dnu}{\,\dee \nu}
\newcommand{\dzeta}{\,\dee \zeta}
\newcommand{\domega}{\,\dee \omega}

\newcommand{\dtx}{\,\dee \tilde{x}}
\newcommand{\dtv}{\,\dee \tilde{v}}
\newcommand{\tx}{\tilde{x}}
\newcommand{\tv}{\tilde{v}}

\newcommand{\dtX}{\,\dee \tilde{X}}
\newcommand{\dtV}{\,\dee \tilde{V}}
% \newcommand{\tX}{\tilde{X}}
% \newcommand{\tH}{\tilde{H}}
% \newcommand{\tV}{\tilde{V}}
% \newcommand{\tY}{\tilde{Y}}
% \newcommand{\tZ}{\tilde{Z}}

\newcommand{\dbarx}{\,\dee \bar{x}}
\newcommand{\dbarv}{\,\dee \bar{v}}
\newcommand{\barx}{\bar{x}}
\newcommand{\barX}{\bar{X}}
\newcommand{\barv}{\bar{v}}
\newcommand{\barR}{\Bar{R}}

\newcommand{\dhx}{\,\dee \hat{x}}
\newcommand{\dhv}{\,\dee \hat{v}}
\newcommand{\hx}{\hat{x}}
\newcommand{\hv}{\hat{v}}

\newcommand{\lyap}{\mathcal{U}}
\newcommand{\uyap}{\mathcal{V}}
\newcommand{\wyap}{\mathcal{W}}

% norms, trace and others
\newcommand{\norm}[1]{\left\lVert#1\right\rVert}
\newcommand{\normf}[1]{\left\lVert#1\right\rVert_{\mathrm{F}}}
\newcommand{\normop}[1]{\left\lVert#1\right\rVert_{\mathrm{op}}}
% \newcommand{\normltwo}[1]{\left\lVert#1\right\rVert_{\mathbb{L}^2}}
\newcommand{\normtwo}[1]{\left\lVert#1\right\rVert_2}
\newcommand{\eq}[1]{\begin{align}#1\end{align}}
\newcommand{\eqn}[1]{\begin{align*}#1\end{align*}}

% \newcommand{\Tr}[1]{\mathrm{Tr}\left(#1\right)}
\newcommand{\op}{\mathrm{op}}
\usepackage[super]{nth}
\newcommand{\lap}[1]{\Vec{\Delta} \left(#1\right)}

\usepackage{nicefrac}

\def\ssum{\mathsmaller{\sum}}
\def\sint{\mathsmaller{\int}}

%%% Local Variables: 
%%% mode: latex
%%% TeX-master: "main"
%%% End:
main.bbl
\begin{thebibliography}{}

\bibitem[Archambeau et~al., 2008]{archambeau2008variational}
Archambeau, C., Opper, M., Shen, Y., Cornford, D., and Shawe-Taylor, J. (2008).
\newblock Variational inference for diffusion processes.
\newblock {\em Advances in Neural Information Processing Systems}.

\bibitem[Balan et~al., 2015]{balan2015bayesian}
Balan, A.~K., Rathod, V., Murphy, K.~P., and Welling, M. (2015).
\newblock Bayesian dark knowledge.
\newblock In {\em Advances in Neural Information Processing Systems}, pages
  3438--3446.

\bibitem[Blundell et~al., 2015]{blundell2015weight}
Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D. (2015).
\newblock Weight uncertainty in neural networks.
\newblock {\em International Conference on Machine Learning}.

\bibitem[Bou{\'e} et~al., 1998]{boue1998variational}
Bou{\'e}, M., Dupuis, P., et~al. (1998).
\newblock A variational representation for certain functionals of brownian
  motion.
\newblock {\em The Annals of Probability}, 26(4):1641--1659.

\bibitem[Chang et~al., 2020]{Chang2020PrincipledWI}
Chang, O., Flokas, L., and Lipson, H. (2020).
\newblock Principled weight initialization for hypernetworks.
\newblock In {\em International Conference on Learning Representations}.

\bibitem[Chen et~al., 2018]{chen2018neural}
Chen, R. T.~Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. (2018).
\newblock Neural ordinary differential equations.
\newblock {\em Advances in Neural Information Processing Systems}.

\bibitem[Dandekar et~al., 2020]{dandekar2020bayesian}
Dandekar, R., Dixit, V., Tarek, M., Garcia-Valadez, A., and Rackauckas, C.
  (2020).
\newblock Bayesian neural ordinary differential equations.
\newblock {\em arXiv preprint arXiv:2012.07244}.

\bibitem[De~Brouwer et~al., 2019]{de2019gru}
De~Brouwer, E., Simm, J., Arany, A., and Moreau, Y. (2019).
\newblock Gru-ode-bayes: Continuous modeling of sporadically-observed time
  series.
\newblock {\em arXiv preprint arXiv:1905.12374}.

\bibitem[Deng, 2012]{deng2012mnist}
Deng, L. (2012).
\newblock The mnist database of handwritten digit images for machine learning
  research.
\newblock {\em IEEE Signal Processing Magazine}, 29(6):141--142.

\bibitem[Dinh et~al., 2016]{dinh2016density}
Dinh, L., Sohl-Dickstein, J., and Bengio, S. (2016).
\newblock Density estimation using real nvp.
\newblock {\em arXiv preprint arXiv:1605.08803}.

\bibitem[Dupont et~al., 2019]{dupont2019augmented}
Dupont, E., Doucet, A., and Teh, Y.~W. (2019).
\newblock Augmented neural odes.
\newblock {\em Advances in Neural Information Processing Systems}.

\bibitem[Dusenberry et~al., 2020]{dusenberry2020efficient}
Dusenberry, M., Jerfel, G., Wen, Y., Ma, Y., Snoek, J., Heller, K.,
  Lakshminarayanan, B., and Tran, D. (2020).
\newblock Efficient and scalable bayesian neural nets with rank-1 factors.
\newblock In {\em International conference on machine learning}, pages
  2782--2792. PMLR.

\bibitem[Duvenaud et~al., 2014]{DuvRipAdaGha2014}
Duvenaud, D., Rippel, O., Adams, R.~P., and Ghahramani, Z. (2014).
\newblock Avoiding pathologies in very deep networks.
\newblock In {\em Artificial Intelligence and Statistics}.

\bibitem[E, 2017]{E2017}
E, W. (2017).
\newblock {A Proposal on Machine Learning via Dynamical Systems}.
\newblock {\em Commun. Math. Stat.}, 5(1):1--11.

\bibitem[Graves, 2011]{graves2011practical}
Graves, A. (2011).
\newblock Practical variational inference for neural networks.
\newblock In {\em Advances in neural information processing systems}, pages
  2348--2356.

\bibitem[Guo et~al., 2017]{guo2017calibration}
Guo, C., Pleiss, G., Sun, Y., and Weinberger, K.~Q. (2017).
\newblock On calibration of modern neural networks.
\newblock In {\em Proceedings of the 34th International Conference on Machine
  Learning-Volume 70}, pages 1321--1330. JMLR. org.

\bibitem[Ha et~al., 2016]{ha2016hypernetworks}
Ha, D., Dai, A., and Le, Q.~V. (2016).
\newblock Hypernetworks.
\newblock {\em arXiv preprint arXiv:1609.09106}.

\bibitem[Haber and Ruthotto, 2017]{haber2017stable}
Haber, E. and Ruthotto, L. (2017).
\newblock Stable architectures for deep neural networks.
\newblock {\em Inverse Problems}, 34(1):014004.

\bibitem[Hegde et~al., 2018]{hegde2018deep}
Hegde, P., Heinonen, M., L{\"a}hdesm{\"a}ki, H., and Kaski, S. (2018).
\newblock Deep learning with differential gaussian process flows.
\newblock {\em arXiv preprint arXiv:1810.04066}.

\bibitem[Hendrycks and Dietterich, 2019]{hendrycks2019robustness}
Hendrycks, D. and Dietterich, T. (2019).
\newblock Benchmarking neural network robustness to common corruptions and
  perturbations.
\newblock {\em Proceedings of the International Conference on Learning
  Representations}.

\bibitem[Hoffman et~al., 2013]{hoffman2013stochastic}
Hoffman, M.~D., Blei, D.~M., Wang, C., and Paisley, J. (2013).
\newblock Stochastic variational inference.
\newblock {\em The Journal of Machine Learning Research}, 14(1):1303--1347.

\bibitem[Innes et~al., 2019]{innes2019zygote}
Innes, M., Edelman, A., Fischer, K., Rackauckus, C., Saba, E., Shah, V.~B., and
  Tebbutt, W. (2019).
\newblock Zygote: A differentiable programming system to bridge machine
  learning and scientific computing.
\newblock {\em arXiv preprint arXiv:1907.07587}, page 140.

\bibitem[Ioffe and Szegedy, 2015]{ioffe2015batch}
Ioffe, S. and Szegedy, C. (2015).
\newblock Batch normalization: Accelerating deep network training by reducing
  internal covariate shift.
\newblock In {\em International conference on machine learning}, pages
  448--456. PMLR.

\bibitem[Izmailov et~al., 2021]{izmailov2021bayesian}
Izmailov, P., Vikram, S., Hoffman, M.~D., and Wilson, A.~G. (2021).
\newblock What are bayesian neural network posteriors really like?
\newblock {\em International Conference on Learning Representations}.

\bibitem[Kelly et~al., 2020]{kelly2020easynode}
Kelly, J., Bettencourt, J., Johnson, M.~J., and Duvenaud, D. (2020).
\newblock Learning differential equations that are easy to solve.
\newblock In {\em Neural Information Processing Systems}.

\bibitem[Kidger et~al., 2021]{kidger2021neural}
Kidger, P., Foster, J., Li, X., Oberhauser, H., and Lyons, T. (2021).
\newblock Neural sdes as infinite-dimensional gans.
\newblock {\em arXiv preprint arXiv:2102.03657}.

\bibitem[Kidger et~al., 2020]{kidger2020neural}
Kidger, P., Morrill, J., Foster, J., and Lyons, T. (2020).
\newblock Neural controlled differential equations for irregular time series.
\newblock {\em arXiv preprint arXiv:2005.08926}.

\bibitem[Kingma and Ba, 2015]{Kingma2015AdamAM}
Kingma, D.~P. and Ba, J. (2015).
\newblock Adam: A method for stochastic optimization.
\newblock {\em CoRR}, abs/1412.6980.

\bibitem[Kong et~al., 2020]{kong2020sde}
Kong, L., Sun, J., and Zhang, C. (2020).
\newblock Sde-net: Equipping deep neural networks with uncertainty estimates.
\newblock {\em arXiv preprint arXiv:2008.10546}.

\bibitem[Krizhevsky et~al., 2014]{krizhevsky2014cifar}
Krizhevsky, A., Nair, V., and Hinton, G. (2014).
\newblock The cifar-10 dataset.
\newblock {\em online: http://www. cs. toronto. edu/kriz/cifar. html}, 55:5.

\bibitem[Krueger et~al., 2018]{krueger2018bayesian}
Krueger, D., Huang, C.-W., Islam, R., Turner, R., Lacoste, A., and Courville,
  A. (2018).
\newblock Bayesian hypernetworks.
\newblock {\em arXiv preprint arXiv:1710.04759}.

\bibitem[Li et~al., 2020]{li2020scalable}
Li, X., Wong, T.-K.~L., Chen, R.~T., and Duvenaud, D. (2020).
\newblock Scalable gradients for stochastic differential equations.
\newblock {\em arXiv preprint arXiv:2001.01328}.

\bibitem[Liu et~al., 2019]{liu2019neural}
Liu, X., Xiao, T., Si, S., Cao, Q., Kumar, S., and Hsieh, C.-J. (2019).
\newblock Neural sde: Stabilizing neural ode networks with stochastic noise.
\newblock {\em arXiv preprint arXiv:1906.02355}.

\bibitem[Look and Kandemir, 2019]{look2019differential}
Look, A. and Kandemir, M. (2019).
\newblock Differential bayesian neural nets.
\newblock {\em arXiv preprint arXiv:1912.00796}.

\bibitem[Louizos and Welling, 2017]{louizos2017multiplicative}
Louizos, C. and Welling, M. (2017).
\newblock Multiplicative normalizing flows for variational bayesian neural
  networks.
\newblock In {\em Proceedings of the 34th International Conference on Machine
  Learning-Volume 70}, pages 2218--2227. JMLR. org.

\bibitem[Mishkin et~al., 2018]{mishkin2018slang}
Mishkin, A., Kunstner, F., Nielsen, D., Schmidt, M., and Khan, M.~E. (2018).
\newblock Slang: Fast structured covariance approximations for bayesian deep
  learning with natural gradient.
\newblock In {\em Advances in Neural Information Processing Systems}, pages
  6245--6255.

\bibitem[Neal, 1996]{neal2012bayesian}
Neal, R.~M. (1996).
\newblock {\em Bayesian learning for neural networks}, volume 118.
\newblock Springer Science \& Business Media.

\bibitem[Oganesyan et~al., 2020]{oganesyan2020stochasticity}
Oganesyan, V., Volokhova, A., and Vetrov, D. (2020).
\newblock Stochasticity in neural odes: An empirical study.
\newblock {\em arXiv preprint arXiv:2002.09779}.

\bibitem[Opper, 2019]{opper2019variational}
Opper, M. (2019).
\newblock Variational inference for stochastic differential equations.
\newblock {\em Annalen der Physik}, 531(3):1800233.

\bibitem[Osawa et~al., 2019]{osawa2019practical}
Osawa, K., Swaroop, S., Jain, A., Eschenhagen, R., Turner, R.~E., Yokota, R.,
  and Khan, M.~E. (2019).
\newblock Practical deep learning with bayesian principles.
\newblock {\em arXiv preprint arXiv:1906.02506}.

\bibitem[Peluchetti and Favaro, 2020a]{peluchetti2020doubly}
Peluchetti, S. and Favaro, S. (2020a).
\newblock Doubly infinite residual networks: a diffusion process approach.

\bibitem[Peluchetti and Favaro, 2020b]{peluchetti2020infinitely}
Peluchetti, S. and Favaro, S. (2020b).
\newblock Infinitely deep neural networks as diffusion processes.
\newblock In {\em International Conference on Artificial Intelligence and
  Statistics}, pages 1126--1136. PMLR.

\bibitem[Rezende et~al., 2014]{rezende2014stochastic}
Rezende, D.~J., Mohamed, S., and Wierstra, D. (2014).
\newblock Stochastic backpropagation and approximate inference in deep
  generative models.
\newblock {\em arXiv preprint arXiv:1401.4082}.

\bibitem[Roeder et~al., 2017]{roeder2017sticking}
Roeder, G., Wu, Y., and Duvenaud, D.~K. (2017).
\newblock Sticking the landing: Simple, lower-variance gradient estimators for
  variational inference.
\newblock In {\em Advances in Neural Information Processing Systems}.

\bibitem[Rubanova et~al., 2019]{rubanova2019latent}
Rubanova, Y., Chen, R. T.~Q., and Duvenaud, D. (2019).
\newblock Latent odes for irregularly-sampled time series.
\newblock {\em arXiv preprint arXiv:1907.03907}.

\bibitem[Ryder et~al., 2018]{ryder2018blackbox}
Ryder, T., Golightly, A., McGough, A.~S., and Prangle, D. (2018).
\newblock Black-box variational inference for stochastic differential
  equations.
\newblock In {\em International Conference on Machine Learning}.

\bibitem[Tzen and Raginsky, 2019a]{tzen2019neural}
Tzen, B. and Raginsky, M. (2019a).
\newblock Neural stochastic differential equations: Deep latent gaussian models
  in the diffusion limit.
\newblock {\em arXiv preprint arXiv:1905.09883}.

\bibitem[Tzen and Raginsky, 2019b]{tzen2019theoretical}
Tzen, B. and Raginsky, M. (2019b).
\newblock Theoretical guarantees for sampling and inference in generative
  models with latent diffusions.
\newblock {\em arXiv preprint arXiv:1903.01608}.

\bibitem[Wang et~al., 2018]{wang2018adversarial}
Wang, K.-C., Vicol, P., Lucas, J., Gu, L., Grosse, R., and Zemel, R. (2018).
\newblock Adversarial distillation of bayesian neural network posteriors.
\newblock {\em arXiv preprint arXiv:1806.10317}.

\bibitem[Wenzel et~al., 2020]{wenzel2020good}
Wenzel, F., Roth, K., Veeling, B., Swiatkowski, J., Tran, L., Mandt, S., Snoek,
  J., Salimans, T., Jenatton, R., and Nowozin, S. (2020).
\newblock How good is the bayes posterior in deep neural networks really?
\newblock In {\em International Conference on Machine Learning}.

\bibitem[Y{\i}ld{\i}z et~al., 2019]{yildiz2019ode}
Y{\i}ld{\i}z, {\c{C}}., Heinonen, M., and L{\"a}hdesm{\"a}ki, H. (2019).
\newblock Ode$^2$vae: Deep generative second order odes with bayesian neural
  networks.
\newblock {\em arXiv preprint arXiv:1905.10994}.

\bibitem[Zadrozny and Elkan, 2002]{zadrozny2002transforming}
Zadrozny, B. and Elkan, C. (2002).
\newblock Transforming classifier scores into accurate multiclass probability
  estimates.
\newblock In {\em Proceedings of the Eighth ACM SIGKDD International Conference
  on Knowledge Discovery and Data Mining}.

\bibitem[Zhang et~al., 2018]{zhang2018noisy}
Zhang, G., Sun, S., Duvenaud, D., and Grosse, R. (2018).
\newblock Noisy natural gradient as variational inference.
\newblock In {\em International Conference on Machine Learning}, pages
  5852--5861.

\bibitem[Zhang et~al., 2019a]{zhang2019fixup}
Zhang, H., Dauphin, Y.~N., and Ma, T. (2019a).
\newblock Fixup initialization: Residual learning without normalization.
\newblock {\em arXiv preprint arXiv:1901.09321}.

\bibitem[Zhang et~al., 2019b]{zhang2019approximation}
Zhang, H., Gao, X., Unterman, J., and Arodz, T. (2019b).
\newblock Approximation capabilities of neural ordinary differential equations.
\newblock {\em arXiv preprint arXiv:1907.12998}.

\bibitem[Zhang et~al., 2019c]{anodev2}
Zhang, T., Yao, Z., Gholami, A., Keutzer, K., Gonzalez, J., Biros, G., and
  Mahoney, M.~W. (2019c).
\newblock {ANODEV2:} {A} coupled neural {ODE} evolution framework.
\newblock {\em CoRR}, abs/1906.04596.

\end{thebibliography}

references.bib
@inproceedings{kelly2020easynode,
  title={Learning Differential Equations that are Easy to Solve},
  author={Kelly, Jacob and Bettencourt, Jesse and Johnson, Matthew James and Duvenaud, David},
  booktitle={Neural Information Processing Systems},
  year={2020},
  url={https://arxiv.org/abs/2007.04504}
}


@inproceedings{wenzel2020good,
  title={How Good is the Bayes Posterior in Deep Neural Networks Really?},
  author={Wenzel, Florian and Roth, Kevin and Veeling, Bastiaan and Swiatkowski, Jakub and Tran, Linh and Mandt, Stephan and Snoek, Jasper and Salimans, Tim and Jenatton, Rodolphe and Nowozin, Sebastian},
  booktitle={International Conference on Machine Learning},
  year={2020},
}

@article{deng2012mnist,
  title={The mnist database of handwritten digit images for machine learning research},
  author={Deng, Li},
  journal={IEEE Signal Processing Magazine},
  volume={29},
  number={6},
  pages={141--142},
  year={2012},
  publisher={IEEE}
}

@article{kidger2020neural,
  title={Neural controlled differential equations for irregular time series},
  author={Kidger, Patrick and Morrill, James and Foster, James and Lyons, Terry},
  journal={arXiv preprint arXiv:2005.08926},
  year={2020}
}

@article{kidger2021neural,
  title={Neural sdes as infinite-dimensional gans},
  author={Kidger, Patrick and Foster, James and Li, Xuechen and Oberhauser, Harald and Lyons, Terry},
  journal={arXiv preprint arXiv:2102.03657},
  year={2021}
}

@article{opper2019variational,
  title={Variational inference for stochastic differential equations},
  author={Opper, Manfred},
  journal={Annalen der Physik},
  volume={531},
  number={3},
  pages={1800233},
  year={2019},
  publisher={Wiley Online Library}
}

@article{archambeau2008variational,
  title={Variational inference for diffusion processes},
  author={Archambeau, C{\'e}dric and Opper, Manfred and Shen, Yuan and Cornford, Dan and Shawe-Taylor, John},
  journal={Advances in Neural Information Processing Systems},
  year={2008},
  publisher={MIT}
}

@article{krizhevsky2014cifar,
  title={The cifar-10 dataset},
  author={Krizhevsky, Alex and Nair, Vinod and Hinton, Geoffrey},
  journal={online: http://www. cs. toronto. edu/kriz/cifar. html},
  volume={55},
  pages={5},
  year={2014}
}

@article{oganesyan2020stochasticity,
  title={Stochasticity in Neural ODEs: An Empirical Study},
  author={Oganesyan, Viktor and Volokhova, Alexandra and Vetrov, Dmitry},
  journal={arXiv preprint arXiv:2002.09779},
  year={2020}
}

@article{hegde2018deep,
  title={Deep learning with differential Gaussian process flows},
  author={Hegde, Pashupati and Heinonen, Markus and L{\"a}hdesm{\"a}ki, Harri and Kaski, Samuel},
  journal={arXiv preprint arXiv:1810.04066},
  year={2018}
}

@article{tzen2019theoretical,
  title={Theoretical guarantees for sampling and inference in generative models with latent diffusions},
  author={Tzen, Belinda and Raginsky, Maxim},
  journal={arXiv preprint arXiv:1903.01608},
  year={2019}
}

@article{tzen2019neural,
  title={Neural stochastic differential equations: Deep latent gaussian models in the diffusion limit},
  author={Tzen, Belinda and Raginsky, Maxim},
  journal={arXiv preprint arXiv:1905.09883},
  year={2019}
}

@inproceedings{ryder2018blackbox,
  title={Black-box variational inference for stochastic differential equations},
  author={Ryder, Tom and Golightly, Andrew and McGough, A Stephen and Prangle, Dennis},
  booktitle={International Conference on Machine Learning},
  year={2018},
}

@inproceedings{Chang2020PrincipledWI,
  title={Principled Weight Initialization for Hypernetworks},
  author={O. Chang and Lampros Flokas and H. Lipson},
  booktitle={International Conference on Learning Representations},
  year={2020}
}

@article{brock2021characterizing,
      title={Characterizing signal propagation to close the performance gap in unnormalized ResNets}, 
      author={Andrew Brock and Soham De and Samuel L. Smith},
      year={2021},
      eprint={2101.08692},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

@article{krueger2018bayesian,
      title={Bayesian Hypernetworks}, 
      author={David Krueger and Chin-Wei Huang and Riashat Islam and Ryan Turner and Alexandre Lacoste and Aaron Courville},
      year={2018},
      eprint={1710.04759},
      journal={arXiv preprint arXiv:1710.04759},
      archivePrefix={arXiv},
      primaryClass={stat.ML}
}

@article{rubanova2019latent,
      title={Latent ODEs for Irregularly-Sampled Time Series}, 
      author={Yulia Rubanova and Ricky T. Q. Chen and David Duvenaud},
      year={2019},
      eprint={1907.03907},
      journal={arXiv preprint arXiv:1907.03907},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

@article{chen2018neural,
  title={Neural Ordinary Differential Equations},
  author={Chen, Ricky T. Q. and Rubanova, Yulia and Bettencourt, Jesse and Duvenaud, David},
  journal={Advances in Neural Information Processing Systems},
  year={2018}
}

@inproceedings{roeder2017sticking,
  title={Sticking the landing: Simple, lower-variance gradient estimators for variational inference},
  author={Roeder, Geoffrey and Wu, Yuhuai and Duvenaud, David K},
  booktitle={Advances in Neural Information Processing Systems},
  year={2017}
}

% reliability diagrams
@inproceedings{niculescu2005predicting,
  title={Predicting good probabilities with supervised learning},
  author={Niculescu-Mizil, Alexandru and Caruana, Rich},
  booktitle={Proceedings of the 22nd international conference on Machine learning},
  pages={625--632},
  year={2005}
}

@inproceedings{zadrozny2002transforming,
  title={Transforming Classifier Scores into Accurate Multiclass Probability Estimates},
  author={Zadrozny, Bianca and Elkan, Charles},
  booktitle={Proceedings of the Eighth ACM SIGKDD International Conference on Knowledge Discovery and Data Mining},
  year={2002}
}

@inproceedings{guo2017calibration,
  title={On calibration of modern neural networks},
  author={Guo, Chuan and Pleiss, Geoff and Sun, Yu and Weinberger, Kilian Q},
  booktitle={Proceedings of the 34th International Conference on Machine Learning-Volume 70},
  pages={1321--1330},
  year={2017},
  organization={JMLR. org}
}

% complex posteriors

@inproceedings{mishkin2018slang,
  title={Slang: Fast structured covariance approximations for bayesian deep learning with natural gradient},
  author={Mishkin, Aaron and Kunstner, Frederik and Nielsen, Didrik and Schmidt, Mark and Khan, Mohammad Emtiyaz},
  booktitle={Advances in Neural Information Processing Systems},
  pages={6245--6255},
  year={2018}
}

@article{zhang2017noisy,
  title={Noisy natural gradient as variational inference},
  author={Zhang, Guodong and Sun, Shengyang and Duvenaud, David and Grosse, Roger},
  journal={arXiv preprint arXiv:1712.02390},
  year={2017}
}

@article{dandekar2020bayesian,
  title={Bayesian Neural Ordinary Differential Equations}, 
  author={Raj Dandekar and Vaibhav Dixit and Mohamed Tarek and Aslan Garcia-Valadez and Chris Rackauckas},
  year={2020},
  eprint={2012.07244},
  archivePrefix={arXiv},
  journal={arXiv preprint arXiv:2012.07244},
  primaryClass={cs.LG}
}

@inproceedings{louizos2017multiplicative,
  title={Multiplicative normalizing flows for variational bayesian neural networks},
  author={Louizos, Christos and Welling, Max},
  booktitle={Proceedings of the 34th International Conference on Machine Learning-Volume 70},
  pages={2218--2227},
  year={2017},
  organization={JMLR. org}
}

@inproceedings{ioffe2015batch,
  title={Batch normalization: Accelerating deep network training by reducing internal covariate shift},
  author={Ioffe, Sergey and Szegedy, Christian},
  booktitle={International conference on machine learning},
  pages={448--456},
  year={2015},
  organization={PMLR}
}

@article{hendrycks2019robustness,
  title={Benchmarking Neural Network Robustness to Common Corruptions and Perturbations},
  author={Dan Hendrycks and Thomas Dietterich},
  journal={Proceedings of the International Conference on Learning Representations},
  year={2019}
}

@article{anodev2,
  author    = {Tianjun Zhang and
               Zhewei Yao and
               Amir Gholami and
               Kurt Keutzer and
               Joseph Gonzalez and
               George Biros and
               Michael W. Mahoney},
  title     = {{ANODEV2:} {A} Coupled Neural {ODE} Evolution Framework},
  journal   = {CoRR},
  volume    = {abs/1906.04596},
  year      = {2019},
  url       = {http://arxiv.org/abs/1906.04596},
  archivePrefix = {arXiv},
  eprint    = {1906.04596},
  timestamp = {Fri, 01 May 2020 13:53:46 +0200},
  biburl    = {https://dblp.org/rec/journals/corr/abs-1906-04596.bib},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}

@article{dupont2019augmented,
  title={Augmented Neural ODEs},
  author={Emilien Dupont and Arnaud Doucet and Yee Whye Teh},
  journal={Advances in Neural Information Processing Systems},
  year={2019}
}

@article{Kingma2015AdamAM,
  title={Adam: A Method for Stochastic Optimization},
  author={Diederik P. Kingma and Jimmy Ba},
  journal={CoRR},
  year={2015},
  volume={abs/1412.6980}
}

@article{zhang2019approximation,
  title={Approximation capabilities of neural ordinary differential equations},
  author={Zhang, Han and Gao, Xi and Unterman, Jacob and Arodz, Tom},
  journal={arXiv preprint arXiv:1907.12998},
  year={2019}
}

@article{mackay1992practical,
  title={A practical Bayesian framework for backpropagation networks},
  author={MacKay, David JC},
  journal={Neural computation},
  volume={4},
  number={3},
  pages={448--472},
  year={1992},
  publisher={MIT Press}
}

@book{neal2012bayesian,
  title={Bayesian learning for neural networks},
  author={Neal, Radford M},
  volume={118},
  year={1996},
  publisher={Springer Science \& Business Media}
}

@inproceedings{hinton1993keeping,
  title={Keeping the neural networks simple by minimizing the description length of the weights},
  author={Hinton, Geoffrey E and Van Camp, Drew},
  booktitle={Proceedings of the sixth annual conference on Computational learning theory},
  pages={5--13},
  year={1993}
}

@inproceedings{he2016deep,
  title={Deep residual learning for image recognition},
  author={He, Kaiming and Zhang, Xiangyu and Ren, Shaoqing and Sun, Jian},
  booktitle={Proceedings of the IEEE conference on computer vision and pattern recognition},
  pages={770--778},
  year={2016}
}

@article{RadialBN,
  title={Radial Bayesian Neural Networks: Beyond Discrete Support In Large-Scale Bayesian Deep Learning.},
  author={Sebastian Farquhar and Michael A. Osborne and Yarin Gal},
  journal={arXiv: Machine Learning},
  year={2019}
}

@article{li2020scalable,
  title={Scalable Gradients for Stochastic Differential Equations},
  author={Li, Xuechen and Wong, Ting-Kam Leonard and Chen, Ricky TQ and Duvenaud, David},
  journal={arXiv preprint arXiv:2001.01328},
  year={2020}
}

@article{ramachandran2017searching,
  title={Searching for activation functions},
  author={Ramachandran, Prajit and Zoph, Barret and Le, Quoc V},
  journal={arXiv preprint arXiv:1710.05941},
  year={2017}
}

@article{hoffman2013stochastic,
  title={Stochastic variational inference},
  author={Hoffman, Matthew D and Blei, David M and Wang, Chong and Paisley, John},
  journal={The Journal of Machine Learning Research},
  volume={14},
  number={1},
  pages={1303--1347},
  year={2013},
  publisher={JMLR. org}
}

@article{rezende2014stochastic,
  title={Stochastic backpropagation and approximate inference in deep generative models},
  author={Rezende, Danilo Jimenez and Mohamed, Shakir and Wierstra, Daan},
  journal={arXiv preprint arXiv:1401.4082},
  year={2014}
}

@inproceedings{louizos2017bayesian,
  title={Bayesian compression for deep learning},
  author={Louizos, Christos and Ullrich, Karen and Welling, Max},
  booktitle={Advances in neural information processing systems},
  pages={3288--3298},
  year={2017}
}

@inproceedings{balan2015bayesian,
  title={Bayesian dark knowledge},
  author={Balan, Anoop Korattikara and Rathod, Vivek and Murphy, Kevin P and Welling, Max},
  booktitle={Advances in Neural Information Processing Systems},
  pages={3438--3446},
  year={2015}
}

@article{wang2018adversarial,
  title={Adversarial distillation of bayesian neural network posteriors},
  author={Wang, Kuan-Chieh and Vicol, Paul and Lucas, James and Gu, Li and Grosse, Roger and Zemel, Richard},
  journal={arXiv preprint arXiv:1806.10317},
  year={2018}
}

@article{massaroli2020dissecting,
  title={Dissecting neural odes},
  author={Massaroli, Stefano and Poli, Michael and Park, Jinkyoo and Yamashita, Atsushi and Asama, Hajime},
  journal={arXiv preprint arXiv:2002.08071},
  year={2020}
}

@article{innes2019zygote,
  title={Zygote: A differentiable programming system to bridge machine learning and scientific computing},
  author={Innes, Mike and Edelman, Alan and Fischer, Keno and Rackauckus, Chris and Saba, Elliot and Shah, Viral B and Tebbutt, Will},
  journal={arXiv preprint arXiv:1907.07587},
  pages={140},
  year={2019}
}

@misc{peluchetti2020doubly,
      title={Doubly infinite residual networks: a diffusion process approach}, 
      author={Stefano Peluchetti and Stefano Favaro},
      year={2020},
      eprint={2007.03253},
      archivePrefix={arXiv},
      primaryClass={stat.ML}
}

@inproceedings{peluchetti2020infinitely,
  title={Infinitely deep neural networks as diffusion processes},
  author={Peluchetti, Stefano and Favaro, Stefano},
  booktitle={International Conference on Artificial Intelligence and Statistics},
  pages={1126--1136},
  year={2020},
  organization={PMLR}
}

@article{kong2020sde,
  title={SDE-Net: Equipping Deep Neural Networks with Uncertainty Estimates},
  author={Kong, Lingkai and Sun, Jimeng and Zhang, Chao},
  journal={arXiv preprint arXiv:2008.10546},
  year={2020}
}

@article{liu2019neural,
  title={Neural sde: Stabilizing neural ode networks with stochastic noise},
  author={Liu, Xuanqing and Xiao, Tesi and Si, Si and Cao, Qin and Kumar, Sanjiv and Hsieh, Cho-Jui},
  journal={arXiv preprint arXiv:1906.02355},
  year={2019}
}


@article{blundell2015weight,
  title={Weight uncertainty in neural networks},
  author={Blundell, Charles and Cornebise, Julien and Kavukcuoglu, Koray and Wierstra, Daan},
  journal={International Conference on Machine Learning},
  year={2015},
  organization={PMLR}
}

@article{casale2018gaussian,
  title={Gaussian Process Prior Variational Autoencoders},
  author={Casale, Francesco Paolo and Dalca, Adrian V and Saglietti, Luca and Listgarten, Jennifer and Fusi, Nicolo},
  journal={32nd Conference on Neural Information Processing Systems},
  year={2018}
}

@article{yildiz2019ode,
  title={ODE$^2$VAE: Deep generative second order ODEs with Bayesian neural networks},
  author={Y{\i}ld{\i}z, {\c{C}}a{\u{g}}atay and Heinonen, Markus and L{\"a}hdesm{\"a}ki, Harri},
  journal={arXiv preprint arXiv:1905.10994},
  year={2019}
}

@article{E2017,
author = {E, Weinan},
journal = {Commun. Math. Stat.},
number = {1},
pages = {1--11},
publisher = {Springer Berlin Heidelberg},
title = {{A Proposal on Machine Learning via Dynamical Systems}},
volume = {5},
year = {2017}
}

@article{haber2017stable,
  title={Stable architectures for deep neural networks},
  author={Haber, Eldad and Ruthotto, Lars},
  journal={Inverse Problems},
  volume={34},
  number={1},
  pages={014004},
  year={2017},
  publisher={IOP Publishing}
}

@inproceedings{zhang2018noisy,
  title={Noisy natural gradient as variational inference},
  author={Zhang, Guodong and Sun, Shengyang and Duvenaud, David and Grosse, Roger},
  booktitle={International Conference on Machine Learning},
  pages={5852--5861},
  year={2018}
}

@inproceedings{graves2011practical,
  title={Practical variational inference for neural networks},
  author={Graves, Alex},
  booktitle={Advances in neural information processing systems},
  pages={2348--2356},
  year={2011}
}

@article{tucker2018doubly,
  title={Doubly reparameterized gradient estimators for monte carlo objectives},
  author={Tucker, George and Lawson, Dieterich and Gu, Shixiang and Maddison, Chris J},
  journal={arXiv preprint arXiv:1810.04152},
  year={2018}
}

@article{zhang2019anodev2,
  title={Anodev2: A coupled neural ode evolution framework},
  author={Zhang, Tianjun and Yao, Zhewei and Gholami, Amir and Keutzer, Kurt and Gonzalez, Joseph and Biros, George and Mahoney, Michael},
  journal={arXiv preprint arXiv:1906.04596},
  year={2019}
}

@article{ha2016hypernetworks,
  title={Hypernetworks},
  author={Ha, David and Dai, Andrew and Le, Quoc V},
  journal={arXiv preprint arXiv:1609.09106},
  year={2016}
}

@article{dinh2016density,
  title={Density estimation using real nvp},
  author={Dinh, Laurent and Sohl-Dickstein, Jascha and Bengio, Samy},
  journal={arXiv preprint arXiv:1605.08803},
  year={2016}
}

@article{ma2015complete,
      title={A Complete Recipe for Stochastic Gradient MCMC}, 
      author={Yi-An Ma and Tianqi Chen and Emily B. Fox},
      year={2015},
      journal={arXiv preprint arXiv:1506.04696},
      eprint={1506.04696},
      archivePrefix={arXiv},
      primaryClass={math.ST}
}

@article{izmailov2021bayesian,
      title={What Are Bayesian Neural Network Posteriors Really Like?}, 
      author={Pavel Izmailov and Sharad Vikram and Matthew D. Hoffman and Andrew Gordon Wilson},
      year={2021},
      journal={International Conference on Learning Representations},
      eprint={2104.14421},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

@article{look2019differential,
  title={Differential Bayesian neural nets},
  author={Look, Andreas and Kandemir, Melih},
  journal={arXiv preprint arXiv:1912.00796},
  year={2019}
}

@inproceedings{glorot2010understanding,
  title={Understanding the difficulty of training deep feedforward neural networks},
  author={Glorot, Xavier and Bengio, Yoshua},
  booktitle={Proceedings of the thirteenth international conference on artificial intelligence and statistics},
  pages={249--256},
  year={2010},
  organization={JMLR Workshop and Conference Proceedings}
}

@article{zhang2019fixup,
  title={Fixup initialization: Residual learning without normalization},
  author={Zhang, Hongyi and Dauphin, Yann N and Ma, Tengyu},
  journal={arXiv preprint arXiv:1901.09321},
  year={2019}
}

@article{de2019gru,
  title={Gru-ode-bayes: Continuous modeling of sporadically-observed time series},
  author={De Brouwer, Edward and Simm, Jaak and Arany, Adam and Moreau, Yves},
  journal={arXiv preprint arXiv:1905.12374},
  year={2019}
}

@inproceedings{DuvRipAdaGha2014,
  title={Avoiding pathologies in very deep networks},
  author={David Duvenaud and Oren Rippel and Ryan P. Adams and Zoubin Ghahramani},
  booktitle={Artificial Intelligence and Statistics},
  year={2014}
}

@article{song2020score,
  title={Score-Based Generative Modeling through Stochastic Differential Equations},
  author={Song, Yang and Sohl-Dickstein, Jascha and Kingma, Diederik P and Kumar, Abhishek and Ermon, Stefano and Poole, Ben},
  journal={International Conference on Learning Representations},
  year={2020}
}

@inproceedings{dusenberry2020efficient,
  title={Efficient and scalable bayesian neural nets with rank-1 factors},
  author={Dusenberry, Michael and Jerfel, Ghassen and Wen, Yeming and Ma, Yian and Snoek, Jasper and Heller, Katherine and Lakshminarayanan, Balaji and Tran, Dustin},
  booktitle={International conference on machine learning},
  pages={2782--2792},
  year={2020},
  organization={PMLR}
}

@article{osawa2019practical,
  title={Practical deep learning with bayesian principles},
  author={Osawa, Kazuki and Swaroop, Siddharth and Jain, Anirudh and Eschenhagen, Runa and Turner, Richard E and Yokota, Rio and Khan, Mohammad Emtiyaz},
  journal={arXiv preprint arXiv:1906.02506},
  year={2019}
}

@article{boue1998variational,
  title={A variational representation for certain functionals of Brownian motion},
  author={Bou{\'e}, Michelle and Dupuis, Paul and others},
  journal={The Annals of Probability},
  volume={26},
  number={4},
  pages={1641--1659},
  year={1998},
  publisher={Institute of Mathematical Statistics}
}