AI-based association analysis for medical imaging using latent-space geometric confounder correction

TL;DR

This work tackles confounding and interpretability in AI-driven medical image analysis by introducing a geometric latent-space framework. It retains confounder information in latent representations while identifying a confounder-free direction ${\vec{p}^*}$ orthogonal to confounders and maximally aligned with the learning target ${\vec{t}}$, guided by a correlation-based loss $L_{corr}$ and a joint objective with reconstruction $L_{rec}$. The method enables semantic visualization by sampling along ${\vec{p}^*}$ and reconstructing images, and supports semi-supervised learning to leverage unlabeled data. Across synthetic, facial-phenotype, and brain-imaging studies, confounder correction reduces spurious associations while preserving image reconstruction quality and providing interpretable visualizations for clinical and epidemiological investigations.

Abstract

This study addresses the challenges of confounding effects and interpretability in artificial-intelligence-based medical image analysis. Whereas existing literature often resolves confounding by removing confounder-related information from latent representations, this strategy risks affecting image reconstruction quality in generative models, thus limiting their applicability in feature visualization. To tackle this, we propose a different strategy that retains confounder-related information in latent representations while finding an alternative confounder-free representation of the image data. Our approach views the latent space of an autoencoder as a vector space, where imaging-related variables, such as the learning target (t) and confounder (c), have a vector capturing their variability. The confounding problem is addressed by searching a confounder-free vector which is orthogonal to the confounder-related vector but maximally collinear to the target-related vector. To achieve this, we introduce a novel correlation-based loss that not only performs vector searching in the latent space, but also encourages the encoder to generate latent representations linearly correlated with the variables. Subsequently, we interpret the confounder-free representation by sampling and reconstructing images along the confounder-free vector. The efficacy and flexibility of our proposed method are demonstrated across three applications, accommodating multiple confounders and utilizing diverse image modalities. Results affirm the method's effectiveness in reducing confounder influences, preventing wrong or misleading associations, and offering a unique visual interpretation for in-depth investigations by clinical and epidemiological researchers. The code is released in the following GitLab repository: https://gitlab.com/radiology/compopbio/ai_based_association_analysis}

Figures (12)

• Figure 1: The proposed AI approach for association analysis in medical imaging. (a) Geometry perspective of correlations between a target and a confounder variable ($\textbf{t}, \textbf{c}$), and its extension ($\vec{t}, \vec{c}$) into the latent space (n=3 latent dimensions) of an autoencoder. Plane O is orthogonal to $\vec{c}$. $\vec{{p}^{*}}$ is the vector projection of $\vec{t}$ onto plane O. $\textbf{d}$ is the latent representation of an input image and $\textbf{d}^{\prime}$ is its projection onto $\vec{{p}^{*}}$. $z_{p}^{*}$ is the distance between $\textbf{d}^{\prime}$ and the origin. For cases with $m$ confounders, the latent dimensions should be $n \geq m + 1$, so as to guarantee there exist a $\vec{{p}^{*}}$ orthogonal to $m$ confounders; (b) a directed acyclic diagram explains the relationships between $\textbf{t}, \textbf{c}$, and image I. We aim to extract image features associated with the learning target while being independent to the confounders. (c) The proposed approach in a neural network perspective. $[z_{1}, z_{2}, ..., z_{n}]$ are the learned latent features by the network, which construct the latent space shown in (a); $\textbf{X}$ and $\textbf{X}^{\prime}$ refer to the input and reconstructed image; $\uptheta_{enc}$, $\uptheta_{dec}$, $\uptheta_{pe}$ are the trainable parameters of encoder, decoder, and projection estimator.
• Figure 2: The distribution of the 2-D latent space for the synthetic images in the test set of Experiment 1, and the eleven reconstructed images sampling along the brightness-related vector, (a) without (i.e., vector $\vec{p}$) and (b) with correction (vector $\vec{{p}^{*}}$) for the confounding of circle radius, together with the predicted brightness $\hat{t}$ derived by Equation \ref{['eq:6_6']} and Equation \ref{['eq:6_7']}. $Z_{1}$-axis: the first dimension of the latent space; $Z_{2}$-axis: the second dimension. Each data point in the latent space represents an input image, which is denoted by its radius and brightness. After training, eleven frames were reconstructed by sampling eleven points along the vector $\vec{p}$ and $\vec{{p}^{*}}$ (Equation \ref{['eq:6_6']} and Equation \ref{['eq:6_7']}) to visualize the confounding effects. Whereas our method does not involve the estimation of vectors $\vec{t}$ and $\vec{c}$, we have manually included them in this figure only for the purpose of enhancing comprehension.
• Figure 3: The input images $\textbf{X}$ (8$\times$5 circle images), and the reconstrued images $\textbf{X}^{\prime}$ of different methods, in Experiment 1.
• Figure 4: Distribution of datapoints in the 2-D latent space via (a) unsupervised training and (b) our proposed supervised training, in Experiment 1.
• Figure 5: Interpretation heatmaps of facial changes in children with PAE using the proposed method: (a) without correction for confounders; (b) with correction for ethnicity, BMI, maternal smoking, maternal age, and sex. Red areas refer to inward changes of the face with respect to the geometric center of the head. Heatmap generation is detailed in Section \ref{['sec:sematic_feature_visualization']}.