Stochastic Siamese MAE Pretraining for Longitudinal Medical Images

TL;DR

STAMP tackles the challenge of learning temporally rich representations from longitudinal 3D medical images by introducing a time-conditioned, stochastic Siamese MAE. It combines a time difference encoding with a learned prior and a future-aware posterior to perform conditional variational inference, enabling robust forecasting of disease progression from a single baseline visit. Empirical results across OCT and MRI datasets show STAMP consistently outperforms prior MAE-based and foundation-model approaches on AMD, GA, and AD progression tasks, including scenarios with irregular visit intervals and domain shifts. The approach offers a scalable, label-efficient pathway for personalized prognosis in progressive diseases, with potential to inform follow-up scheduling and early interventions.

Abstract

Temporally aware image representations are crucial for capturing disease progression in 3D volumes of longitudinal medical datasets. However, recent state-of-the-art self-supervised learning approaches like Masked Autoencoding (MAE), despite their strong representation learning capabilities, lack temporal awareness. In this paper, we propose STAMP (Stochastic Temporal Autoencoder with Masked Pretraining), a Siamese MAE framework that encodes temporal information through a stochastic process by conditioning on the time difference between the 2 input volumes. Unlike deterministic Siamese approaches, which compare scans from different time points but fail to account for the inherent uncertainty in disease evolution, STAMP learns temporal dynamics stochastically by reframing the MAE reconstruction loss as a conditional variational inference objective. We evaluated STAMP on two OCT and one MRI datasets with multiple visits per patient. STAMP pretrained ViT models outperformed both existing temporal MAE methods and foundation models on different late stage Age-Related Macular Degeneration and Alzheimer's Disease progression prediction which require models to learn the underlying non-deterministic temporal dynamics of the diseases.

Figures (5)

• Figure 1: Top: HARBOR dataset for wet-AMD conversion from iAMD, arrow pointing the subretinal fluid. Middle: PINNACLE dataset for GA conversion from iAMD, arrow pointing increasing light transmission due to atrophy. Bottom: ADNI dataset, CN to MCI to AD with difference map, red/green means decrease/increase in pixel intensity.
• Figure 2: An overview of STAMP. Two 3D volumes (${\bm{x}}_t$ and ${\bm{x}}_{t+\Delta t}$) $\Delta t$ apart from a patient are used as input. After patchifying both scans, only the future visit is masked ($\Tilde{{\bm{x}}}_{t+\Delta t}$) and a learnable CLS token is attached to both branches. $\Delta t$ and the subsequent summation indicate TE added to the CLS token in $x_t$ branch. A ViT-based encoder $f$ embeds the past visit and visible patches of the future into $h_t$ and $\Tilde{h}_{t+\Delta t}$. For the stochasticity, the posterior ($q_\phi$) is learned from the embeddings $CLS_t$ and $CLS_{t+\Delta t}$ of the partially visible future, while the prior ($p_\psi$) is learned from $CLS_t$ and TE. After sampling ($\Tilde{z}_{t+\Delta t}$) from the posterior $q_\phi$, a cross-attention-based decoder queries $\Tilde{h}_{t+\Delta t}$ against $[\Tilde{z}_{t+ \Delta t}, h_t]$ to reconstruct the future visit ($\hat{{\bm{x}}}_{t+\Delta t}$). Once pretrained, the components within the dashed blue line are available during inference for the downstream task.
• Figure 3: Left: PCA projection of 20 sampled tokens from the learned prior $p_\psi$ per discrete $\Delta t$. Right: PCA projection of 20 sampled tokens from the learned prior $p_\psi$ per fractional $\Delta t$. The temporal alignment indicates diverse but time sensitive stochastic sampling in both cases.
• Figure 4: Attention map visualization of STAMP w.r.t pretrained CLS. Each row is from a different volume. Left: Attention maps without using TE. Middle: Attention maps with 3 months prompted TE. Right: Attention maps with 15 months prompted TE.
• Figure 5: PCA of TE per discrete ($\bullet$) and fractional (X) $\Delta t$. The number indicates deviation of a fractional from the interpolated value between the discrete TEs.