 All convolutional layers use $3\times3\times3$ filters, and we choose PReLU~\cite{PReLU} as the non-linear activation layer. Pooling and unpooling layers are usually used for multi-scale image processing, \xy{but unlike~\cite{yang2016} where unpooling based on pooling index is straightforward}, they are not suitable for our formulation since the two encoders perform pooling independently. Therefore, we follow the idea of~\cite{SpringenbergDBR14} and use convolutional layers with a stride of 2 and deconvolution layers~\cite{Long2015} as surrogates for pooling and unpooling layers. I.e., the network learns its pooling and unpooling operations during training. For training, we use the L1 norm as our similarity criterion. To predict the full image momentum during testing, we use a sliding window approach, predicting the momentum patch-by-patch, followed by averaging the momenta in overlapping areas. We refer to this architecture as our \emph{deterministic network.}

\begin{figure}[tb]
\centering
\includegraphics[width=.7\columnwidth]{ISBI_network_AS.pdf}
\caption{3D Bayesian network structure. The network takes 2 3D patches from moving/target images and generates 3 3D initial momentum patches (in 3 dim.). Numbers for the input and output denote input/output patch size, and numbers in the network indicate the number of features (output channels) for the conv. layers. To create a deterministic network, we remove all dropout layers.}
\label{fig:network}
\vspace{-0.4cm}
\end{figure}
We extend the network to a \emph{Bayesian network} by adding dropout layers~\cite{srivastava14a} after all convolution+PReLU layers except for the ``pooling/unpooling" (convolutional) layers~\cite{Gal2015Bayesian}. Given the input($I$)/output($O$) of a convolutional layer and the convolutional layer's weight, $W$, adding dropout can be regarded as performing variational inference for the posterior distribution of the conv layer's weight $p(W|I, O)$ using a Bernoulli distribution $q(W) = W * (\text{Bernoulli}(p)_i)|_{i=1}^N$, where $p$ is the dropout probability, $N$ is the number of nodes in the conv layer, and $\text{Bernoulli}(p)_i$ randomly sets the $i$th node in the conv layer to 0. According to \cite{Gal2015Bayesian}, training this dropout network is equivalent to minimizing the KL-divergence between the variational and true posterior $\text{KL}(q(W)|p(W|I, O))$. During testing, we keep the dropout layers and sample the network to obtain multiple momentum predictions for a single image. We choose the sample mean as the final momentum prediction and perform LDDMM shooting (by integrating Eq.~\eqref{eqn:forward}) for the momentum samples to generate multiple deformation fields. The variances of the deformation fields then estimate the predicted deformation \emph{uncertainty}. We set the dropout probability to $0.3$.

\noindent
\textbf{Patch pruning.}
We achieve a substantial speedup in computation time by using a \emph{patch pruning} strategy. Specifically, we apply a large sliding window stride and ignore patches only containing the background of both the moving image and the target image. In our experiments, these technique reduced the number of patches to predict by approximately 99.85\% for $229\times193\times193$ 3D brain images with $15\times15\times15$ patch size and a sliding window stride of 14.

\vspace{-0.2cm}
\section{EXPERIMENTS}
\label{sec:experiments}
\begin{figure*}[t!]
  \footnotesize
\centering
\begin{tabularx}{0.78\textwidth}{|r|c *{6}{|Y}|}
\hline
& \multicolumn{7}{c|}{\textbf{Deformation Error w.r.t LDDMM optimization on T1w-T1w data} [voxel]}\\ \hline
\textsl{Data Percentile} & 0.3\% & 5\% & 25\% & 50\% & 75\% & 95\% & 99.7\%\\ \hline
Affine (Baseline) & 0.1664 & 0.46 & 0.9376 & 1.4329 & 2.0952 & 3.5037 & 6.2576\\ \hline
\textbf{Ours}, T1w-T1w data  & \cellcolor{green!40}{0.0353} & \cellcolor{green!40}{0.0951} & \cellcolor{green!40}{0.1881} & \cellcolor{green!40}{0.2839} & \cellcolor{green!40}{0.416} & \cellcolor{green!40}{0.714} & \cellcolor{green!40}{1.4409} \\ \hline
\cite{yang2016}, T1w-T2w data & 0.0582 & 0.1568 & 0.3096 & 0.4651 & 0.6737 & 1.1106 & 2.0628 \\ \hline
\textbf{Ours}, T1w-T2w data & \cellcolor{green!40}{0.0551} & \cellcolor{green!40}{0.1484} & \cellcolor{green!40}{0.2915} & \cellcolor{green!40}{0.4345} & \cellcolor{green!40}{0.6243} & \cellcolor{green!40}{1.0302} & \cellcolor{green!40}{2.0177} \\ \hline
\textbf{Ours}, T1w-T2w data, 10 images & \cellcolor{green!40}{0.0663} & \cellcolor{green!40}{0.1782} & \cellcolor{green!40}{0.3489} & \cellcolor{green!40}{0.5208} & \cellcolor{green!40}{0.752} & \cellcolor{green!40}{1.2421} & \cellcolor{green!40}{2.3454} \\ \hline
& \multicolumn{7}{c|}{\textbf{Prediction/Optimization error between T1w-T2w and T1w-T1w} [voxel]} \\ \hline
\textsl{Data Percentile} & 0.3\% & 5\% & 25\% & 50\% & 75\% & 95\% & 99.7\% \\ \hline
\textbf{Ours} & \cellcolor{green!40}{0.0424} & \cellcolor{green!40}{0.1152}  & \cellcolor{green!40}{0.2292} & \cellcolor{green!40}{0.3444} & \cellcolor{green!40}{0.4978} & \cellcolor{green!40}{0.8277} & \cellcolor{green!40}{1.6959}\\ \hline
\texttt{NiftyReg} (Baseline) & 0.2497 & 0.7463 & 1.8234 & 3.1719 & 5.1124 & 8.9522 & 14.4666 \\ \hline
\end{tabularx}
\captionof{table}{Evaluation results for the 3D dataset.}
\label{table:3D}
\centering
\includegraphics[width=0.78\textwidth]{example_image_other_8figure_small.pdf}
\captionof{figure}{Exemplary test case. (a) T1w moving image; (b) T2w target image; (c)-(d) deformation prediction result from T1w-T1w/T1w-T2w data; (e) deformation result by LDDMM optimization for T1w-T1w registration; (f) uncertainty of predicted T1w-T2w deformation as the 2-norm of the sum of variances of deformation fields in all spatial directions, mapped on the predicted T1w-T2w wrapped image. Yellow = more uncertainty, blue = less uncertainty; (g)+(h) \texttt{NiftyReg} registration result for T1w-T1w/T1w-T2w pair.}
\label{fig:example}
\vspace{-0.3cm}
\end{figure*}
We assess our approach on the IBIS 3D Autism Brain image dataset~\cite{Wolff2046}, 
containing 375 T1w/T1w brain images ($229\times193\times193$) of 2 years subject with pediatric disease. We first register all images affinely to the ICBM 152 atlas~\cite{FonovEBAMC11} and select 359 images for training and the remaining 16 images for testing. For training, we randomly select 359 T1w-T1w image pairs and perform LDDMM registration using PyCA\footnote{\url{https://bitbucket.org/scicompan