{'fraction_non_alphanumeric': 0.03339652096342551, 'fraction_numerical': 0.009756913470115968, 'mean_word_length': 5.1937154696132595, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 1, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Clinically Plausible Pathology-Anatomy Disentanglement in Patient Brain MRI with Structured Variational Priors

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.


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 re-gions (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 Figure 1(a), 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 overgeneralisation (Hu et al., 2018). This may lead to lower synthesis quality or an inability to preserve crucial finer details in medical images. Figure 1(b) 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. In this work, we propose to address these issues by using structured variational inference (Hoffman and Blei, 2015) for finegrained 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 ( Figure 2). 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.

Model
Our model accepts image space observations sampled from a dataset x ∈ R wx×hx ∼ 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 Figure 2. Each group or layer consists of spatial latent variables (You et al., 2019) at various resolutions scales, denoted as z = {z l ∈ R w l ×h l ×c l ; l ∈ {0, 1, ..., L}}. The inference and generative models can therefore be expressed as follows: Observed or latent (pathological or anatomical) variables

BCE( f(zP), P)
Optional Supervision Figure 2: The hierarchically structured latent space with (L + 1) disjoint variable groups (layers), encoder distribution parameters (∆µ l , ∆σ l ) and decoder distribution parameters (µ l , σ l ). Dashed lines are active during inference. Solid lines are active during generation and, if residual parameterisation is used, also during bidirectional inference.
In this paper, we examine four ways to construct the ELBO objective, summerized in Table 1.
(1) VAE (M.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: (2) nVAE (M.2): 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 θ in nVAE not only characterize the generative distribution p θ (3b) but are also a part of the inference model and hence play an important role in characterizing the posterior distribution q φ,θ (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. Figure 3: Network Architecture

Model ELBO
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(z l ) is dependent on those of the preceding layers p(z <l ), each element within the same latent group z l := {z i l , ...} still adhere to the independence assumption as ∀l ∈ {0, 1, ..., L}, z i l |= z j l ; i = j. We can hence calculate the relative KLD loss for each element with a simple analytic expression: 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 (M.3) replaces the standard-Gaussian prior of the topmost layer of nVAE with a K-component multimodal VamPrior 1 K K k=1 q φ (z l |u k ) characterized by trainable pseudo-inputs u and encoder parameters φ. The subsequent layers still adhere to the hier-archical residual parameterisation of nVAE (we retain the residual Gaussian parameterisation for subsequent layers). The implication is that only z 0 is multi-modal whereas the distributions of all lower level "residual deviations" are assumed to be Gaussian.
(4) nVMP+ (M.4) 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.

Experiments and Results
We validate our approach on two brain MRI datasets: the publically available Alzheimer's Disease Neuroimaging Initiative 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 2, VAEs with more expressive structured priors indeed outperform their mean-field counter-part at the same model capacity in terms of image reconstruction fidelity. We additionally examine model behaviours in a supervised learning setting depicted in the bottom-left (purple) block in Figure 2 and Figure 3. In this setting, we supplement the MS model with a lesion segmentation objective between a chosen "pathological" latent subset z P and expert pathology (lesion) segmentation labels P. The rest of the latent space (that remains unsupervised) are regarded as the anatomical latent subsets, denoted as 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 z P comapred to z A , as shown in Figure 6, 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 z A , as in, those unsupervised latent units attain a higher linear predictability (Lasso regression R 2 scores (Eastwood and Williams, 2018), Table 3) with respect to lesion volume. This is in contrast to the behaviour of the baseline meanfield VAE, where information from the supervision task is constrained within the supervised group 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.  We can qualitatively evaluate pathologyanatomy 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 Figure 4 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 Figure 5: Visualising conditional distributions. Images in the top row are generated by fixing z A to that of Subject B1 and resampling the layer corresponding to z P . Images in the bottom row are generated by obtaining z A from other real samples and fixing z P to that of Subject B1.
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 Figure 4, we may visualise many possible disease states given this subject's anatomy ( Figure 5, top row) or explore how this subject's lesions would manifest on other subjects' brain anatomies ( Figure 5, bottom row).

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 pathologicallysound details, which could be meaningful for various downstream tasks. 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 (CI-FAR) 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. In Table 3, 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 z 2 (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.

References
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.
In this particular example (Figure 6), the appearance of the hyper-intense MS lesions in the synthesised images is relatively insensitive to multiplicative perturbation in all but one latent layer, z 2 . The layer with the highest pathological attribute sensitivity, z 2 , is hence considered to be a disentangled "pathological" latent subset 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 z P and the rest as 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 Figure 2 and Figure 3), where the additional supervision is given to a chosen layer z P .