# MGAug: Multimodal Geometric Augmentation in
Latent Spaces of Image Deformations

Tonmoy Hossain1, Jian Wang3 , Miaomiao Zhang1,2

###### Abstract

Geometric transformations have been widely used to augment the size of
training images. Existing methods often assume a unimodal distribution of the
underlying transformations between images, which limits their power when data
with multimodal distributions occur. In this paper, we propose a novel model,
Multimodal Geometric Augmentation (MGAug), that for the first time generates
augmenting transformations in a multimodal latent space of geometric
deformations. To achieve this, we first develop a deep network that embeds the
learning of latent geometric spaces of diffeomorphic transformations (a.k.a.
diffeomorphisms) in a variational autoencoder (VAE). A mixture of multivariate
Gaussians is formulated in the tangent space of diffeomorphisms and serves as
a prior to approximate the hidden distribution of image transformations. We
then augment the original training dataset by deforming images using randomly
sampled transformations from the learned multimodal latent space of VAE. To
validate the efficiency of our model, we jointly learn the augmentation
strategy with two distinct domain-specific tasks: multi-class classification
on 2D synthetic datasets and segmentation on real 3D brain magnetic resonance
images (MRIs). We also compare MGAug with state-of-the-art transformation-
based image augmentation algorithms. Experimental results show that our
proposed approach outperforms all baselines by significantly improved
prediction accuracy. Our code is publicly available at
https://github.com/tonmoy-hossain/MGAug.

## Introduction

Recent advances in deep learning (DL) have led to remarkable progress in image
analysis tasks, including but not limited to image classification (Krizhevsky,
Sutskever, and Hinton 2017; Wang and Zhang 2022), segmentation (Hossain et al.
2019; Chen et al. 2021), object detection (Kim and Pavlovic 2016; Pang et al.
2019), and 3D reconstruction (Jayakumar, Hossain, and Zhang 2023). The
literature has shown that deep neural networks (DNNs) are most powerful when
trained on vast amounts of manually labeled data. Unfortunately, many image
applications have relatively small annotated datasets available, due to
privacy concerns, regulations, costs, and incidents of disease. The DNNs, in
such cases, tend to overfit due to their poor generalizations to large data
variances. Data augmentation (DA) is a simple, yet efficient and commonly used
strategy to overcome this issue by synthesizing more training examples to
artificially increase the dataset (Cubuk et al. 2019; Zhao et al. 2019; Zoph
et al. 2020; Shen et al. 2020). It is often employed as a preprocessing step
applied to the training data to expose the DNNs to various types of
approximated invariance; hence reducing the problem of overfitting. In the
image domain, common augmentation strategies include adding variations to
image intensities (Gao et al. 2021), or deforming the image by randomly
generated geometric transformations (i.e., translation, rotation, or scaling)
(Moreno-Barea et al. 2018). This paper focuses on DA, particularly in
geometric spaces.

In practice, DA is usually a manual process where a human specifies a set of
affine or deformable transformations believed to make DL models invariant
(Hendrycks et al. 2019; Yun et al. 2019). However, such methods can pose a
high risk of producing unrealistic or implausible variations of the image. For
example, rotating an image 180 degrees makes sense for natural images, but
produces ambiguities on digital numbers (i.e., the digits ‘6’ and ‘9’). In
medical domain, randomly generated deformations may transform images from a
healthy to a diseased group. To alleviate this issue, researchers have
developed a group of automatic DA methods that learn an optimal augmenting
policy from the data itself (Cubuk et al. 2019; Lim et al. 2019; Zoph et al.
2020; Araslanov and Roth 2021; Muralikrishnan et al. 2022). One approach is to
learn a generator via generative adversarial networks (GANs) to synthesize
data from scratch (Mirza and Osindero 2014; Chen et al. 2022). However, these
models are difficult to train due to the high complexity of largely unknown
parameter spaces (Jaderberg et al. 2015). Despite showing encouraging results
on relatively easy datasets (e.g., $28^{2}$ MNIST), they do not appear to be
applicable to more complex datasets with large structural and geometric
variances. Another common approach is to learn a probabilistic distribution of
a set of deformable transformations derived from the training images (Hauberg
et al. 2016; Krebs et al. 2019; Dalca et al. 2019; Zhao et al. 2019; Olut et
al. 2020; Shen et al. 2020). Most importantly, these methods assume a unimodal
distribution of the underlying transformations between images. They do not
provide an appropriate way to learn the full space of geometric
transformations among the training data. This greatly limits their power when
multimodal data (e.g., time-series/longitudinal images, or images of multiple
sub-populations) occurs.

In this paper, we present a novel model, MGAug, that for the first time learns
augmenting transformations in a multimodal latent space of geometric
deformations. In contrast to existing approaches limited to data with unimodal
distributions, our MGAug features the main advantages of

  1. (i)

Generalizing the automatic DA process into a geometric space that allows
hidden transformations with multiple modes.

  2. (ii)

Developing a new generative model that learns the latent distribution of
augmented transformations via a mixture of multivariate Gaussians defined as a
prior in the tangent space of diffeomorphisms.

  3. (iii)

Improving the quality of augmented samples from learned knowledge of training
data itself with intra-group constraints.

We demonstrate the effectiveness of our proposed model in two distinct domain-
specific tasks: multi-class classification and segmentation on both 2D
synthetic data and real 3D brain magnetic resonance images (MRIs),
respectively. We also compare MGAug with the state-of-the-art transformation-
based augmentation algorithms (Olut et al. 2020; Dalca et al. 2019).
Experimental results show that our method outperforms all baseline algorithms
by significantly improved prediction accuracy.

## Background: Geometric Deformations via Template-based Image Registration

We first briefly review the concept of template-based deformable image
registration (Joshi et al. 2004; Zhang, Singh, and Fletcher 2013), which is
widely used to estimate deformable geometric deformations from groupwise
images. In this work, we adopt the algorithm of large deformation
diffeomorphic metric mapping (LDDMM), which encourages the smoothness of
deformation fields and preserves the topological structures of images (Beg et
al. 2005).

Given a number of $N$ images $\\{I_{1},\cdots,I_{N}\\}$, the problem of
template-based image registration is to find the shortest paths (a.k.a.
geodesics) between a template image $I$ and the rest of image data with
derived associate deformation fields $\phi_{1},\cdots,\phi_{N}$ that minimizes
the energy function

$E(I,\phi_{n})=\sum_{n=1}^{N}\frac{1}{\sigma^{2}}\text{Dist}[I\circ\phi_{n}(v_{n}(t)),I_{n}]\,+\,\text{Reg}[\phi_{n}(v_{n}(t))],$
(1)

where $\sigma^{2}$ is a noise variance and $\circ$ denotes an interpolation
operator that deforms image $I$ with an estimated transformation $\phi_{n}$.
The $\text{Dist}[\cdot,\cdot]$ is a distance function that measures the
dissimilarity between images, i.e., sum-of-squared differences (Beg et al.
2005), normalized cross correlation (Avants et al. 2008), and mutual
information (Wells III et al. 1996). The ${\rm Reg}[\cdot]$ is a regularizer
that guarantees the smoothness of transformations, i.e.,

$\displaystyle{\rm
Reg}[\phi_{n}(v_{n}(t))]=\int_{0}^{1}(Lv_{n}(t),v_{n}(t))\,dt,$
$\displaystyle\text{with}\quad\frac{d\phi_{n}(t)}{dt}=-D\phi_{n}(t)\cdot
v_{n}(t),$ (2)

where $L:V\rightarrow V^{*}$ is a symmetric, positive-definite differential
operator that maps a tangent vector $v_{n}(t)\in V$ into its dual space as a
momentum vector $m_{n}(t)\in V^{*}$. We write $m_{n}(t)=Lv_{n}(t)$, or
$v_{n}(t)=Km_{n}(t)$, with $K$ being an inverse operator of $L$. The operator
$D$ denotes a Jacobian matrix and $\cdot$ represents element-wise matrix
multiplication. In this paper, we use a metric of the form
$L=-\alpha\Delta+\mathbb{I}$, in which $\Delta$ is the discrete Laplacian
operator, $\alpha$ is a positive regularity parameter, and $\mathbb{I}$
denotes an identity matrix.

The minimum of Eq. (2) is uniquely determined by solving a Euler-Poincaré
differential equation (EPDiff) (Arnold 1966; Miller, Trouvé, and Younes 2006)
with a given initial condition of velocity fields, noted as $v_{n}(0)$,

$\displaystyle\frac{\partial v_{n}(t)}{\partial t}$
$\displaystyle=-K[(Dv_{n}(t))^{T}\cdot m_{n}(t)+Dm_{n}(t)\cdot v_{n}(t)$
$\displaystyle+m_{n}(t)\cdot\operatorname{div}v_{n}(t)],$ (3)

where the operator $D$ denotes a Jacobian matrix, $\operatorname{div}$ is the
divergence, and $\cdot$ represents an element-wise matrix multiplication. This
is known as the geodesic shooting algorithm (Vialard et al. 2012), which
nicely proves that the deformation-based shape descriptor, $\phi_{n}$, can be
fully characterized by an initial velocity field, $v_{n}(0)$.

We are now able to equivalently minimize the energy function in Eq. (1) as

$\displaystyle E(I,v_{n}(0))$
$\displaystyle=\sum_{n=1}^{N}\frac{1}{\sigma^{2}}\text{Dist}[I\circ\phi_{n}(v_{n}(0)),I_{n}]$
$\displaystyle+(Lv_{n}(0),v_{n}(0)),\,\,\text{s.t. Eq.~{}\eqref{eq:epdiff}.}$
(4)

For notation simplicity, we will drop the time index in the following
sections.

## Our Method: MGAug

This section presents a novel algorithm, MGAug, that learns a multimodal
latent distribution of geometric transformations represented by the initial
velocity fields $v_{0}$. We first introduce a generative model in latent
geometric spaces for DA based on VAE, followed by an approximated variational
inference. The augmented data is then used to train various image analysis
tasks, i.e., image classification or segmentation. Similar to (Zhao et al.
2019; Chen et al. 2022), we employ a joint learning scheme to optimize the
MGAug and IAT network parameters. To simplify the notation, we omit the time
index of $v_{0}$ in the following sections.

### Generative Model in Latent Multimodal Geometric Spaces

Given a group of $N$ input images $\mathbf{J}$ with a template $I$, we assume
the latent geometric transformations $\mathbf{x}$ are generated from a number
of $C$ clusters. The distribution of $\mathbf{x}$ can be formulated as a
Gaussian mixture with means $\boldsymbol{\mu}$ and variances
$\mathbf{\Sigma}$, specified by a VAE encoder parameterized by $\beta$, i.e.,
$p_{\theta}(\mathbf{x}|{\boldsymbol{\epsilon}})$, where
$\boldsymbol{\epsilon}$ is a random number generated from
$\mathcal{N}(0,\mathbb{I})$. That is to say, the VAE encoder outputs a set of
$C$ means $\boldsymbol{\mu}$ and $C$ variances $\mathbf{\Sigma}$. A one-hot
vector $\mathbf{z}$ is sampled from the mixing probability
$\mathbf{\boldsymbol{\pi}}\in[0,1]$ with
$\sum_{c=1}^{C}\mathbf{\boldsymbol{\pi}}^{c}=1$, which chooses one component
from the Gaussian mixture. In this paper, we set the parameter
$\mathbf{\boldsymbol{\pi}}=C^{-1}$ to make the variable $\mathbf{z}$ uniformly
distributed, which is similar to (Dilokthanakul et al. 2016). The initial
velocity field $\mathbf{v}$ is generated from a VAE decoder parameterized by
$\theta$ and the continuous latent variable $\mathbf{x}$. The entire process
to generate the observed image data $\mathbf{J}$ from a set of latent
variables
$\\{\mathbf{v},\mathbf{x},\mathbf{\boldsymbol{\epsilon}},\mathbf{z}\\}$ is:

$\displaystyle\mathbf{\boldsymbol{\epsilon}}$
$\displaystyle\sim\mathcal{N}(0,\mathbb{I}),$ (5) $\displaystyle\mathbf{z}$
$\displaystyle\sim\text{Mult}(\mathbf{\boldsymbol{\pi}}),$ (6)
$\displaystyle\mathbf{x}|\mathbf{z},\mathbf{\boldsymbol{\epsilon}}$
$\displaystyle\sim\prod_{c=1}^{C}\mathcal{N}(\mathbf{\boldsymbol{\mu}}_{\mathbf{z}_{c}}(\mathbf{\boldsymbol{\epsilon}},\beta),\mathbf{\Sigma}_{\mathbf{z}_{c}}(\mathbf{\boldsymbol{\epsilon}},\beta))^{\mathbf{z}_{c}},$
(7) $\displaystyle\mathbf{v}|\mathbf{x}$
$\displaystyle\sim\mathcal{N}(\mathbf{v}_{\theta}(\mathbf{x})\,|\,0,\mathbf{\Sigma}),$
(8) $\displaystyle\mathbf{J}|\mathbf{v},I,\mathbf{\lambda}$
$\displaystyle\sim\mathcal{N}(\mathbf{J}\,|\,I\circ\mathbf{\phi}(\mathbf{v}),\boldsymbol{\lambda}^{2}\mathbb{I}),$
(9)

where $\text{Mult}(\cdot)$ denotes a multinomial distribution,
$\mathbf{\boldsymbol{\lambda}}^{2}$ is a noise variance in the observed image
space, and an identity matrix $\mathbb{I}$ maps a matrix to be diagonal. A
graphical representation of our proposed generative model is shown in Fig. 1.

Figure 1: Graphical representation of our proposed generative model in
multimodal latent space.

To guarantee the smoothness of the geometric transformations, we define a
multivariate Gaussian distribution on the generated initial velocity fields
$\mathbf{v}$ in Eq. (8) as

$p(\mathbf{v}|\mathbf{x})\propto\exp\left(-\frac{1}{2}(\mathbf{\Sigma}^{-1}\mathbf{v}_{\theta}(\mathbf{x}),\mathbf{v}_{\theta}(\mathbf{x}))\right).$
(10)

Similar to (Zhang and Fletcher 2013; Zhang, Shao, and Fletcher 2015), we
define the covariance matrix $\boldsymbol{\Sigma}$ as an inverse operator of
the differential operator, $L$, which is used in the registration regularity
in Eq. (2).

Considering that our input images are measured on a discrete grid, we
formulate our image noise model as an i.i.d. Gaussian at each image voxel,
i.e.,

$\displaystyle
p(\mathbf{J}|\mathbf{v},I,\mathbf{\boldsymbol{\lambda}})=\frac{1}{(2\pi)^{M/2}(\boldsymbol{\lambda})^{M}}\exp\left(-\frac{\|I\circ\mathbf{\phi}(\mathbf{v})-\mathbf{J}\|_{2}^{2}}{2\boldsymbol{\lambda}^{2}}\right),$
(11)

where $M$ is the total number of image voxels. Putting together Eq. (5) to Eq.
(11), we arrive at a generative model

$p_{\beta,\theta}(\mathbf{J},\mathbf{v},\mathbf{x},\mathbf{\boldsymbol{\epsilon}},\mathbf{z})=p(\mathbf{J}|\mathbf{v},I,\mathbf{\boldsymbol{\lambda}})\,p(\mathbf{v}|\mathbf{x})\,p(\mathbf{x}|\mathbf{z},\mathbf{\boldsymbol{\epsilon}})\,p(\mathbf{z})\,p(\mathbf{\boldsymbol{\epsilon}})$
(12)

### Variational Lower Bound

We will train the generative model in Eq. (12) with the variational inference
objective, i.e. the log-evidence lower bound (ELBO). With an assumption that
the mean-field variational family
$q(\mathbf{v},\mathbf{x},\mathbf{\boldsymbol{\epsilon}},\mathbf{z}\,|\,\mathbf{J})$
is an approximate to the posterior, we are able to formulate the ELBO as

$\mathcal{L}_{\text{ELBO}}=\mathbb{E}\left[\log\frac{p_{\beta,\theta}(\mathbf{J},\mathbf{v},\mathbf{x},\mathbf{\boldsymbol{\epsilon}},\mathbf{z})}{q(\mathbf{v},\mathbf{x},\mathbf{\boldsymbol{\epsilon}},\mathbf{z}\,|\,\mathbf{J})}\right]\\\
$ (13)

We factorize the mean-field variational family as
$\mathbf{q}=\prod_{i=1}^{n}q(\mathbf{v}_{i}|\mathbf{J}_{i},I,\boldsymbol{\lambda})\,q_{\psi_{\mathbf{x}}}(\mathbf{x}_{i}|\mathbf{v}_{i},\mathbf{J}_{i})\,q_{\psi_{\mathbf{\boldsymbol{\epsilon}}}}(\mathbf{\boldsymbol{\epsilon}}_{i}|\mathbf{v}_{i})\,p_{\beta}(\mathbf{z}_{i}|\mathbf{x}_{i},\mathbf{\boldsymbol{\epsilon}}_{i}),$
where $i$ denotes the $i$-th image. Each variational factor is parameterized
with the recognition networks $\psi_{\mathbf{x}}$ and
$\psi_{\mathbf{\boldsymbol{\epsilon}}}$ that output the parameters of the
variational Gaussian distributions. Note that we will drop the index $i$ in
the following sections for the purpose of notation simplicity.

Figure 2: An overview of our model MGAug for image analysis tasks.

Using previously defined Eq. (10) and Eq. (11), the posterior
$q(\mathbf{v}_{i}|\mathbf{J}_{i},I,\boldsymbol{\lambda})\propto
p(\mathbf{J}|\mathbf{v},I,\mathbf{\boldsymbol{\lambda}})\cdot
p(\mathbf{v}|\mathbf{x})$ is formulated by Baye’s rule. Following (Jiang et
al. 2016), we can derive the posterior
$p_{\beta}(\mathbf{z}|\mathbf{x},\mathbf{\boldsymbol{\epsilon}})$ as

$\displaystyle\centering
p_{\beta}(\mathbf{z}_{j}=1|\mathbf{x},\mathbf{\boldsymbol{\epsilon}})\@add@centering$
$\displaystyle=\frac{p(\mathbf{z}_{j}=1)\,p(\mathbf{x}|\mathbf{z}_{j}=1,\boldsymbol{\epsilon})}{\sum_{c=1}^{C}p(\mathbf{z}_{c}=1)\,p(\mathbf{x}|\mathbf{z}_{j}=1,\boldsymbol{\epsilon})}$
$\displaystyle=\frac{\boldsymbol{\pi}_{j}\,\mathcal{N}(\mathbf{x}|\boldsymbol{\mu}_{j}(\boldsymbol{\epsilon};\beta),\boldsymbol{\Sigma}_{j}(\boldsymbol{\epsilon};\beta))}{\sum_{c=1}^{C}\boldsymbol{\pi}_{\mathbf{z}_{c}}\,\mathcal{N}(\mathbf{x}|\mu_{\mathbf{z}_{c}}(\boldsymbol{\epsilon};\beta),\Sigma_{\mathbf{z}_{c}}(\boldsymbol{\epsilon};\beta))},$

where $\mathbf{z}_{j},j\in\\{1,\cdots,C\\}$ denotes the $j$-th element of
$\mathbf{z}$. The variational lower bound can be now formulated as

$\displaystyle\mathcal{L}_{ELBO}$
$\displaystyle=\mathbb{E}_{q(\mathbf{x}|\mathbf{J})}[\log
P_{\theta}(\mathbf{J}|\mathbf{x})]$
$\displaystyle-\mathbb{E}_{q(\boldsymbol{\epsilon}|\mathbf{J})\,p(\mathbf{z}|\mathbf{x},\boldsymbol{\epsilon})}[KL(q_{\psi_{\mathbf{x}}}(\mathbf{x}|\mathbf{J})\,||\,p_{\beta}(\mathbf{x}|\boldsymbol{\epsilon},\mathbf{z}))]$
$\displaystyle-
KL(q_{\psi_{\mathbf{\boldsymbol{\epsilon}}}}(\boldsymbol{\epsilon}|\mathbf{J})\,||\,p(\boldsymbol{\epsilon}))$
$\displaystyle-\mathbb{E}_{q(\mathbf{x}|\mathbf{J})\,q(\boldsymbol{\epsilon}|\mathbf{J})}[KL(p_{\beta}(\mathbf{z}|\mathbf{x},\boldsymbol{\epsilon})\,||\,p(\mathbf{z}))]$
(14)

### Joint Learning of MAGug With Image Analysis

We demonstrate the effectiveness of MGAug in image analysis tasks (IAT) by
developing a joint learning framework. It has been shown in (Zhao et al. 2019;
Chen et al. 2022) that these two tasks can be mutually beneficial. The MGAug
effectively synthesizes image samples to augment the training dataset, while
the IAT can provide extra information on rectifying the learned distributions
of latent geometric spaces. An overview of our proposed network architecture
is shown in Fig. 2.

#### Network design of MGAug.

MGAug is mainly designed based on the previously introduced inference model,
where the approximate posteriors are simulated by neural networks. Analogous
to (Kingma and Welling 2013), we let the encoder of MGAug, parameterized by
$\beta$, approximate the posterior as a mixture of multivariate Gaussians. In
the latent space $\mathbf{z}$, we use non-linear hidden layers to learn the
parameters (i.e., means and covariances) of the mixture Gaussians. The
decoder, parameterized by $\theta$, outputs augmented velocity fields
$\mathbf{v}_{n}$, followed by generating their associated deformations
$\phi_{n}$ (via Eq. (3)) and augmented deformed images. In this paper, we
adopt UNet architecture (Ronneberger, Fischer, and Brox 2015) as the network
backbone for MGAug. However, other networks such as UNet++ (Zhou et al. 2018)
and TransUNet (Chen et al. 2021) can be easily applied.

Defining
$\Theta=(\mathbf{J},\beta,\theta,\mathbf{x},\mathbf{z},\mathbf{v},\boldsymbol{\epsilon},I)$
as posterior parameters, we formulate the network loss of MGAug as a
combination of dissimilarity between deformed images and templates, a
regularity to ensure the smoothness of augmenting transformations, and the
variational inference objective as

$\displaystyle\ell_{\text{MGAug}}(\Theta)$ $\displaystyle=-\log
p_{\beta,\theta}(\mathbf{J},\mathbf{v},\mathbf{x},\mathbf{\boldsymbol{\epsilon}},\mathbf{z})$
$\displaystyle\approx\frac{\|I\circ\mathbf{\phi}(\mathbf{v}(\beta,\theta))-\mathbf{J}\|_{2}^{2}}{2\boldsymbol{\lambda}^{2}}+\mathcal{L}_{ELBO}$
$\displaystyle+\frac{1}{2}(L\mathbf{v}(\beta,\theta),\mathbf{v}(\beta,\theta))+\text{reg}(\beta,\theta),$

where $\text{reg}(\cdot)$ denotes network parameter regularizations.

#### Network design of IAT.

We demonstrate MGAug on the task of image classification and segmentation
separately. Other domain-specific tasks, such as image reconstruction, object
recognition, and tracking can be integrated similarly.

For classification, assuming we have augmented images (consist of $C$ classes
with $N_{c}$ images per class) and their labels denoted as $\hat{y}$. The
classifier is trained by minimizing the loss

$\ell_{\text{IAT}}^{clf}=-\gamma\sum_{n=1}^{N_{c}}\sum_{c=1}^{C}\hat{y}_{nc}\cdot\log[y_{nc}(\beta,\theta,w_{1})]+\text{reg}(\beta,\theta,w_{1}),$
(15)

where $y_{nc}$ denotes the predicted label, $\gamma$ is the weighting
parameter, and $w_{1}$ is a classification network parameter.

For segmentation, we employ a loss function as a combination of
Sørensen$-$Dice coefficient (Dice 1945) and cross-entropy (CE). Let $Q$ denote
the number of segmentation labels, $S_{q}$ denotes the augmented ground truth,
and $R_{q}$ denotes the prediction with the associated probabilistic map
$R^{p}_{q}$ for $q$-th label. We formulate the segmentation loss as

$\displaystyle\ell_{\text{IAT}}^{seg}$
$\displaystyle=\tau\sum_{q=1}^{Q}[1-\text{Dice}(S_{q},R_{q}(\beta,\theta,w_{2}))$
$\displaystyle+\text{CE}(S_{q},R^{p}_{q}(\beta,\theta,w_{2}))]+\text{reg}(\beta,\theta,w_{2}),$
(16)

where $\text{Dice}(S_{q},R_{q})=2(|S_{q}|\cap|R_{q}|)/(|S_{q}|+|R_{q}|)$ and
$\text{CE}(S_{q},R^{p}_{q})=-S_{q}\cdot\log R_{q}$. The $\tau$ is the
weighting parameter and $w_{2}$ denotes the segmentation network parameters.

The total loss of our joint learning framework of MGAug and IAT is
$\ell=\ell_{\text{MGAug}}+\ell_{\text{IAT}}$. We employ an alternative
optimization scheme (Nocedal and Wright 1999) to minimize the loss by
alternating between the training of the MGAug network and the IAT network
without sharing weights, making it an end-to-end network architecture (Zhao et
al. 2019; Chen et al. 2022).

## Experimental Evaluation

We validate our proposed model, MGAug, on both 2D synthetic datasets and real
3D brain MRI scans, highlighting its effectiveness on training data with
limited availability.

### Datasets

2D silhouette shapes. We choose Kimia216 (Sebastian, Klein, and Kimia 2004),
which is a relatively smaller dataset with seven classes of 2D silhouette
shape data (bone, bottle, children, face, fish, heart, and key). Each shape
class includes approximately $20$ binary images, with the size of $28^{2}$.
The total number of our training images for multi-classification is about
$140$.

2D handwritten digits. We use $70000$ handwritten digits of $10$ classes
($0$-$9$) from the MNIST repository (LeCun 1998). Each class includes
approximately $7000$ images with the size of $28^{2}$.

3D brain MRIs. We include $803$ subjects of public T1-weighted longitudinal
brain MRIs with manually delineated anatomical structures from Open Access
Series of Imaging Studies (OASIS) (LaMontagne et al. 2019), Alzheimer’s
Disease Neuroimaging Initiative (ADNI) (Jack Jr et al. 2008), Autism Brain
Imaging Data Exchange (ABIDE) (Di Martino et al. 2014), and LONI Probabilistic
Brain Atlas Individual Subject Data (LPBA40) (Shattuck et al. 2008). Due to
the difficulty of preserving the diffeomorphic property across individual
subjects with large age variations, we carefully evaluate images from subjects
aged from $60$ to $90$. All MRIs were pre-processed as $256\times 256\times
256$, $1.25mm^{3}$ isotropic voxels, underwent skull-stripping, intensity
normalization, bias field correction and pre-alignment with affine
transformations.

### Experiments and Implementation Details

We validate the proposed model, MGAug, on both image classification and
segmentation tasks. As MGAug falls into the category of transformation-based
DA schemes, we mainly focuses on comparing its performance with similar
baselines: randomly generated affine transformations (ATA, a commonly used
augmentation strategy), and two state-of-the-art unimodal generative models
based on diffeomorphic registrations (AdvEigAug or AEA, an augmentation model
based on the statistics of a learned latent space of velocity fields (Olut et
al. 2020); and UDA, a unimodal deformation-based augmentation using random
stationary velocity fields (Dalca et al. 2019)).

Evaluation of classification on small datasets. We first validate our model
and other baselines on a relatively small dataset of 7-class 2D silhouette
shapes. We adopt six classification network backbones, including a three-layer
CNN, ResNet18 (He et al. 2016), VGGNet16 (Simonyan and Zisserman 2014),
SqueezeNet (Iandola et al. 2016), DenseNet121 (Huang et al. 2017), and Vision
Transformer (ViT) (Dosovitskiy et al. 2020).

We carefully determine the optimal number of modes, $C$, for our multimodal
generative model by performing cross-validation with various values of $C$.
The value of C will vary for different image analysis tasks depending on the
dataset distributions. For classification, we assume that the value of C is
close to the number of classes. Furthermore, we conduct a comprehensive
assessment of the augmentation efficiency offered by our proposed approach.
This involves augmenting the training dataset to various magnitudes, ranging
from $2$ to $10$ times the size of the original dataset.

To examine the quality of the generated augmenting transformations, we compute
the determinant of Jacobian of the deformation fields. Throughout the
evaluation process, we employ multiple metrics such as average accuracy (Acc),
precision (Prec), and F1-score (F1) to validate the classification
performance.

Evaluation of segmentation on real 3D brain MRIs. We validate the performance
of all methods on real 3D brain images with a decent size of training data. In
this experiment, we employ TransUNet as our network backbone, given its proven
superiority in brain segmentation performance (Chen et al. 2021). To ensure a
fair comparison of the final segmentation accuracy, we augment the training
data three times for all methods, resulting in a total number of approximately
$2000$ augmented images. We then compare the segmentation dice scores of all
augmentation methods over different brain structures.

Evaluation of the quality of augmented images. To further evaluate the quality
of the augmented images, we validate MGAug on two relatively large datasets
(2D MNIST and 3D brain MRIs) with various levels of reduced training sample
sizes (i.e., $10\%$, $20\%$, and $50\%$ of the original size). We augment the
training images three times of the original size for all methods.

Parameter setting. We set parameter $\alpha=3.0$ for the operator $L$ and the
number of time steps for Euler integration in EPDiff (Eq. (3)) as $10$. We set
the noise variance $\sigma$ as $0.02$, the batch size as $16$, and use the
cosine annealing learning rate scheduler that starts from a learning rate
$\eta=1e^{-3}$ for network training. We run all the models on training and
validation images with Adam optimizer until obtaining the best validation
performance. The training and prediction procedure of all learning-based
methods are performed on a 40GB NVIDIA A100 Tensor Core GPU. For both 2D and
3D datasets, we split the images by using $70\%$ as training, $15\%$ as
validation, and $15\%$ as testing images.

Our method is flexible to either automatically estimated or pre-selected
template(s) for group-wise data. Therefore, we demonstrate estimated with-in-
class template(s) for classification tasks vs. pre-selected templates for
segmentation tasks, respectively. In all experiments presented in this paper,
we run cross-validation to determine the optimal number of modes $C$.

### Results

The left panel of Fig. 3 reports our model MGAug’s classification accuracy
(using ViT as the network backbone) on 2D silhouette shape dataset with
different number of modes, $C$, along with diverse levels of augmented images
quantities (i.e., $5\times,7\times,$ and $10\times$ the original training data
size). It suggests that the utilization of a multimodal latent space $(C>1)$
consistently yields superior performance compared to an unimodal distribution
$(C=1)$ by more than $2$-$7\%$. Our model achieves an optimal performance when
$C=7$, which represents the number of classes in the training data. The right
panel of Fig. 3 shows a comprehensive evaluation of classification accuracy
across all augmentation methods with an increased number of augmented images
in training data. Our method significantly outperforms other baselines by more
than $7\%$ of accuracy with different numbers of augmented images. Another
interesting observation is that our MGAug achieves comparable classification
accuracy when using roughly $3\times$ augmented images for training vs. other
baselines using $10\times$ of ground truth images. This indicates the potency
of our method in efficiently leveraging augmented data to achieve competitive
results.

Table 1: Accuracy comparison on 2D shapes over all models with various network
backbones.
| ATA | AEA | UDA | MGAug
---|---|---|---|---
| Acc | Prec | F1 | Acc | Prec | F1 | Acc | Prec | F1 | Acc | Prec | F1
CNN | 71.73 | 73.25 | 66.89 | 83.93 | 88.82 | 83.07 | 85.71 | 87.14 | 83.73 | 91.07 | 92.14 | 90.78
ResNet18 | 76.78 | 79.22 | 73.44 | 80.36 | 82.17 | 78.76 | 87.50 | 90.07 | 86.83 | 89.28 | 90.43 | 89.15
VGGNet16 | 73.21 | 72.46 | 71.23 | 85.71 | 88.04 | 85.85 | 85.71 | 88.57 | 81.58 | 89.28 | 91.21 | 89.22
SqueezeNet | 75.00 | 75.85 | 72.99 | 76.78 | 79.90 | 74.53 | 83.93 | 87.30 | 82.73 | 92.86 | 93.97 | 92.59
DenseNet121 | 75.89 | 77.54 | 73.88 | 87.50 | 90.79 | 86.18 | 89.28 | 90.86 | 89.08 | 94.64 | 95.55 | 94.28
ViT | 78.57 | 82.62 | 76.53 | 89.28 | 92.20 | 89.33 | 91.07 | 91.93 | 96.84 | 98.21 | 98.41 | 98.20
Figure 3: Left: Accuracy comparison on various modes $C$ on 2D shape dataset.
Right: A comparison of classification performance for all models over
increasing number of augmented data taking all ground truth images.

Tab. 1 reports the multi-class classification accuracy on a relatively smaller
2D silhouette shape data across all augmentation methods with various
classification network backbones. For a fair comparison, we augment the data
$10\times$ of the original training images. MGAug achieves the highest
performance, outperforming all baselines by $2$-$9\%$ of accuracy on six
different network backbones.

Fig. 4 visualizes examples of augmented images and the determinant of Jacobian
(DetJac) of augmented transformations produced by MGAug. The value of DetJac
indicates volume changes when applying augmented deformations on templates.
For example, there is no volume change when DetJac$=$1, while volume shrinks
when DetJac$<$1 and expands when DetJac$>$1\. The value of DetJac smaller than
zero indicates an artifact or singularity in the transformation field, i.e., a
failure to preserve the diffeomorphic property when the effect of folding and
crossing grids occurs. Our method nicely preserves the topological structures
of augmented images without introducing artifacts.

Figure 4: Examples of DetJac generated by MGAug. Left to right:
source/template, target images, augmented DetJac overlaid with template
images.

Fig. 5 reports the statistics of dice scores (mean and variance) of all
augmentation methods over fifteen brain structures experimenting on $30\%$ of
training images with $2000$ augmented images. Our MGAug outperforms all the
baselines with a significant dice improvement in all anatomical structures (at
least $2\%$ or higher).

Figure 5: Dice score comparison over fifteen anatomical brain structures, (LWM
- Left White Matter, RWM - Right White Matter, LLV - Left Lateral Ventricle,
RLV - Right Lateral Ventricle, LCC - Left Cerebellum Cortex, RCC - Right
Cerebellum Cortex, LT - Left Thalamus, RT - Right Thalamus, LC - Left Cudate,
RC - Right Caudate, LP - Left Putamen, RP - Right Putamen, LH - Left
Hippocampus, RH - Right Hippocampus, BS - Brain Stem) for four augmentation
methods.

The left panel of Fig. 6 showcases the classification results on the 2D MNIST
dataset for all methods, using different percentages of ground truth data
during the training phase. Our MGAug method consistently achieves the highest
accuracy across all augmentation techniques. Remarkably, MGAug reaches an
accuracy of $98.02\%$ (which is the baseline accuracy without any
augmentation) using only $20\%$ of the ground truth data, whereas other
methods require at least $50\%$ of ground truth data to achieve similar
results. Similarly, the right panel of Fig. 6 presents the dice evaluation for
3D brain MRI segmentations, as the percentage of ground truth data used in
training is progressively increased. In a manner mirroring the classification
results, our MGAug consistently outperforms the baseline algorithms by
yielding the highest dice scores. These results indicate that our augmentation
method has great potential to better alleviate the heavy need for labor and
time-consuming ground-truths.

Figure 6: Left to right: classification evaluation on 2D MNIST dataset;
segmentation evaluation on 3D brain MRIs with various percentages of ground
truth used in the training inference. The horizontal dashed line (\- -) serves
as a reference, representing the performance achieved when utilizing all
available training images without any augmentation.

#### Additional discussion.

It is to our attention that diffusion models have recently become a powerful
generative model (Ho, Jain, and Abbeel 2020). In this context, we initiate a
preliminary comparison between our MGAug model and data augmentation through a
diffusion model operating in the image space (by manipulating image
intensities and textures) (Trabucco et al. 2023). This comparison was carried
out with a specific focus on classification tasks using the 2D MNIST dataset.

Experimental results show that the diffusion model achieves an accuracy of
$97.56\%$ when using $50\%$ of the ground truth labels. This accuracy,
interestingly, is slightly lower than what our MGAug method accomplishes
($98.02\%$) when only using $20\%$ of the ground truth data. This may indicate
the advantages of multimodal distributions of latent geometric space developed
in our model. Our future work will further explore the power of diffusion
models in the latent geometric space for efficient data augmentation, which
has not been investigated so far.

MGAug vs. intensity-based augmentations. While the primary emphasis of this
paper does not fall into the realm of intensity-based augmentations, we
compare MGAug with a few notable DA schemes that manipulate image intensities
or textures. These methods include AutoAugment (Cubuk et al. 2019),
RandAugment (Cubuk et al. 2020), AugMix (Hendrycks et al. 2019), and
TrivialAugment (Müller and Hutter 2021).

Tab. 2 reports the accuracy comparison of a variety of DA schemes for image
classification networks based on the ViT backbone. The testing results on both
MNIST digits and 2D shape data show that our model consistently outperforms
the baselines.

Table 2: Accuracy comparison on both 2D digits and 2D shape datasets over a mix of augmentation methods (including intensity-based vs. transformation-based DA) on ViT. Datasets | 2D Digits | 2D Shape
---|---|---
Models | Acc | F1 | Acc | F1
AutoAugment | 90.82 | 90.73 | 76.05 | 75.25
RandAugment | 95.25 | 95.22 | 81.25 | 81.70
AugMix | 97.31 | 97.30 | 74.28 | 69.35
TrivialAug | 95.48 | 95.49 | 83.93 | 82.75
Affine | 95.90 | 95.84 | 78.57 | 76.53
AEA | 96.95 | 96.91 | 89.28 | 89.33
UDA | 97.85 | 97.83 | 91.07 | 96.84
MGAug | 98.62 | 98.78 | 98.21 | 98.20

## Conclusion & Discussion

In this paper, we present a novel data augmentation model, MGAug, which learns
augmenting transformations in a multimodal latent space of geometric
deformations reparameterized by low-dimensional tangent spaces of
diffeomorphisms. In contrast to existing augmentation models that adopt a
unimodal distribution to augment data, our augmentation method fully
characterizes the multimodal distribution of image transformations via
introducing a mixture of multivariate Gaussians as the prior. Our method
significantly improves the performance of multi-class classification (on 2D
dataset) and image segmentation (on 3D brain MRIs) and eliminates the
artifacts of geometric transformations by preserving topological structures of
augmented images. While MGAug focuses on LDDMM, the theoretical tool developed
in our model is widely applicable to various registration frameworks, for
example, stationary velocity fields (Modat et al. 2012).

To the best of our knowledge, we are the first to learn geometric deformations
in a multimodal manner via deep neural networks. Our augmentation model can be
easily applied to different domain-specific tasks, e.g., image reconstruction
and object tracking. Interesting future directions based on this work can be
i) investigating the model interpretability for classification tasks when
geometric deformations are learned in a multimodal manner, and ii) the
robustness study of MGAug when image intensity variations occur, i.e., under
adversarial noise attacks.

#### Acknowledgments.

This work was supported by NSF CAREER Grant 2239977.

## References

  * Araslanov and Roth (2021) Araslanov, N.; and Roth, S. 2021.  Self-supervised augmentation consistency for adapting semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 15384–15394.
  * Arnold (1966) Arnold, V. 1966.  Sur la géométrie différentielle des groupes de Lie de dimension infinie et ses applications à l’hydrodynamique des fluides parfaits.  In _Annales de l’institut Fourier_ , volume 16, 319–361.
  * Avants et al. (2008) Avants, B. B.; Epstein, C. L.; Grossman, M.; and Gee, J. C. 2008.  Symmetric diffeomorphic image registration with cross-correlation: evaluating automated labeling of elderly and neurodegenerative brain.  _Medical image analysis_ , 12(1): 26–41.
  * Beg et al. (2005) Beg, M. F.; Miller, M. I.; Trouvé, A.; and Younes, L. 2005.  Computing large deformation metric mappings via geodesic flows of diffeomorphisms.  _International journal of computer vision_ , 61(2): 139–157.
  * Chen et al. (2022) Chen, C.; Qin, C.; Ouyang, C.; Li, Z.; Wang, S.; Qiu, H.; Chen, L.; Tarroni, G.; Bai, W.; and Rueckert, D. 2022.  Enhancing MR image segmentation with realistic adversarial data augmentation.  _Medical Image Analysis_ , 82: 102597.
  * Chen et al. (2021) Chen, J.; Lu, Y.; Yu, Q.; Luo, X.; Adeli, E.; Wang, Y.; Lu, L.; Yuille, A. L.; and Zhou, Y. 2021.  Transunet: Transformers make strong encoders for medical image segmentation.  _arXiv preprint arXiv:2102.04306_.
  * Cubuk et al. (2019) Cubuk, E. D.; Zoph, B.; Mane, D.; Vasudevan, V.; and Le, Q. V. 2019.  Autoaugment: Learning augmentation strategies from data.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 113–123.
  * Cubuk et al. (2020) Cubuk, E. D.; Zoph, B.; Shlens, J.; and Le, Q. V. 2020.  Randaugment: Practical automated data augmentation with a reduced search space.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops_ , 702–703.
  * Dalca et al. (2019) Dalca, A. V.; Balakrishnan, G.; Guttag, J.; and Sabuncu, M. R. 2019.  Unsupervised learning of probabilistic diffeomorphic registration for images and surfaces.  _Medical image analysis_ , 57: 226–236.
  * Di Martino et al. (2014) Di Martino, A.; Yan, C.-G.; Li, Q.; Denio, E.; Castellanos, F. X.; Alaerts, K.; Anderson, J. S.; Assaf, M.; Bookheimer, S. Y.; Dapretto, M.; et al. 2014.  The autism brain imaging data exchange: towards a large-scale evaluation of the intrinsic brain architecture in autism.  _Molecular psychiatry_ , 19(6): 659–667.
  * Dice (1945) Dice, L. R. 1945.  Measures of the amount of ecologic association between species.  _Ecology_ , 26(3): 297–302.
  * Dilokthanakul et al. (2016) Dilokthanakul, N.; Mediano, P. A.; Garnelo, M.; Lee, M. C.; Salimbeni, H.; Arulkumaran, K.; and Shanahan, M. 2016.  Deep unsupervised clustering with gaussian mixture variational autoencoders.  _arXiv preprint arXiv:1611.02648_.
  * Dosovitskiy et al. (2020) Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.; et al. 2020.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_.
  * Gao et al. (2021) Gao, Y.; Tang, Z.; Zhou, M.; and Metaxas, D. 2021.  Enabling data diversity: efficient automatic augmentation via regularized adversarial training.  In _International Conference on Information Processing in Medical Imaging_ , 85–97. Springer.
  * Hauberg et al. (2016) Hauberg, S.; Freifeld, O.; Larsen, A. B. L.; Fisher, J.; and Hansen, L. 2016.  Dreaming more data: Class-dependent distributions over diffeomorphisms for learned data augmentation.  In _Artificial intelligence and statistics_ , 342–350. PMLR.
  * He et al. (2016) He, K.; Zhang, X.; Ren, S.; and Sun, J. 2016.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 770–778.
  * Hendrycks et al. (2019) Hendrycks, D.; Mu, N.; Cubuk, E. D.; Zoph, B.; Gilmer, J.; and Lakshminarayanan, B. 2019.  Augmix: A simple data processing method to improve robustness and uncertainty.  _arXiv preprint arXiv:1912.02781_.
  * Ho, Jain, and Abbeel (2020) Ho, J.; Jain, A.; and Abbeel, P. 2020.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33: 6840–6851.
  * Hossain et al. (2019) Hossain, T.; Shishir, F. S.; Ashraf, M.; Al Nasim, M. A.; and Shah, F. M. 2019.  Brain tumor detection using convolutional neural network.  In _2019 1st international conference on advances in science, engineering and robotics technology (ICASERT)_ , 1–6. IEEE.
  * Huang et al. (2017) Huang, G.; Liu, Z.; Van Der Maaten, L.; and Weinberger, K. Q. 2017.  Densely connected convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 4700–4708.
  * Iandola et al. (2016) Iandola, F. N.; Han, S.; Moskewicz, M. W.; Ashraf, K.; Dally, W. J.; and Keutzer, K. 2016.  SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and¡ 0.5 MB model size.  _arXiv preprint arXiv:1602.07360_.
  * Jack Jr et al. (2008) Jack Jr, C. R.; Bernstein, M. A.; Fox, N. C.; Thompson, P.; Alexander, G.; Harvey, D.; Borowski, B.; Britson, P. J.; L. Whitwell, J.; Ward, C.; et al. 2008.  The Alzheimer’s disease neuroimaging initiative (ADNI): MRI methods.  _Journal of Magnetic Resonance Imaging: An Official Journal of the International Society for Magnetic Resonance in Medicine_ , 27(4): 685–691.
  * Jaderberg et al. (2015) Jaderberg, M.; Simonyan, K.; Zisserman, A.; et al. 2015.  Spatial transformer networks.  _Advances in neural information processing systems_ , 28.
  * Jayakumar, Hossain, and Zhang (2023) Jayakumar, N.; Hossain, T.; and Zhang, M. 2023.  SADIR: Shape-Aware Diffusion Models for 3D Image Reconstruction.  In _International Workshop on Shape in Medical Imaging_ , 287–300. Springer.
  * Jiang et al. (2016) Jiang, Z.; Zheng, Y.; Tan, H.; Tang, B.; and Zhou, H. 2016.  Variational deep embedding: An unsupervised and generative approach to clustering.  _arXiv preprint arXiv:1611.05148_.
  * Joshi et al. (2004) Joshi, S.; Davis, B.; Jomier, M.; and Gerig, G. 2004.  Unbiased Diffeomorphic Atlas Construction for Computational Anatomy.  _NeuroImage_ , 23, Supplement1: 151–160.
  * Kim and Pavlovic (2016) Kim, J.; and Pavlovic, V. 2016.  A shape preserving approach for salient object detection using convolutional neural networks.  In _Pattern Recognition (ICPR), 2016 23rd International Conference on_ , 609–614. IEEE.
  * Kingma and Welling (2013) Kingma, D. P.; and Welling, M. 2013.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_.
  * Krebs et al. (2019) Krebs, J.; Delingette, H.; Mailhé, B.; Ayache, N.; and Mansi, T. 2019.  Learning a probabilistic model for diffeomorphic registration.  _IEEE transactions on medical imaging_ , 38(9): 2165–2176.
  * Krizhevsky, Sutskever, and Hinton (2017) Krizhevsky, A.; Sutskever, I.; and Hinton, G. E. 2017.  Imagenet classification with deep convolutional neural networks.  _Communications of the ACM_ , 60(6): 84–90.
  * LaMontagne et al. (2019) LaMontagne, P. J.; Benzinger, T. L.; Morris, J. C.; Keefe, S.; Hornbeck, R.; Xiong, C.; Grant, E.; Hassenstab, J.; Moulder, K.; Vlassenko, A. G.; et al. 2019.  OASIS-3: longitudinal neuroimaging, clinical, and cognitive dataset for normal aging and Alzheimer disease.  _MedRxiv_.
  * LeCun (1998) LeCun, Y. 1998.  The MNIST database of handwritten digits.  _http://yann.lecun.com/exdb/mnist/_.
  * Lim et al. (2019) Lim, S.; Kim, I.; Kim, T.; Kim, C.; and Kim, S. 2019.  Fast autoaugment.  _Advances in Neural Information Processing Systems_ , 32.
  * Miller, Trouvé, and Younes (2006) Miller, M. I.; Trouvé, A.; and Younes, L. 2006.  Geodesic shooting for computational anatomy.  _Journal of mathematical imaging and vision_ , 24(2): 209–228.
  * Mirza and Osindero (2014) Mirza, M.; and Osindero, S. 2014.  Conditional generative adversarial nets.  _arXiv preprint arXiv:1411.1784_.
  * Modat et al. (2012) Modat, M.; Daga, P.; Cardoso, M. J.; Ourselin, S.; Ridgway, G. R.; and Ashburner, J. 2012.  Parametric non-rigid registration using a stationary velocity field.  In _2012 IEEE Workshop on Mathematical Methods in Biomedical Image Analysis_ , 145–150. IEEE.
  * Moreno-Barea et al. (2018) Moreno-Barea, F. J.; Strazzera, F.; Jerez, J. M.; Urda, D.; and Franco, L. 2018.  Forward noise adjustment scheme for data augmentation.  In _2018 IEEE symposium series on computational intelligence (SSCI)_ , 728–734. IEEE.
  * Müller and Hutter (2021) Müller, S. G.; and Hutter, F. 2021.  Trivialaugment: Tuning-free yet state-of-the-art data augmentation.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 774–782.
  * Muralikrishnan et al. (2022) Muralikrishnan, S.; Chaudhuri, S.; Aigerman, N.; Kim, V. G.; Fisher, M.; and Mitra, N. J. 2022.  GLASS: Geometric latent augmentation for shape spaces.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 18552–18561.
  * Nocedal and Wright (1999) Nocedal, J.; and Wright, S. J. 1999.  _Numerical optimization_.  Springer.
  * Olut et al. (2020) Olut, S.; Shen, Z.; Xu, Z.; Gerber, S.; and Niethammer, M. 2020.  Adversarial data augmentation via deformation statistics.  In _European Conference on Computer Vision_ , 643–659. Springer.
  * Pang et al. (2019) Pang, J.; Chen, K.; Shi, J.; Feng, H.; Ouyang, W.; and Lin, D. 2019.  Libra r-cnn: Towards balanced learning for object detection.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 821–830.
  * Ronneberger, Fischer, and Brox (2015) Ronneberger, O.; Fischer, P.; and Brox, T. 2015.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18_ , 234–241. Springer.
  * Sebastian, Klein, and Kimia (2004) Sebastian, T. B.; Klein, P. N.; and Kimia, B. B. 2004.  Recognition of shapes by editing their shock graphs.  _IEEE Transactions on pattern analysis and machine intelligence_ , 26(5): 550–571.
  * Shattuck et al. (2008) Shattuck, D. W.; Mirza, M.; Adisetiyo, V.; Hojatkashani, C.; Salamon, G.; Narr, K. L.; Poldrack, R. A.; Bilder, R. M.; and Toga, A. W. 2008.  Construction of a 3D probabilistic atlas of human cortical structures.  _Neuroimage_ , 39(3): 1064–1080.
  * Shen et al. (2020) Shen, Z.; Xu, Z.; Olut, S.; and Niethammer, M. 2020.  Anatomical data augmentation via fluid-based image registration.  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , 318–328. Springer.
  * Simonyan and Zisserman (2014) Simonyan, K.; and Zisserman, A. 2014.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_.
  * Trabucco et al. (2023) Trabucco, B.; Doherty, K.; Gurinas, M.; and Salakhutdinov, R. 2023.  Effective data augmentation with diffusion models.  _arXiv preprint arXiv:2302.07944_.
  * Vialard et al. (2012) Vialard, F.-X.; Risser, L.; Rueckert, D.; and Cotter, C. J. 2012.  Diffeomorphic 3D image registration via geodesic shooting using an efficient adjoint calculation.  _International Journal of Computer Vision_ , 97(2): 229–241.
  * Wang and Zhang (2022) Wang, J.; and Zhang, M. 2022.  Geo-SIC: Learning Deformable Geometric Shapes in Deep Image Classifiers.  _arXiv preprint arXiv:2210.13704_.
  * Wells III et al. (1996) Wells III, W. M.; Viola, P.; Atsumi, H.; Nakajima, S.; and Kikinis, R. 1996.  Multi-modal volume registration by maximization of mutual information.  _Medical image analysis_ , 1(1): 35–51.
  * Yun et al. (2019) Yun, S.; Han, D.; Oh, S. J.; Chun, S.; Choe, J.; and Yoo, Y. 2019.  Cutmix: Regularization strategy to train strong classifiers with localizable features.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 6023–6032.
  * Zhang and Fletcher (2013) Zhang, M.; and Fletcher, T. 2013.  Probabilistic principal geodesic analysis.  _Advances in neural information processing systems_ , 26.
  * Zhang, Shao, and Fletcher (2015) Zhang, M.; Shao, H.; and Fletcher, P. T. 2015.  A mixture model for automatic diffeomorphic multi-atlas building.  In _Med. Image Comput. Comput.-Assisted Intervention Workshop-BAMBI_.
  * Zhang, Singh, and Fletcher (2013) Zhang, M.; Singh, N.; and Fletcher, P. T. 2013.  Bayesian estimation of regularization and atlas building in diffeomorphic image registration.  In _International conference on information processing in medical imaging_ , 37–48. Springer.
  * Zhao et al. (2019) Zhao, A.; Balakrishnan, G.; Durand, F.; Guttag, J. V.; and Dalca, A. V. 2019.  Data augmentation using learned transformations for one-shot medical image segmentation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 8543–8553.
  * Zhou et al. (2018) Zhou, Z.; Rahman Siddiquee, M. M.; Tajbakhsh, N.; and Liang, J. 2018.  Unet++: A nested u-net architecture for medical image segmentation.  In _Deep learning in medical image analysis and multimodal learning for clinical decision support_ , 3–11. Springer.
  * Zoph et al. (2020) Zoph, B.; Cubuk, E. D.; Ghiasi, G.; Lin, T.-Y.; Shlens, J.; and Le, Q. V. 2020.  Learning data augmentation strategies for object detection.  In _European conference on computer vision_ , 566–583. Springer.