:
1 1 2022 Machine Learning for Health (ML4H) 2022

# Clinically Plausible Pathology-Anatomy Disentanglement in Patient Brain MRI
with Structured Variational Priors

Anjun Hu1,2 Jean-Pierre R. Falet1,2,3 Brennan S. Nichyporuk1,2 Changjian
Shui1,2 Douglas L. Arnold3 Sotirios A. Tsaftaris4,5 Tal Arbel1,2
{anjun, jpfalet, maxshui<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
1Centre for Intelligent Machines  McGill University  Montréal  Canada
2MILA (Québec Artificial Intelligence Institute)  Montréal  Canada
3Brain Imaging Center  Dept. Neurology and Neurosurgery  McGill University
Montréal  Canada
4School of Engineering  University of Edinburgh  UK
5The Alan Turing Institute  UK

###### Abstract

We propose a hierarchically structured variational inference model for
accurately disentangling observable evidence of disease (e.g. brain lesions or
atrophy) from subject-specific anatomy in brain MRIs. With flexible, partially
autoregressive priors, our model (1) addresses the subtle and fine-grained
dependencies that typically exist between anatomical and pathological
generating factors of an MRI to ensure the clinical validity of generated
samples; (2) preserves and disentangles finer pathological details pertaining
to a patient’s disease state. Additionally, we experiment with an alternative
training configuration where we provide supervision to a subset of latent
units. It is shown that (1) a partially supervised latent space achieves a
higher degree of disentanglement between evidence of disease and subject-
specific anatomy; (2) when the prior is formulated with an autoregressive
structure, knowledge from the supervision can propagate to the unsupervised
latent units, resulting in more informative latent representations capable of
modelling anatomy-pathology interdependencies.

###### keywords:

Variational Inference, Disentanglement, Medical Image Analysis

##  1 Introduction

In medical image analysis, various representation learning methods have been
introduced to disentangle particular anatomical or pathological structures
from the the rest of the image under different observations (Liu et al.,
2021a; Fragemann et al., 2022). Although disentanglement of observable
evidence of disease (e.g. brain lesions or atrophy) from subject-specific
anatomy has been shown to be helpful in various downstream tasks (Qin et al.,
2019; Hu et al., 2020; Zhou et al., 2021; Zuo et al., 2021; Liu et al., 2021b,
c; Yang et al., 2022), maintaining the clinical plausibility of the learned
representations is particularly challenging in the context of brain MRI
analysis, for a number of reasons. Firstly, pathological and anatomical
features in brain MRIs typically exhibits some degree of dependency, but the
exact structure is generally unknown a priori or requires extensive domain
knowledge to accurately describe. For example, multiple sclerosis (MS) is a
chronic neurological disease characterized by T2 hyperintense lesions and
gadolinium-enhancing (Gd+) lesions in the brain and spinal cord. MS lesions
are typically found in periventricular white matter but cannot occur in
certain brain regions (e.g. within the ventricles). It is often the case that
the observable pathological features in the brain (e.g. hyper-intense MS
lesions) cannot be easily disentangled from subject-specific anatomical
structures (e.g. sulcal pattern, ventricular shape) in a clinically plausible
way due to dependencies between the anatomical and pathological generating
factors. Secondly, fine-grained pathological features in brain MRIs may have
very high clinical significance. Although the omission of finer details does
not preclude generative models from achieving a good approximation of the
overall data distribution (i.e. a good test set log likelihood), such details
could be highly meaningful in certain downstream tasks. This is, again,
exemplified by MS lesions that can be as small as 3 mm but still represent a
significant marker of disease activity (Wang et al., 1997). A robust and
clinically useful representation in the context of brain MRI analysis must
therefore: 1) faithfully model the spatial distribution of the lesions and
their dependency on patients’ brain anatomical structures, and 2) accurately
capture and disentangle lesions of all sizes, along with other potential
imaging markers characterized by fine-grained details in the images (e.g.
white matter texture).

Many existing methods for learning disentangled representations are based on
the variational auto-encoder (VAE) (Kingma and Welling, 2014; Burgess et al.,
2018; Higgins et al., 2016). However, a straight adoption of VAE for
pathology-anatomy disentanglement in patient brain MRIs is often
unsatisfactory due to the aforementioned challenges. There are two common
failure modes. Firstly, mean-field variational inference poses the unlikely
assumption that all generating factors are independent (Zhang et al., 2019),
thereby failing to capture the inherent dependencies that exist between the
anatomical and pathological generative factors in patient brain images. This
may result in the synthesis of clinically implausible samples, as depicted in
fig:IntroInvalidMSLesions, where lesions appear in clinically impossible
regions (red arrow). Moreover, the independence assumption may exacerbates
another pitfall of VAEs (Zhang et al., 2019), namely, the tendency to suffer
from latent mode-covering or over-generalisation (Hu et al., 2018). This may
lead to lower synthesis quality or an inability to preserve crucial finer
details in medical images. fig:IntroBlurredMSLesions depicts such behaviour in
a mean-field model where small lesions are obscured in the reconstructed
images. Failure to capture fine details in the learned representation could
lead to significantly poorer performance in downstream tasks.

fig:IntroFailureModes [][b] [angle=0]InvalidLesions [][b]
[angle=0]BlurredLesionsWithArrows

Figure 1: Failure modes of mean-field VAE. (a) A sample drawn from a mean-
field model with lesions in clinically invalid locations (red arrow). (b)
Mean-field model leads to missed small lesions in the reconstructed image
(yellow arrows). Our proposed structured model does not suffer from these
issues

In this work, we propose to address these issues by using structured
variational inference (Hoffman and Blei, 2015) for fine-grained pathology-
anatomy disentanglement in brain MRI. Specifically we model the dependencies
that typically exist between pathological and anatomical features via multi-
scale VAEs with a hierarchical latent structure (fig:GraphicalModel). We
evaluate the effect of different priors in the reconstruction quality and the
degree of disentanglement both quantitatively and qualitatively. We find that
more expressive structured priors indeed lead to higher reconstruction quality
and the preservation of important small pathological details. We verify that
the model is capable of pathology disentanglement in an unsupervised setting.
With an optional supervision objective, the model is shown to achieve a higher
degree of disentanglement and to be capable of capturing latent dependency.

##  2 Model

Our model accepts image space observations sampled from a dataset
$\mathbf{x}\in\mathbb{R}^{w_{\mathbf{x}}\times h_{\mathbf{x}}}\sim\mathcal{D}$
as inputs. The model adheres to the customary setup of VAEs with a structured
latent space consisting of $(L+1)$ disjoint variable groups (layers) that
follow a hierarchical structure, as portrayed in fig:GraphicalModel. Each
group or layer consists of spatial latent variables (You et al., 2019) at
various resolutions scales, denoted as
$\mathbf{z}=\\{\mathbf{z}_{l}\in\mathbb{R}^{w_{l}\times h_{l}\times
c_{l}};l\in\\{0,1,...,L\\}\\}$. The inference and generative models can
therefore be expressed as follows:

$\begin{gathered}q_{\phi,\theta}(\mathbf{z}|\mathbf{x}):=q_{\phi_{L},\theta_{L}}(\mathbf{z}_{L}|\mathbf{x})\prod_{l=0}^{L-1}q_{\phi_{l},\theta_{l}}(\mathbf{z}_{l}|\mathbf{z}_{l+1})\\\
p_{\theta}(\mathbf{x},\mathbf{z}):=p_{\theta_{x}}(\mathbf{x}|\mathbf{z}_{L})p(\mathbf{z}_{0})\prod_{l=1}^{L}p_{\theta_{l}}(\mathbf{z}_{l}|\mathbf{z}_{l-1})\end{gathered}$
(1)

fig:GraphicalModel

Figure 2: The hierarchically structured latent space with $(L+1)$ disjoint
variable groups (layers), encoder distribution parameters
($\Delta\mu_{l},\Delta\sigma_{l}$) and decoder distribution parameters
($\mu_{l},\sigma_{l}$). Dashed lines are active during inference. Solid lines
are active during generation and, if residual parameterisation is used, also
during bidirectional inference.

fig:ArchDiagram

Figure 3: Network Architecture

tab:ModelParametrizations  Model ELBO VAE (M.1)
$\begin{array}[]{r@{}l@{}}\mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\Bigl{[}\mathbb{E}_{\mathbf{z}\sim
q_{\phi}(\mathbf{z}|\mathbf{x})}[\log
p_{\theta}(\mathbf{x}|\mathbf{z})]-\sum_{l=0}^{L}\mathrm{KLD}[q_{\phi}(\mathbf{z}_{l}|\mathbf{x})||\mathcal{N}(\mathbf{0},\mathbf{I})]\Bigr{]}\end{array}$
nVAE (M.2)
$\begin{array}[]{r@{}l@{}}\mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\Bigl{[}\mathbb{E}_{\mathbf{z}\sim
q_{\phi,\theta}(\mathbf{z}|\mathbf{x})}[\log
p_{\theta}(\mathbf{x}|\mathbf{z})]-\mathrm{KLD}[q_{\phi}(\mathbf{z}_{0}|\mathbf{x})||\mathcal{N}(\mathbf{0},\mathbf{I})]\\\
-\sum_{l=1}^{L}\mathbb{E}_{q_{\phi,\theta}(\mathbf{z}_{<l}|\mathbf{x})}\mathrm{KLD}[q_{\phi,\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l},\mathbf{x})||p_{\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l})]\Bigr{]}\end{array}$
nVMP (M.3)
$\begin{array}[]{r@{}l@{}}\mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\Biggl{[}\mathbb{E}_{\mathbf{z}\sim
q_{\phi,\theta}(\mathbf{z}|\mathbf{x})}[\log
p_{\theta}(\mathbf{x}|\mathbf{z})]-\mathrm{KLD}\bigl{[}q_{\phi}(\mathbf{z}_{0}|\mathbf{x})||\frac{1}{K}\sum_{k=1}^{K}q_{\phi}(\mathbf{z}_{0}|\mathbf{u}_{k})\bigr{]}\\\
-\sum_{l=1}^{L}\mathbb{E}_{q_{\phi,\theta}(\mathbf{z}_{<l}|\mathbf{x})}\mathrm{KLD}[q_{\phi,\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l},\mathbf{x})||p_{\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l})]\Biggr{]}\end{array}$
nVMP+ (M.4)
$\begin{array}[]{r@{}l@{}}\mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\Biggl{[}\mathbb{E}_{\mathbf{z}\sim
q_{\phi,\theta}(\mathbf{z}|\mathbf{x})}[\log
p_{\theta}(\mathbf{x}|\mathbf{z})]-\mathrm{KLD}\bigl{[}q_{\phi}(\mathbf{z}_{0}|\mathbf{x})||\frac{1}{K}\sum_{k=1}^{K}q_{\phi}(\mathbf{z}_{0}|\mathbf{u}_{k})\bigr{]}\\\
-\sum_{l=1}^{L}\mathbb{E}_{q_{\phi,\theta}(\mathbf{z}_{<l}|\mathbf{x})}\biggl{[}\mathrm{KLD}[q_{\phi,\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l},\mathbf{x})||p_{\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l})]+\mathrm{KLD}[p_{\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l})||\frac{1}{K}\sum_{k=1}^{K}q_{\phi}(\mathbf{z}_{l}|\mathbf{u}_{k})]\biggr{]}\Biggr{]}\end{array}$

Table 1: ELBOs for all four model parameterisation options examined in this
work.

In this paper, we examine four ways to construct the ELBO objective,
summerized in tab:ModelParametrizations.

(1) VAE (1): a vanilla multi-scale VAE with mean-field, parameter-free
standard Gaussian priors at each layer. This model has a “hierarchy” in the
sense of having latent representations at various resolution scales, but is
not “hierarchical” in its distributional parametrization as there is no
explicit inter-group dependency nor explicit information sharing between the
encoder and the decoder. In this parametrization:

$\displaystyle\begin{split}&q_{\phi_{l}}(\mathbf{z}_{l}):=\mathcal{N}(\Delta\mu_{\phi_{>l}}(\mathbf{x}),\Delta\sigma_{\phi_{>l}}(\mathbf{x}))\end{split}$
(2a) $\displaystyle
p_{\theta_{l}}(\mathbf{z}_{l}):=\mathcal{N}(\mathbf{0},\mathbf{I})\;\;\forall
l\in\\{0,1,...,L\\}$ (2b)

(2) nVAE (1): a hierarchical model with residual normal parameterisation
proposed by Vahdat and Kautz (2020) and Rasmus et al. (2015). The features
that set this model apart from its vanilla counterpart are the explicit
information sharing between the encoder and the decoder networks, as well as
its partially auto-regressive nature. Firstly, unlike conventional VAEs,
decoder parameters $\theta$ in nVAE not only characterize the generative
distribution $p_{\theta}$ (3b) but are also a part of the inference model and
hence play an important role in characterizing the posterior distribution
$q_{\phi,\theta}$ (3a). For latent groups other than the topmost one $l>0$,
the inference model is bidirectional. It estimates the relative variational
posteriors (3a) that characterize the deviation from priors obtained from
preceeding layers of the decoder. With this design, KL optimization is
expected to be simpler than when posteriors predict the absolute mean and
variances at each layer.

$q_{\phi,\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l},\mathbf{x}):=\mathcal{N}\Bigl{(}\mu_{\theta_{<l}}(\mathbf{z}_{<l})+\Delta\mu_{\phi_{>l}}(\mathbf{x}),\\\\[-5.0pt]
\sigma_{\theta_{<l}}(\mathbf{z}_{<l})\cdot\Delta\sigma_{\phi_{>l}}(\mathbf{x})\Bigr{)}$
(3a) $\displaystyle
p_{\theta}(\mathbf{z}_{l}|\mathbf{z}_{<l}):=\mathcal{N}(\mu_{\theta_{<l}}(\mathbf{z}_{<l}),\sigma_{\theta_{<l}}(\mathbf{z}_{<l}))$
(3b)

Furthermore, nVAE is considered to be partially auto-regressive and hence a
more expressive prior than the standard mean-field parametrization. While the
prior for each group $p(\mathbf{z}_{l})$ is dependent on those of the
preceding layers $p(\mathbf{z}_{<l})$, each element within the same latent
group $\mathbf{z}_{l}:=\\{z^{i}_{l},...\\}$ still adhere to the independence
assumption as $\forall
l\in\\{0,1,...,L\\},z^{i}_{l}\rotatebox[origin={c}]{90.0}{$\models$}z^{j}_{l};i\neq
j$. We can hence calculate the relative KLD loss for each element with a
simple analytic expression:

$\mathrm{(Relative)\;KLD}\;[q_{\phi,\theta}(z^{i}_{l}|\mathbf{x})||p_{\theta}(z^{i}_{l})]=\\\
\frac{1}{2}\Biggl{(}\frac{(\Delta\mu_{l}^{i})^{2}}{(\sigma_{l}^{i})^{2}}+(\Delta\sigma_{l}^{i})^{2}-\log(\sigma_{l}^{i})^{2}-1\Biggr{)}$
(4)

We propose two extensions to nVAE by integrating a VamPrior (Tomczak and
Welling, 2017) into the hierarchical VAE setup for extra flexibility and
expressiveness in the hierarchical latent structure. By incorporating encoder
parameters in the trainable prior, we expect the following two models to
achieve a greater extent of “coupling” or “collaboration” between the priors
and the posteriors.

(3) nVMP (1) replaces the standard-Gaussian prior of the topmost layer of nVAE
with a $K$-component multimodal VamPrior
$\frac{1}{K}\sum_{k=1}^{K}q_{\phi}(\mathbf{z}_{l}|\mathbf{u}_{k})$
characterized by trainable pseudo-inputs $\mathbf{u}$ and encoder parameters
$\phi$. The subsequent layers still adhere to the hierarchical residual
parameterisation of nVAE (we retain the residual Gaussian parameterisation for
subsequent layers). The implication is that only $\mathbf{z}_{0}$ is multi-
modal whereas the distributions of all lower level “residual deviations” are
assumed to be Gaussian.

(4) nVMP+ (1) extends nVMP with one more KL term between the encoder-driven
VamPriors and the decoder priors is imposed for the entire hierarchy. In this
case, the decoder “priors” are regarded as “intermediate posteriors” and
encouraged to imitate the encoder-driven multi-modal distribution throughout
the hierarchy. We postulate that this configuration adds an extra layer of
information sharing between the encoder and the decoder networks which can
potentially lead to further improvement in representation quality.

##  3 Experiments and Results

We validate our approach on two brain MRI datasets: the publically available
Alzheimer’s Disease Neuroimaging Initiative (ADNI) dataset (Mueller et al.,
2005) ($N=864$), and a proprietary MS dataset from a MS clinical trial
($N=815$). The central 16 2-D slices of T1-weighted sequences were used for
the AD experiments, while the central 24 2-D slices of Fluid Attenuated
Inverse Recovery (FLAIR) sequence were used for the MS experiments. Expert T2
lesion segmentation labels for the MS experiments were provided. Both datasets
were divided into non-overlapping training (60 %), validation (20 %) and
testing (20 %) sets. Additional acquisition and pre-processing details are
described in Appendix A.

We first train the model under an unsurpervised setting and evaluate the
effect of incorporating additional prior structures on synthesis quality. As
shown in Table LABEL:tab:MainResults, VAEs with more expressive structured
priors indeed outperform their mean-field counterpart at the same model
capacity in terms of image reconstruction fidelity.

tab:MainResults  [MS]  Model LL$\uparrow$ PSNR$\uparrow$ SSIM$\uparrow$
FID$\downarrow$ VAE 2758 25.1 0.72 0.058 nVAE 2458 25.7 0.75 0.023 nVMP 2374
25.9 0.75 0.031 nVMP+ 1953 26.6 0.79 0.035 [AD]  Model LL$\uparrow$
PSNR$\uparrow$ SSIM$\uparrow$ FID$\downarrow$ VAE 2386 24.6 0.70 0.030 nVAE
2105 25.2 0.73 0.013 nVMP 1863 25.5 0.75 0.011 nVMP+ 842 26.8 0.80 0.007

Table 2: Reconstruction quality metrics. (Wang et al., 2004; Heusel et al.,
2017)

We additionally examine model behaviours in a supervised learning setting
depicted in the bottom-left (purple) block in fig:GraphicalModel and
fig:ArchDiagram. In this setting, we supplement the MS model with a lesion
segmentation objective between a chosen “pathological” latent subset
$\mathbf{z}_{P}$ and expert pathology (lesion) segmentation labels
$\mathbf{P}$. The rest of the latent space (that remains unsupervised) are
regarded as the anatomical latent subsets, denoted as $\mathbf{z}_{A}$.

Firstly, as one might expect, supervision is shown to enhance latent
disentanglement as one may anticipate. Disease-related features in the
synthesized images are noticeably more sensitive (Rolínek et al., 2018) to
perturbations in $\mathbf{z}_{P}$ comapred to $\mathbf{z}_{A}$, as shown in
fig:AttributeSensitivity, Appendix B. Such a disparity in attribute
sensitivity is appreciable in unsupervised models, but is made much more
pronounced by the selective latent supervision.

Secondly and more importantly, _supervision helps to verify that the model is
indeed actively using the latent structures._ In models with autoregressive
structures (nVAE, nVMP, nVMP+), knowledge from the supervision is propagated
to the unsupervised “anatomical” latent units $\mathbf{z}_{A}$, as in, those
unsupervised latent units attain a higher linear predictability (Lasso
regression $R^{2}$ scores (Eastwood and Williams, 2018), Table
LABEL:tab:Informativeness) with respect to lesion volume. This is in contrast
to the behaviour of the baseline mean-field VAE, where information from the
supervision task is constrained within the supervised group $\mathbf{z}_{P}$.
This observation shows that the model is indeed taking advantage of the extra
structures brought by the autoregressive priors and the residual
parameterisation and hence, indeed capable of modelling the dependencies
between anatomical and pathological generating factors.

fig:StyleMixing [angle=0]StyleMixing

Figure 4: Pathology-anatomy disentanglement visualised by “style-mixing”
between pairs of Alzheimer’s Disease (AD) and Multiple Sclerosis (MS) patient
brain MRIs.

We can qualitatively evaluate pathology-anatomy disentanglement by swapping
anatomical and pathological latent features between a pair of subjects in a
manner similar to “style-mixing” (Karras et al., 2018). As shown in
fig:StyleMixing for representative examples, brain atrophy in AD patients
(left), and T2 lesions in MS patients (right), are disentangled from the
subject’s anatomical particularities (such as sulcal pattern), thus enabling
the mixing the pathology of one patient with the anatomy of the other.

We may also leverage conditional distributions learned by the model to examine
subject-specific pathology distributions. For example, based on learned
representations of Subject B1 in fig:StyleMixing, we may visualise many
possible disease states given this subject’s anatomy
(fig:ConditionalResampling, top row) or explore how this subject’s lesions
would manifest on other subjects’ brain anatomies (fig:ConditionalResampling,
bottom row).

fig:ConditionalResampling

Figure 5: Visualising conditional distributions. Images in the top row are
generated by fixing $\mathbf{z}_{A}$ to that of Subject B1 and resampling the
layer corresponding to $\mathbf{z}_{P}$. Images in the bottom row are
generated by obtaining $\mathbf{z}_{A}$ from other real samples and fixing
$\mathbf{z}_{P}$ to that of Subject B1.

##  4 Conclusions

We propose hierarchical VAEs with structured priors for learning pathology-
anatomy disentangled representations of brain MRIs. Our model can faithfully
capture imaging features, including fine-grained details, while accounting for
pathology-anatomy dependencies to ensure sample validity. We additionally
examine model bevaviours in a supervised learning setting. Supervision is
shown to (1) further enhance latent disentanglement; and (2) enable the
inspection of information propagation between latent groups for modelling
pathology-anatomy interdependencies. Our model allows for robust and
controllable brain MRI synthesis rich in high-frequency and pathologically-
sound details, which could be meaningful for various downstream tasks.

The authors are grateful to the International Progressive MS Alliance for
supporting this work (grant number: PA-1412-02420), and to the companies who
generously provided the clinical trial data that made it possible: Biogen,
BioMS, MedDay, Novartis, Roche / Genentech, and Teva. Funding was also
provided by the Natural Sciences and Engineering Research Council of Canada,
the Canadian Institute for Advanced Research (CIFAR) Artificial Intelligence
Chairs program, and a technology transfer grant from Mila - Quebec AI
Institute. S.A. Tsaftaris acknowledges the support of Canon Medical and the
Royal Academy of Engineering and the Research Chairs and Senior Research
Fellowships scheme (grant RCSRF1819 / 8 / 25). Supplementary computational
resources and technical support were provided by Calcul Québec and the Digital
Research Alliance of Canada. Falet, J.-P., was supported by an end MS
Personnel Award from the Multiple Sclerosis Society of Canada, by a Canada
Graduate Scholarship-Masters Award from the Canadian Institutes of Health
Research, and by the Fonds de recherche du Québec - Santé / Ministère de la
Santé et des Services sociaux training program for specialty medicine
residents with an interest in pursuing a research career, Phase 1. This work
was made possible by the end-to-end deep learning experimental pipeline
developed in collaboration with our colleagues Justin Szeto, Eric Zimmerman,
and Kirill Vasilevski. Additionally, the authors would like to thank Louis
Collins and Mahsa Dadar for preprocessing the MRI data.

## References

  * Burgess et al. (2018) Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner.  Understanding disentangling in $\beta$-vae.  _arXiv preprint arXiv:1804.03599_ , 2018.
  * Chen et al. (2016) Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel.  Variational lossy autoencoder.  _arXiv preprint arXiv:1611.02731_ , 2016.
  * Eastwood and Williams (2018) Cian Eastwood and Christopher KI Williams.  A framework for the quantitative evaluation of disentangled representations.  In _International Conference on Learning Representations_ , 2018.
  * Fragemann et al. (2022) Jana Fragemann, Lynton Ardizzone, Jan Egger, and Jens Kleesiek.  Review of disentanglement approaches for medical applications–towards solving the gordian knot of generative models in healthcare.  _arXiv preprint arXiv:2203.11132_ , 2022.
  * Fu et al. (2019) Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin.  Cyclical annealing schedule: A simple approach to mitigating kl vanishing.  In _Proceedings of NAACL-HLT_ , pages 240–250, 2019.
  * Heusel et al. (2017) Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, Günter Klambauer, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a nash equilibrium.  _CoRR_ , abs/1706.08500, 2017.  URL http://arxiv.org/abs/1706.08500.
  * Higgins et al. (2016) Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.  $beta$-vae: Learning basic visual concepts with a constrained variational framework.  2016\.
  * Hoffman and Blei (2015) Matthew Hoffman and David Blei.  Stochastic Structured Variational Inference.  In Guy Lebanon and S. V. N. Vishwanathan, editors, _Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics_ , volume 38 of _Proceedings of Machine Learning Research_ , pages 361–369, San Diego, California, USA, 09–12 May 2015. PMLR.  URL https://proceedings.mlr.press/v38/hoffman15.html.
  * Hu et al. (2020) Dan Hu, Han Zhang, Zhengwang Wu, Fan Wang, Li Wang, J. Keith Smith, Weili Lin, Gang Li, and Dinggang Shen.  Disentangled-multimodal adversarial autoencoder: Application to infant age prediction with incomplete multimodal neuroimages.  _IEEE Transactions on Medical Imaging_ , 39(12):4137–4149, 2020.  10.1109/TMI.2020.3013825.
  * Hu et al. (2018) Zhiting Hu, Zichao Yang, Ruslan Salakhutdinov, and Eric P. Xing.  On unifying deep generative models.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=rylSzl-R-.
  * Karras et al. (2018) Tero Karras, Samuli Laine, and Timo Aila.  A style-based generator architecture for generative adversarial networks, 2018.  URL https://arxiv.org/abs/1812.04948.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma and Welling (2014) Diederik P. Kingma and Max Welling.  Auto-Encoding Variational Bayes.  In _2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings_ , 2014.
  * Liu et al. (2021a) Xiao Liu, Pedro Sanchez, Spyridon Thermos, Alison Q O’Neil, and Sotirios A Tsaftaris.  A tutorial on learning disentangled representations in the imaging domain.  _arXiv preprint arXiv:2108.12043_ , 2021a.
  * Liu et al. (2021b) Xiao Liu, Spyridon Thermos, Alison O’Neil, and Sotirios A Tsaftaris.  Semi-supervised meta-learning with disentanglement for domain-generalised medical image segmentation.  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , pages 307–317. Springer, 2021b.
  * Liu et al. (2021c) Xiaofeng Liu, Fangxu Xing, Georges El Fakhri, and Jonghye Woo.  A unified conditional disentanglement framework for multimodal brain mr image translation.  In _2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI)_ , pages 10–14. IEEE, 2021c.
  * Mueller et al. (2005) Susanne G Mueller, Michael W Weiner, Leon J Thal, Ronald C Petersen, Clifford R Jack, William Jagust, John Q Trojanowski, Arthur W Toga, and Laurel Beckett.  Ways toward an early diagnosis in alzheimer’s disease: the alzheimer’s disease neuroimaging initiative (adni).  _Alzheimer’s & Dementia_, 1(1):55–66, 2005\.
  * Qin et al. (2019) Chen Qin, Bibo Shi, Rui Liao, Tommaso Mansi, Daniel Rueckert, and Ali Kamen.  Unsupervised deformable registration for multi-modal images via disentangled representations.  In _International Conference on Information Processing in Medical Imaging_ , pages 249–261. Springer, 2019.
  * Rasmus et al. (2015) Antti Rasmus, Harri Valpola, Mikko Honkala, Mathias Berglund, and Tapani Raiko.  Semi-supervised learning with ladder network.  _CoRR_ , abs/1507.02672, 2015.  URL http://arxiv.org/abs/1507.02672.
  * Rolínek et al. (2018) Michal Rolínek, Dominik Zietlow, and Georg Martius.  Variational autoencoders pursue PCA directions (by accident).  _CoRR_ , abs/1812.06775, 2018.  URL http://arxiv.org/abs/1812.06775.
  * Tomczak and Welling (2017) Jakub M. Tomczak and Max Welling.  VAE with a vampprior.  _CoRR_ , abs/1705.07120, 2017.  URL http://arxiv.org/abs/1705.07120.
  * Vahdat and Kautz (2020) Arash Vahdat and Jan Kautz.  Nvae: A deep hierarchical variational autoencoder.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 19667–19679. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/e3b21256183cf7c2c7a66be163579d37-Paper.pdf.
  * Vahdat et al. (2018) Arash Vahdat, William Macready, Zhengbing Bian, Amir Khoshaman, and Evgeny Andriyash.  Dvae++: Discrete variational autoencoders with overlapping transformations.  In _International conference on machine learning_ , pages 5035–5044. PMLR, 2018.
  * Wang et al. (1997) Liqun Wang, H-M Lai, A J Thompson, and D H Miller.  Survey of the distribution of lesion size in multiple sclerosis: implication for the measurement of total lesion load.  _Journal of Neurology, Neurosurgery & Psychiatry_, 63(4):452–455, 1997.  ISSN 0022-3050.  10.1136/jnnp.63.4.452.  URL https://jnnp.bmj.com/content/63/4/452.
  * Wang et al. (2004) Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli.  Image quality assessment: from error visibility to structural similarity.  _IEEE transactions on image processing_ , 13(4):600–612, 2004.
  * Yang et al. (2022) Qiushi Yang, Xiaoqing Guo, Zhen Chen, Peter Y. M. Woo, and Yixuan Yuan.  D2-net: Dual disentanglement network for brain tumor segmentation with missing modalities.  _IEEE Transactions on Medical Imaging_ , pages 1–1, 2022.  10.1109/TMI.2022.3175478.
  * You et al. (2019) Suhang You, Kerem C. Tezcan, Xiaoran Chen, and Ender Konukoglu.  Unsupervised lesion detection via image restoration with a normative prior.  In M. Jorge Cardoso, Aasa Feragen, Ben Glocker, Ender Konukoglu, Ipek Oguz, Gozde Unal, and Tom Vercauteren, editors, _Proceedings of The 2nd International Conference on Medical Imaging with Deep Learning_ , volume 102 of _Proceedings of Machine Learning Research_ , pages 540–556. PMLR, 08–10 Jul 2019.  URL https://proceedings.mlr.press/v102/you19a.html.
  * Zhang et al. (2019) C. Zhang, J. Butepage, H. Kjellstrom, and S. Mandt.  Advances in variational inference.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 41(08):2008–2026, aug 2019.  ISSN 1939-3539.  10.1109/TPAMI.2018.2889774.
  * Zhou et al. (2021) Lei Zhou, Joseph Bae, Huidong Liu, Gagandeep Singh, Jeremy Green, Dimitris Samaras, and Prateek Prasanna.  Chest radiograph disentanglement for covid-19 outcome prediction.  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , pages 345–355. Springer, 2021.
  * Zuo et al. (2021) Lianrui Zuo, Blake E Dewey, Aaron Carass, Yihao Liu, Yufan He, Peter A Calabresi, and Jerry L Prince.  Information-based disentangled representation learning for unsupervised mr harmonization.  In _International Conference on Information Processing in Medical Imaging_ , pages 346–359. Springer, 2021.

##  Appendix A Data Acquisition, Implementation and Training Details

All MRI sequences were acquired at a resolution of 1 mm $\times$ 1 mm $\times$
1 mm, Each 2-D slice was downsampled to a resolution of 2 mm $\times$ 2 mm.
These were standardized to have zero-mean and unit variance.

We compare the four parameterisations in Table tab:ModelParametrizations with
a 5-layer model ($L=4$) the exact same capacity. For each dataset, the latent
space capacity is set to $\\{\mathbf{z}_{L}\in\mathbb{R}^{w_{\mathbf{x}}\times
h_{\mathbf{x}}\times
2},\mathbf{z}_{L-1}\in\mathbb{R}^{(w_{\mathbf{x}}/2)\times(h_{\mathbf{x}}/2)\times
2},...,\mathbf{z}_{0}\in\mathbb{R}^{(w_{\mathbf{x}}/2^{L})\times(h_{\mathbf{x}}/2^{L})\times
2}\\}$. We use the Adam optimizer (Kingma and Ba, 2014) with a learning rate
of 5e-5 and s weight decay of 1e-8.

Two loss re-weighting mechanisms are used in our training procedure: (1) We
use a linear annealing schedule (Fu et al., 2019) for KLD losses with a cycle
length of 10000 iterations. The initial KLD learning rate is set to 2e-7. (2)
To avoid posterior collapse, we use a KL Balancing trick suggested by (Vahdat
and Kautz, 2020). We re-scale each KL term of the hierarchy with a coefficient
proportional to the size of each latent layer as well as the KLD value of that
layer. This mechanism encourages more balanced information attribution to each
latent layer (Vahdat et al., 2018; Chen et al., 2016).

##  Appendix B Additional Results

As discussed in Section 3, we evaluate layer-wise latent pathology
informativeness of MS models by examining each layer’s linear predictability
of a salient pathological attribute, T2 lesion volume. To quantify linear
predictability, we train Lasso regressors ($\alpha=10$) with latent
representations obtained from each individual latent layer of each model and
compute each Lasso regressors’s $R^{2}$ scores with respect to T2 lesion
volume based on expert segmentation labels.

tab:Informativeness  Model $\mathbf{z}_{0}$ $\mathbf{z}_{1}$ $\mathbf{z}_{P}$
$\mathbf{z}_{3}$ $\mathbf{z}_{4}$ VAE 0.20 0.08 0.28 0.01 0.00 nVAE 0.11 0.12
0.00 0.00 0.00 nVMP 0.06 0.23 0.01 0.01 0.00 nVMP+ 0.03 0.00 0.20 0.00 0.00
psVAE 0.00 0.00 0.62 0.08 0.01 psNVAE 0.31 0.20 0.65 0.02 0.00 psNVMP 0.54
0.23 0.63 0.02 0.01 psNVMP+ 0.21 0.37 0.56 0.09 0.00

Table 3: (MS) Layer-wise latent informativeness with respect to T2 lesion
volume. Models with prefix “ps-” have partially supervised latent spaces (i.e.
are $\mathbf{z}_{P}$-supervised).

In tab:Informativeness, rows 1-4 show the $R^{2}$ scores of the unsupervised
models, which are generally poor; rows 5-8 show the same metrics for
supervised models where supervision is provided to $\mathbf{z}_{2}$
($\mathbf{z}_{P}$) as an additional lesion segmentation objective. Models with
autoregressive structures (nVAE, nVMP, nVMP+) benefit more from the
supervision - knowledge from the supervision is propagated to the unsupervised
“anatomical” latent units, resulting in higher $R^{2}$ scores even in the
unsupervised latent subsets. This shows that the model is indeed actively
using the latent structures.

Furthermore, by separately scaling each latent subgroup and seeing the changes
in the generated images, we can examine the features captured by each
individual latent group. Latent disentanglement, as indicated by the
remarkable disparity in layer-wise pathological attribute sensitivity to
scaling, is made evident with such visualisation.

fig:AttributeSensitivity
$\mathbf{z}_{0}$$\mathbf{z}_{1}$$\mathbf{z}_{2}$$\left(\mathbf{z}_{P}\right)$$\mathbf{z}_{3}$$\mathbf{z}_{4}$

Figure 6: (MS) Layer-wise pathological attribute sensitivity visualised by
individually scaling each layer in the latent hierarchy, from $\mathbf{z}_{0}$
(top row) to $\mathbf{z}_{4}$ (bottom row).

In this particular example (fig:AttributeSensitivity), the appearance of the
hyper-intense MS lesions in the synthesised images is relatively insensitive
to multiplicative perturbation in all but one latent layer, $\mathbf{z}_{2}$.
The layer with the highest pathological attribute sensitivity,
$\mathbf{z}_{2}$, is hence considered to be a disentangled “pathological”
latent subset $\mathbf{z}_{P}$.
We note that even in the unsupervised setting, disease-related features in the
synthesized images are noticeably more sensitive to changes in a small subset
of latent variables than the rest, which allows us to identify such a subset
as $\mathbf{z}_{P}$ and the rest as $\mathbf{z}_{A}$ (anatomical latent
subsets) in a post-hoc manner. Such disparity in pathological attribute
sensitivity is much more pronounced in the “selective supervision” setting
(bottom-left purple block in fig:GraphicalModel and fig:ArchDiagram), where
the additional supervision is given to a chosen layer $\mathbf{z}_{P}$.

fig:VariationPerLayerADNI

Figure 7: (AD) Variations captured by each layer of the model. Images at the
top row are fully resampled at each level of the hierarchy. On each subsequent
row $n$, we show the residual variation of layer $n$ by fixing latent codes at
the top $(n-1)$ layers.

fig:VamPriorClusters

Figure 8: (MS) Clusters discovered by VamPrior.