Surreal-GAN:Semi-Supervised Representation Learning via GAN for uncovering heterogeneous disease-related imaging patterns

1 INTRODUCTION . Neuroimaging , conventional machine learning , and deep learning have offered unprecedented opportunities to understand the underlying mechanism of brain disorders over the past decades ( Davatzikos ( 2018 ) ) , and pave the road towards individualized precision medicine ( Bzdok & MeyerLindenberg ( 2018 ) ) . A large body of case-control studies leverage mass univariate group comparisons to derive structural or functional neuroimaging signatures ( Habeck et al . ( 2008 ) , Hampel et al . ( 2008 ) , Ewers et al . ( 2011 ) ) . However , these studies suffer from under-powered statistical inferences since the heterogeneous nature of neurological diseases usually violate the assumption that each group population are pathologically homogeneous . A developing body of methodologies aim to disentangle this heterogeneity with various machine learning methods . Zhang et al . ( 2016 ) used the Bayesian Latent Dirichlet Allocation ( LDA ) model to identify latent atrophy patterns from voxel-wise gray matter ( GM ) density maps . However , due to the nature of the LDA model , the density maps need to be first discretized , and the model hypothesized to exclusively capture brain atrophy while ignoring the potential regional enlargement . Young et al . ( 2018 ) proposed another method , Sustain , to uncover temporal and phenotypic heterogeneity by inferring both subtypes and stages . However , the model only handles around 10-20 large brain regions without more detailed information . Notably , both methods applied directly in the PT domain , confront main limitations in avoiding potential disease-unrelated brain variations . In con- trast , semi-supervised clustering methods ( Varol et al . ( 2016 ) , Dong et al . ( 2015 ) , Wen et al . ( 2021 ) , Yang et al . ( 2021 ) ) were proposed to cluster patients via the patterns or transformations between the reference group ( CN ) and the target patient group ( PT ) . A recent proposed deep learning-based semi-supervised clustering method , termed Smile-GAN ( Yang et al . ( 2021 ) ) , achieved better clustering performance by learning multiple mappings from CN to PT with an inverse clustering function for both regularization and clustering . Albeit these methods demonstrated potential in clinical applications , they have a major limitation . In particular , these methods model disease heterogeneity as a dichotomous process and derive a discrete output , i.e. , subtype , ignoring the fact that brain disorders progress along a continuum . Moreover , variations among subtypes were contributed by both spatial and temporal differences in disease patterns , thus confounding further clinical analysis. ( Fig . 1a ) To address the aforementioned limitations , we propose a novel method , Surreal-GAN ( SemiSupervised Representation Learning via GAN ) , for deriving heterogeneity-refined imaging signatures . Surreal-GAN is a significant extension of Smile-GAN under the same principle of semisupervised learning by transforming data from CN domain X to PT domain Y. Herein , the key extension is to represent complex disease-related heterogeneity with low dimensional representations with each dimension indicating the severity of one relatively homogeneous imaging patterns ( Fig . 1b ) . We refer to these low dimensional representations as R-indices ( ri , where i represent the ith dimension ) of disease pattern . The first key point of the method is modelling disease as a continuous process and learning infinite transformation directions from CN to PT , with each direction capturing a specific combination of patterns and severity . The idea is realized by learning one transformation function which takes both normal data and a continuous latent variable as inputs and output synthesized-PT data whose distribution is indistinguishable from that of real PT data . The second key point is controlling the monotonicity ( from light to severe with increasing index value ) of disease patterns through one monotonicity loss and double sampling procedure . The third key point is boosting the separation of captured disease patterns through one orthogonality regularization . The fourth key point is introducing a inverse mapping function consisting of a decomposition part and a reconstruction part which further ensure that patterns synthesized through transformation are exactly captured for inferring representations of disease patterns , the R-indices . Besides the above mentioned regularization , function continuity , transformation sparsity and inverse consistency further guide the model to capture meaningful imaging patterns . To support our claims , we first validated Surreal-GAN on semi-synthetic data with known ground truth of disease patterns and severity . We compared performance of Surreal-GAN to NMF ( Lee & Seung ( 1999 ) ) , Factor Analysis and LDA ( Blei et al . ( 2001 ) ) models , and subsequently studied the robustness of model under various conditions . Finally , we applied the model to one real dataset for Alzheimer ’ s disease ( AD ) , and demonstrated significant correlations between the proposed imaging signatures and clincial variables . 2 METHODS . The schematic diagram of Surreal-GAN is shown in Fig . 1d . The model is applied on regional volume data ( volumes of different graymatter ( GM ) and whitematter ( WM ) regions ) derived from structural MRI . The essential idea is to learn one transformation function f : X ∗ Z → Y which transforms CN data x to different synthesized PT data y′ = f ( x , z ) , with latent variable z specifying distinct combinations of patterns and severity . Here , the Latent ( LAT ) domain , Z = { z : 0 ≤ zi ≤ 1 ∀1 ≤ i ≤ M } , is a class of vectors with dimension M ( predefined number of patterns ) . We denote four different data distributions as x ∼ pcn ( x ) , y ∼ ppt ( y ) , y′ ∼ psyn ( y′ ) and z ∼ plat ( z ) , respectively , where z ∼ plat ( z ) is sampled from a multivariate uniform distribution U [ 0 , 1 ] M , rather than a categorical distribution as in Smile-GAN . A continuous z variable lay the foundation for learning continuous representations . In addition to the transformation function , an adversarial discriminator D is introduced to distinguish between real PT data y and synthesized PT data y′ , thereby ensuring that the synthesized PT data are indistinguishable from real PT data ( i.e . minimizing the distance between real PT data distribution and synthesized PT data distribution ) . A continuous latent variable does not necessarily lead to desired properties of CN to PT transformation ( e.g . separation/monotonicity of captured patterns ) . That is , there are a number of functions potentially achieving equality in distributions , thus making it hard to guarantee that the transformation function learned by the model is closely related to the underlying pathology progression process . Moreover , during the training procedure , the transformation function backboned by the neural network tends to ignore the latent variable z . Even when the Latent variable z is emphasized through regularization , different components of z variable do not naturally lead to separate disease patterns and are not positively correlated with pattern severity . Therefore , with the assumption that there is one true underlying function for real PT data y = h ( x , z ) , Surreal-GAN aims to boost the transformation function f to approximate the true underlying function h , while having the latent variable , z , contributing to recognizable imaging patterns rather than being random ” activators ” of the GAN . For this purpose , we constrain the function class of f via different types of regularization as follows : 1 ) encourage sparse transformations ; 2 ) enforce Lipschitz continuity of functions ; 3 ) introduce one inverse function g : Y → Z serving to decompose and reconstruct ; 4 ) boost orthogonality among synthesized/captured patterns ; 5 ) enforce monotonicity of synthesized patterns and their positive correlations with components of z . The training procedure of Surreal-GAN is thus a adversarial learning process constrained by these regularization . 2.1 ADVERSARIAL LOSS . The adversarial loss ( Goodfellow et al . ( 2014 ) ) is used for iterative training of discriminator D and the transformation function f , which can be written as : LGAN ( D , f ) = Ey∼ppt ( y ) [ log ( D ( y ) ) ] + Ex∼pcn ( x ) , z∼Plat ( z ) [ log ( 1−D ( f ( x , z ) ) ) ] ( 1 ) = Ey∼ppt ( y ) [ log ( D ( y ) ) ] + Ey′∼psyn ( y′ ) [ log ( 1−D ( y ′ ) ) ] ( 2 ) The transformation function f attempts to transform CN data to synthesized PT data so that they follow similar distributions as real PT data . The discriminator , providing a probability that y comes from the real data rather than transformation function , tries to identify the synthesized PT data and distinguish it from real PT data . Therefore , the discriminator D attempts to maximize the adversarial loss function while the transformation function f attempts to minimize against it . The training process can be denoted as : min f max D LGAN ( D , f ) = Ey∼ppt ( y ) [ log ( D ( y ) ) ] + Ey′∼psyn ( y′ ) [ log ( 1−D ( y ′ ) ) ] ( 3 ) 2.2 REGULARIZATION . 2.2.1 SPARSE TRANSFORMATION . We assumed that disease process will not change brain anatomy dramatically and primarily only affect certain regions throughout most of disease stages , which means that the true underlying transformation only moderately changes some regions while keeping the rest mildly affected with most of z variable values . Therefore , to encourage sparsity and control the distance of transformations , we defined one change loss to be the l1 distance between the synthesized PT data and the original CN data : Lchange ( f ) = Ex∼pcn ( x ) , z∼plat ( z ) [ ||f ( x , z ) − x||1 ] ( 4 ) 2.2.2 LIPSCHITIZ CONTINUITY AND INVERSE MAPPING . First , by enforcing the transformation function f to be K1-Lipschitz continuous , we can derive that , for fixed latent variable z = a and ∀x1 , x2 ∈ X , ||f ( x1 , a ) − f ( x2 , a ) ||2 ≤ K1||x1−x2||2 . Thus , by controlling the constant K1 , we constrained the transformation function to preserve the original distances among CN data instead of scattering them into the PT domain if a same combination of disease patterns is imposed to them . Second , to prevent latent variable z from being ignored , we introduced an inverse mapping function from PT domain to Latent domain , g : Y → Z . By Lemma1 ( section A.1 ) , with function g being K2-Lipschitz continuous and d ( · , · ) being any distance metric satisfying triangular inequality , we can derive that , ∀z1 , z2 ∼ plat ( z ) , z1 6= z2 , and x̄ ∼ pcn ( x ) , d ( f ( x̄ , z1 ) , f ( x̄ , z2 ) ) is lower bounded by ( d ( z1 , z2 ) K2 − 1 K2 ( d ( g ( f ( x̄ , z1 ) ) , z1 ) + d ( g ( f ( x̄ , z2 ) ) , z2 ) ) ) . Therefore , by minimizing the distance between sampled latent variable z and reconstructed latent variable g ( f ( x , z ) ) , we can control differences between synthesized patterns to be positively related to the distances between latent variables , and thus to be non-trivial ( i.e. , same CN data is transformed to significantly different PT data with different z variables ) . Therefore , we define a reconstruction loss as the l2 distance between g ( f ( x , z ) ) and z : Lrecons ( f , g ) = Ex∼pcn ( x ) , z∼plat ( z ) [ ||g ( f ( x , z ) ) − z||2 ] ( 5 )

Summary: Surreal-GAN aims to create fake pathological data with a latent variable and a healthy input and includes an inverse function that predicts the latent variable from a fake/real pathological data. However, unlike its predecessor model Smile-GAN, this adds 5 regularizations (L1 loss between generated output and input, L2 loss between the latent variable and the reverse generated latent variable, decomposing the latent variable, making the latent variables orthogonal to each other, and adding positive correlation to disease pattern and the latent variable) to make the latent variable more interpretable and continuous. This would increase the interpretability since instead of clustering the real data it could be used as understanding how severe the patient’s disease progress is. The results show that the model has a higher c-index than NMF, LDA and FA model predictions. It also shows that the best number of patterns for the model to predict is two, and it indeed has positive correlation to the severity of the disease.