Self-supervised Disentanglement of Disease Effects from Aging in 3D Medical Shapes

Abstract

Disentangling pathological changes from physiological aging in 3D medical shapes is crucial for developing interpretable biomarkers and patient stratification. However, this separation is challenging when diagnosis labels are limited or unavailable, since disease and aging often produce overlapping effects on shape changes, obscuring clinically relevant shape patterns. To address this challenge, we propose a two-stage framework combining unsupervised disease discovery with self-supervised disentanglement of implicit shape representations. In the first stage, we train an implicit neural model with signed distance functions to learn stable shape embeddings. We then apply clustering on the shape latent space, which yields pseudo disease labels without using ground-truth diagnosis during discovery. In the second stage, we disentangle factors in a compact variational space using pseudo disease labels discovered in the first stage and the ground truth age labels available for all subjects. We enforce separation and controllability with a multi-objective disentanglement loss combining covariance and a supervised contrastive loss. On ADNI hippocampus and OAI distal femur shapes, we achieve near-supervised performance, improving disentanglement and reconstruction over state-of-the-art unsupervised baselines, while enabling high-fidelity reconstruction, controllable synthesis, and factor-based explainability. Code and checkpoints are available at https://github.com/anonymous-submission01/medical-shape-disentanglement

Figures (4)

• Figure 1: Stage 1: Learns per-shape codes ($\mathbf{z}_i$) and a INR-nased SDF decoder ($G_{\theta}$). Unsupervised clustering is applied on the learned codes to create pseudo-labels. Stage 2: A variational autoencoder (VAE) models the distribution of shape codes and learns latents ($\mathbf{z}^v_i$). Pseudo-labels and age labels are used for disentangling specific latent variables, while frozen $G_{\theta}$ is used for reconstruction.
• Figure 1: Stage-1 ablation: purity (Pur (%)$\uparrow$), and mean-volume gap ($\Delta\overline{V}\uparrow$, mm$^3$) for shapes between the two clusters. Average volume of AD is lower compared to CN which is clearly represented by stage-1 clusters.
• Figure 2: Volume distribution of shapes in stage-1 cluslters.
• Figure 3: Rendering healthy and diseased shapes at different ages by latent traversal (trained by both real pseudo labels) shows consistent volume changes in ADNI (left) and deformation in OAI (right) that reflect the original dataset.