Smile-GANs: Semi-supervised clustering via GANs for dissecting brain disease heterogeneity from medical images

The adversarial loss is applied for training of discriminator $D$ and the mapping function $f$, which can be written as:

The mapping $f$ attempts to transform CN to corresponding fake PT data so that they follow similar distributions as real PT data. The discriminator $D$, representing the probability that $y$ come from the real data rather than generator, is trying to identify the fake PT data from the real PT data. Therefore, the discriminator attempts to maximize the adversarial loss function while the mapping $f$ attempts to minimize against it. The training process can be denoted as:

The change loss is to control the distance of transformations. As our model is applied to a lower dimensional representation of imaging data, region of interests (ROIs), we assume that only some specific regions will be affected by the disease process, whereas the rest should remain unchanged. To encourage sparsity, we define the change loss to be the $l_{1}$ distance between the fake PT data and the original CN data: $L_{\text{change}}(f)=E_{x\sim p_{\text{CN}},z\sim p_{\text{SUB}}}[||f(x,z)-x||_{1}]$.

Moreover, we formulate the cluster loss as $L_{\text{cluster}}(f,g)=E_{x\sim p_{\text{CN}},z\sim p_{\text{SUB}}}[||g(f(x,z))-z||_{2}]$. By controlling the distance between the sampled SUB variable z and the reconstructed SUB variable, we enforce the function $g\circ f$ to be an Identity function of $z$. This leads to the first property that $f(a,z)$ is an injective mapping for fixed $x=a$ such that different values of SUB variable $z$ will transform the same CN data to different PT data. With minimization of the cluster loss, there arises the second property: for $z_{1}\neq z_{2}$ and different CN data, $x_{1}\neq x_{2}$ $f(x_{1},z_{1})\neq f(x_{2},z_{2})$. These two properties of the mapping function $f$ are important, since they guarantee that the SUB variable $z$ is not ignored during training process and that there is no intersection among mapping directions (i.e., one PT data is only assigned to one subtype).

More importantly, with the assumption that $p_{\text{PT}}=p_{\text{Y'}}$ after $f$ and $g$ are properly trained [7], we can derive the equality between PT domain $Y$ and Fake-PT domain $Y^{\prime}$ and thus for $y\sim p_{\text{PT}}$, $g(y)=g(y^{\prime})=g(f(x,z))=z$, where $z\sim p_{\text{SUB}}$ indicates the subtype membership. In other words, $g$ can be a clustering function for quickly clustering unseen data.

With the aforementioned losses, we can write the full objective as:

with $\mu$ and $\lambda$ be two parameters controlling the relative importance of each Loss function during the training process. Through this objective, we want to find the mapping f and the clustering function g such that:

For faster convergence of model in implementation, the mapping function, instead of directly transforming the CN data to the fake PT data, first learns a change in the CN data and takes the sum of them to obtain the fake PT data. Therefore, the architecture of the mapping function $f$ can be divided into two phases as shown in Fig. 1(C). In the first phase, the CN data and the SUB variable are mapped to latent representations with same dimension through encoder and decoder [10], respectively. The second phase has one decoding structure mapping the dot-product of two representations to the change $\tilde{x}$, which is added to the CN data $x$ to acquire the fake PT data. The discriminator $D$ and the clustering function $g$ have similar encoding structures, with $D$ mapping PT/fake PT data to prediction vector with dimension 2 while the encoder $g$ mapping the fake PT data to a SUB representation.

First, we rewrite the procedure in section 2.1 as: $\min_{D}L_{\text{GAN}}(D)=E_{y\sim p_{\text{PT}}}[(D(y)-1)^{2}]+E_{z\sim p_{\text{SUB}},x\sim p_{\text{CN}}}[D(f(x,z))^{2}]$ and $\min_{f}L_{\text{GAN}}(f)=E_{z\sim p_{\text{SUB}},x\sim p_{\text{CN}}}[D(f(x,z)-1)^{2}]$. Using this least square loss instead of original log likelihood boosts stability of the training process [11]. Second, the SUB variable z is constructed as a one-hot latent variable with dimension K instead of one single value and thus the cross-entrophy loss is computed for the cluster loss defined in section 2.2.

We set two parameters to be $\mu=5$ and $\lambda=9$ for all experiments. Also,we performed gradient clip for each iteration to avoid the explosion of gradient during the training process. For optimization, we used ADAM optimizer [12] with learning rate 0.0004 for Discriminator $D$ and 0.002 for mapping $f$ and clustering function $g$. $\beta_{1}$ and $\beta_{2}$ are 0.5 and 0.999, respectively. More details about architectures and training procedures are present in Supplementary.

For real application, since the ground truth of subtypes is unknown, we adopt an approximation of the Wasserstein distance (WD) [13] as one metric for monitoring the training process and choosing the stopping point. For Smile-GANs, instead of deriving WD from optimization as the original paper introduced, we used the closed form formula to compute the distance. For stopping criteria, we assume that, in CN domain and in all subpopulations of PT domain, the lower dimensional representation of each data point (ROIs) is sampled from a multivariate Gaussian distribution. Though this assumption might be strong, it does enable us to estimate the WD quickly.

To be more specific, for each mapping direction $z=i$ and for all samples in CN domain $X=\{x_{1},x_{2},\cdots,x_{n}\}$, we calculate the mean vector $m_{1}^{i}$ and covariance matrix $C_{1}^{i}$ of $f(X,i)$. Also, from samples in PT domain Y, we take out the subset $Y_{i}=\{y_{j}\}^{i}$ such that $g(y_{j})=i$ and calculate the mean vector $m_{2}^{i}$ and covariance matrix $C_{2}^{i}$. With mean vectors and covariance matrices, we can compute the 2nd Wasserstein distance using the formula for two multivaraite gaussian measure:

If we further assume that all features are independent, we can derive diagonal covariance matrices which make the computation even faster and, based on our experiments, this assumption does not affect monitoring training process.

Moreover, to deal with rare cases when inconsistencies exist between WD and model performance, we also derive two other metrics: alteration quantity (AQ), which represents the number of subjects whose subtype memberships alter in the last five epochs. A small AQ represents high stability of the model. Lastly, the cluster loss, indicating the performance of clustering function $g$, is also considered as part of the stopping criteria.

Simulated Data Generation Simulated data were generated in a low dimensional space (i.e., 145 ROIs). For each subject, the 145 ROIs were simulated by sampling from a normal distribution $\mathcal{N}(1,0.1)$. In total, 1200 subjects were generated independently and then randomly split into two half-split sets, with each (600) being CN and pseudo-PT group, respectively. The atrophy simulation was only introduced for pseudo-PT subjects, which were further divided into 3 subtypes with the same number of subjects (200). For each subtype, the values of specific pre-selected ROIs were deceased by 10 to 20% randomly to simulate the severity of the atrophy. Moreover, to simulate covariate effects, we randomly sampled 200 subjects from both CN and pseudo-PT respectively and decreased the values by 10 to 20% in some other ROIs. The simulation ground truth for the covariate patterns and the 3 subtypes are shown in Fig.  2(c)(i) and (ii), respectively. Note that overlapping of ROIs across subtypes was imposed to better follow the nature of atrophy.

Experiments on Simulated Data To add variability of the clustering performance, we repeated the simulated data generation and ran the experiment independently 20 times. We first investigated the potential of WD for monitoring training process. Then we compared the clustering performance of Smile-GANs with traditional K-means [14] and Gaussian Mixtrue Model (GMM) [15]. Finally, the mapping function ($f$) of Smile-GANs can be used to visualize the subtypes’ atrophy patterns captured for fake PT data generation. We calculated the mean difference between CN and fake PT in K mapping directions and inspected ROIs with significant decreasing in values.

Data and Image Processing For experiments on real data, we first included 297 CN and 602 AD/MCI subjects for baseline T1-weighted (T1) MRIs from ADNI 2 database. The trained model was further applied to longitudinal data, undergoing more than one visit in 2 to 23 years from the baseline, from ADNI1 and BLSA dataset. The longitudinal data consist of 1323 CN and 610 MCI/AD subjects. For all T1 MRIs, brain tissue segmentation was performed using a multi-atlas segmentation technique [16] and 145 ROIs were derived as features for Smile-GANs. Those features were first harmonized to remove site effect [17], and then age and gender effects were corrected in a pooled sample of matched controls using a voxel-wise linear model. Moreover, gray matter (GM) tissue maps [18] were also segmented for voxel-wise statistical mapping.

Experiments on baseline data We first chose the optimal K. Smile-GANs was run for 20 times for different K (K = 2 to 5). The optimal K was chosen by the highest Adjusted Rand Index (ARI) [19], which quantifies the clustering stability across the 20 repetitions/models, and also guided by prior knowledge in literature. For each model, it assigns each patient to its corresponding subtype membership with the highest probability. The final clustering membership was determined by a consensus clustering strategy across the 20 models. The mapping function ($f$) of Smile-GANs was first applied to visualize the subtypes’ atrophy patterns captured for fake PT data generation. Moreover, Voxel-wise group comparison between CN and subtypes of real PT data were performed with AFNI 3dttest [20] with GM tissue maps [18].

Experiments on longitudinal data The trained 20 models from ANDI2 was applied to longitudinal data for determining the subtype membership for scans from all visits. Note that only subjects who were consistently assigned into the same subtype from more than 60% of the models (i.e., 12) were finally taken into account, which were studied for the disease longitudinal progression pathways from prodromal to full AD stages.

WD for training monitoring Fig. 2(a) shows the change of WD and clustering error (i.e., 1 - clustering accuracy) during training. Generally, the two metrics are consistent in monitoring the training process. Therefore, WD could be used as a surrogate of clustering accuracy when the latter is not available in real applications. Note that inconsistencies at the beginning of training or oscillations exist. Such circumstances can be filtered out by other two metrics, alteration quantity (AQ) and cluster loss. We propose to use WD, along with AQ and cluster loss, as metrics for monitoring the training process.

Clustering performance comparison across models Fig. 2(b) shows the clustering accuracy of Smile-GANs, K-means and GMM. Smile-GANs ($0.916\pm 0.025$) plainly outperforms the K-means ($0.615\pm 0.028$) and GMM ($0.611\pm 0.039$) for clustering those three subtypes.

Mapping function for visualization of atrophy patterns Fig. 2(c) shows the atrophy patterns captured by all different mapping directions (K=3). Smile-GANs can automatically identify the ground truth of atrophy patterns, while not capturing any covariate effects. Together with Fig. 2(b), our results clearly indicates that Smile-GANs can not only cluster subtypes accurately, but also automatically identify the underlying atrophy patterns for better interpretation.

For baseline experiments, the stopping criteria during training is that the highest WD of all mappings, AQ and cluster loss are smaller than 0.22, 35 and 0.001, respectively. The maximum number of epochs was set to be 6000 and models failing to converge were discarded.

Clustering stability for optimal K Fig. 3(A) shows the results of ARI for different K. Though K=2 gave the highest ARI, Smile-GANs roughly divided PT into one subtype with barely atrophy and another subtype with whole-brain atrophy. Based on the prior knowledge from the literature [4][5] and relatively higher ARI for K=4, we therefore selected K=4 for following experiments.

Neuroanatomical Heterogeneity between Subtypes and CN Fig. 3(B) shows regions with atrophy identified by the four mapping directions for fake data generation (Detail of the name of ROIs are present in Supplementary). Fig. 3(C) shows the voxel-based group comparison results for each subtype group of real subtypes versus CN group. The two approaches converge to the same findings: i) Subtype 1, referred as normal brain, exhibits no atrophy over the whole brain; ii) Subtype 2, denoted as diffuse atrophy atypical of AD, shows widespread atrophy patterns in frontal and temporal lobe, but medial temporal lobe is spared; iii) Subtype 3, referred as focal medial temporal lobe atrophy, shows localized atrophy in the hippocampus and the anterior-medial temporal cortex; iv) Subtype 4, denoted as typical-AD, displays severe atrophy over the whole brain.

Clinical Characteristics of clustering subtypes The clinical characteristics of the four subtypes are summarized in Table 1. Most of subjects in subtype 1 are MCI subjects while more than half of AD subjects are in subtype 4. All three clinical variables (i.e., Abeta, T-tau and WML) of subtype 1 significantly differ to those of the other subtypes. Subtype 2 and 3 also show significant difference in Abeta (p=0.018) and T-tau (p=0.003). Though not significant, subtype 2 has substantially higher WML load than subtype 3.

Table 2 shows the subtype membership conversion from baseline assignment to future longitudinal assignment using longitudinal AD, MCI and CN subjects. Most of subjects in subtype 1 at baseline remained unchanged for future membership assignment, but some of them progressed to other subtypes in later visits. A substantial proportion of subjects who were assigned into subtype 2 and 3 at baseline finally transformed into subtype 4, but the transformation between these two subtypes are very rare.

Alternatively, we also studied the change of probability for subtype membership assignment in AD and MCI subjects in a window of 6 years (Fig. 4). For subjects assigned into subtype 1 at baseline, they show increasing probabilities belonging to subtype 2 or 3 at early stage, and then to subtype 4 at later stage (Fig. 4 (a)). For subjects of subtype 2 and 3, they all have increasing probabilities belonging to subtype 4 in later visits (Fig. 4 (b) and (c)). Those results (Table 2 and Fig. 4) potentially indicate two distinct longitudinal disease progression pathways: i) subtypes 1 - 2 - 4, and ii) subtypes 1 - 3 - 4.

Detailed architectures of mapping function, clustering function and discriminator are provided in table 3 and table 4.

Detailed training procedure of Smile-GANs is disclosed by Algorithm 1.

Full names of ROIs shown in Fig. 3(B) are displayed in table 5.