Stochastic Siamese MAE Pretraining for Longitudinal Medical Images

Taha Emre; Arunava Chakravarty; Thomas Pinetz; Dmitrii Lachinov; Martin J. Menten; Hendrik Scholl; Sobha Sivaprasad; Daniel Rueckert

cs.LGDec 29, 2025

Stochastic Siamese MAE Pretraining for Longitudinal Medical Images

Authors: Taha Emre, Arunava Chakravarty, Thomas Pinetz, Dmitrii Lachinov, Martin J. Menten, Hendrik Scholl, Sobha Sivaprasad, Daniel Rueckert, +4 more

Organizations: Institute of Artificial Intelligence, Center for Medical Data Science, Medical University of Vienna, Austria · Department of Ophthalmology and Optometry, Medical University of Vienna, Austria · BioMedIA, Department of Computing, Imperial College London, London, United Kingdom · Chair for AI in Healthcare and Medicine, Technical University of Munich, Munich, Germany · Department of Clinical Pharmacology, Medical University of Vienna, Vienna, Austria · Pallas Kliniken AG, Pallas Klinik Zürich, Zürich, Switzerland · European Vision Institute, Basel, Basel-Stadt, Switzerland · Moorfields National Institute for Health and Care Biomedical Research Centre, Moorfields Eye Hospital, London, United Kingdom · Institute of Ophthalmology, University College London, London, United Kingdom · Faculty of Medicine, University of Southampton, Southampton, Hampshire, United Kingdom · Ophthalmic Image Analysis Group (OPTIMA), Medical University of Vienna, Austria

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 & tables

Fig. 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.

Fig. 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. TE and the subsequent summation indicate Temporal Encoding ( TE ) added to the CLS token with $\Delta t$ prompting 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 stochastic embedding ( SE ) ( $\hat{z}_{t+\Delta t}$ ) from the posterior $q_{\phi}$ , a cross-attention-based decoder queries $\tilde{h}_{t+\Delta t}$ against $[\hat{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. $\sim\!\!\triangleright$ indicates sampling operation.

Fig. 3: Intra-visit intervals of ADNI dataset. Dashed lines indicate 1-, 2-, and 3-years.

Dataset	Class	# of Patients	# of Visits	Interval (months)
HARBOR	pretraining	579	12770	$1.0\pm 1.0$
	iAMD $\rightarrow$ iAMD	431	8141
	iAMD $\rightarrow$ wet-AMD	117	1967
PINNACLE	iAMD $\rightarrow$ iAMD	451	4260	$3.2\pm 3.3$
PINNACLE	iAMD $\rightarrow$ GA	99	621	$3.2\pm 3.3$
ADNI	CN $\rightarrow$ CN	584	2665	$12.9\pm 8.5$

TABLE I: Statistics of the datasets. Disease pathways are defined as a detected conversion during the study.

	Pretraining		Downstream		6-months			12-months
	SE	TE	SE	TE	AUROC ↑	PRAUC ↑	BACC ↑	AUROC ↑	PRAUC ↑	BACC ↑
1	✗	✓	-	✗	$0.634\pm 0.019$	$0.111\pm 0.019$	$0.587\pm 0.012$	$0.602\pm 0.064$	$0.156\pm 0.039$	$0.564\pm 0.052$
2	✗	✓	-	✓	$0.675\pm 0.043$	$0.120\pm 0.028$	$0.612\pm 0.036$	$0.640\pm 0.053$	$\underline{0.188\pm 0.041}$	$0.591\pm 0.033$
3	✓	✗	✗	-	$0.637\pm 0.048$	$0.100\pm 0.017$	$0.592\pm 0.028$	$0.601\pm 0.038$	$0.138\pm 0.011$	$0.571\pm 0.033$
4	✓	✗	✓	-	$0.627\pm 0.033$	$0.088\pm 0.015$	$0.579\pm 0.016$	$0.602\pm 0.028$	$0.134\pm 0.012$	$0.561\pm 0.024$
5	✓	✓	✗	✓	$\underline{0.690\pm 0.036}$	$\underline{0.127\pm 0.027}$	$0.612\pm 0.043$	$\underline{0.652\pm 0.061}$	$0.180\pm 0.063$	$\underline{0.601\pm 0.052}$

TABLE II: Ablation for the proposed temporal (TE) and stochastic (SE) components during pretraining and the downstream task on the HARBOR dataset. Row 7 corresponds to STAMP. In the Pretraining column, (✗) indicates a method without the component. In the Downstream column, (✗) indicates a component added in the pretraining but not used in the downstream task. PRAUC baselines for 6 and 12 months are 0.05 and 0.11, respectively.

	Pretraining		Downstream		1-year			3-years
	SE	TE	SE	TE	AUROC ↑	PRAUC ↑	BACC ↑	AUROC ↑	PRAUC ↑	BACC ↑
1	✗	✓	-	✗	$0.751\pm 0.002$	$0.195\pm 0.002$	$0.699\pm 0.002$	$0.746\pm 0.001$	$0.372\pm 0.001$	$0.696\pm 0.005$
2	✗	✓	-	✓	$0.779\pm 0.001$	$0.214\pm 0.002$	$0.702\pm 0.012$	$0.762\pm 0.006$	$0.362\pm 0.016$	$0.665\pm 0.019$
3	✓	✗	✗	-	$0.770\pm 0.006$	$\mathbf{0.300\pm 0.014}$	$0.692\pm 0.005$	$0.795\pm 0.002$	$0.438\pm 0.005$	$\mathbf{0.714\pm 0.002}$
4	✓	✗	✓	-	$0.787\pm 0.001$	$0.248\pm 0.004$	$0.698\pm 0.006$	$\mathbf{0.799\pm 0.002}$	$\mathbf{0.460\pm 0.002}$	$0.710\pm 0.016$
5	✓	✓	✗	✓	$\underline{0.799\pm 0.001}$	$0.269\pm 0.002$	$0.698\pm 0.003$	$0.771\pm 0.007$	$0.398\pm 0.021$	$0.693\pm 0.010$

TABLE III: Ablation for the proposed temporal (TE) and stochastic (SE) components during pretraining and the downstream task on the ADNI dataset. Row 7 corresponds to STAMP. In the Pretraining column, (✗) indicates a method without the component. In the Downstream column, (✗) indicates a component added in the pretraining but not used in the downstream task. PRAUC baselines for 1 and 3 years are 0.07 and 0.18, respectively.

HARBOR	Model	6-months			12-months
	Model	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$
	STAMP embedding	$0.700\pm 0.047$	$0.140\pm 0.042$	$0.615\pm 0.049$	$0.671\pm 0.057$	$0.200\pm 0.066$	$0.607\pm 0.042$
	STAMP token	$0.694\pm 0.035$	$0.139\pm 0.033$	$0.620\pm 0.030$	$0.674\pm 0.045$	$0.198\pm 0.051$	$0.607\pm 0.037$
ADNI		1-year			3-years
	STAMP embedding	$0.812\pm 0.001$	$0.277\pm 0.004$	$0.736\pm 0.009$	$0.773\pm 0.001$	$0.416\pm 0.013$	$0.691\pm 0.001$
	STAMP token	$0.793\pm 0.002$	$0.211\pm 0.001$	$0.762\pm 0.011$	$0.796\pm 0.002$	$0.371\pm 0.001$	$0.732\pm 0.008$

TABLE IV: Comparison of time encoding as an additive bias, i.e., the original STAMP design denoted as STAMP embedding , and as a separate input token, denoted as STAMP token . Results are reported over 4 folds on HARBOR and over repeated experiments on ADNI.

Model	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$
0.8-0.2	$0.788\pm 0.003$	$0.258\pm 0.002$	$0.723\pm 0.013$
25%	$0.778\pm 0.002$	$0.251\pm 0.008$	$0.718\pm 0.014$
STAMP	$\mathbf{0.812\pm 0.001}$	$\mathbf{0.277\pm 0.004}$	$\mathbf{0.736\pm 0.009}$

TABLE V: Ablation studies for Masking ratio (75% and 25%) and symmetric KL term weighting (0.2–0.8 and 0.8–0.2 for posterior-prior respectively). The models are pretrained with the ADNI dataset.

Model	6-months			12-months
Model	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$
VisionFM [ 45 ]	$0.631\pm 0.036$	$0.090\pm 0.010$	$0.583\pm 0.075$	$0.593\pm 0.053$	$0.126\pm 0.017$	$0.553\pm 0.058$
ImageNet	$0.665\pm 0.029$	$0.096\pm 0.026$	$\underline{0.614\pm 0.033}$	$\underline{0.640\pm 0.022}$	$0.151\pm 0.052$	$0.578\pm 0.013$
RetFound [ 61 ]	$0.636\pm 0.043$	$0.091\pm 0.027$	$0.560\pm 0.024$	$0.551\pm 0.024$	$0.121\pm 0.042$	$0.531\pm 0.033$
TC [ 21 ]	$0.643\pm 0.033$	$0.101\pm 0.011$	$0.611\pm 0.050$	$0.599\pm 0.067$	$0.143\pm 0.042$	$0.564\pm 0.057$
MAE [ 27 ]	$0.620\pm 0.027$	$0.087\pm 0.014$	$0.562\pm 0.025$	$0.602\pm 0.056$	$0.137\pm 0.033$	$0.531\pm 0.029$

TABLE VI: Linear evaluation of wet-AMD conversion prediction on pretrained models within two time‐windows: 6 and 12 months on the HARBOR Dataset. ImageNet, VisionFM and RetFound are fully finetuned. The results are reported over 4-folds

Model	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$
TC [ 21 ]	0.770 $\pm 0.058$	$0.150\pm 0.054$	$0.702\pm 0.046$
MAE [ 27 ]	$0.830\pm 0.025$	$0.208\pm 0.042$	$0.752\pm 0.046$
SiamMAE [ 24 ]	$\underline{0.833\pm 0.036}$	$\underline{0.229\pm 0.053}$	$\underline{0.755\pm 0.040}$
CropMAE [ 22 ]	$0.818\pm 0.026$	$0.174\pm 0.041$	$0.744\pm 0.023$
RSP [ 32 ]	$0.769\pm 0.031$	$0.155\pm 0.036$	$0.690\pm 0.039$
STAMP	$\mathbf{0.848\pm 0.020}$	$\mathbf{0.240\pm 0.061}$	$\mathbf{0.769\pm 0.047}$

TABLE VII: Linear evaluation of GA conversion prediction within a 12 months time window on the PINNACLE Dataset. The models are pretrained with the HARBOR dataset. The results are reported over 5-folds

Model	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$
10%	$0.679\pm 0.016$	$0.111\pm 0.025$	$0.624\pm 0.015$
25%	$0.691\pm 0.044$	$0.140\pm 0.025$	$0.617\pm 0.015$
50%	$0.719\pm 0.034$	$0.140\pm 0.024$	$0.642\pm 0.037$
100%	$0.700\pm 0.047$	$0.140\pm 0.042$	$0.615\pm 0.049$

TABLE VIII: Label-efficiency analysis for 6-month wet-AMD conversion prediction on HARBOR using different portions of the labeled downstream training data. The pretrained STAMP encoder is frozen and only the downstream classifier is trained. Results are reported over 4 folds.

Fig. 4: 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.

Fig. 5: 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 .

Fig. 6: 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 TE s.

Model	GFLOP	# of Params (M)	Runtime (min)
MAE [ 27 ]	$45.1$	$104.8$	$1.9$
SiamMAE [ 24 ]	$153.6$	$109.6$	$10.9$
CropMAE [ 22 ]	$153.6$	$109.6$	$10.9$
RSP [ 32 ]	$281.9$	$116.1$	$19.1$
STAMP	$153.7$	$117.4$	$10.6$

TABLE IX: Number of parameters, FLOP per iteration and runtime per epoch of pretraining

Model	1-year			3-years
Model	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$	AUROC $\uparrow$	PRAUC $\uparrow$	BACC $\uparrow$
Random init	$0.690\pm 0.006$	$0.142\pm 0.002$	$0.628\pm 0.010$	$0.713\pm 0.030$	$0.317\pm 0.039$	$0.652\pm 0.015$
ImageNet	$0.736\pm 0.007$	$0.188\pm 0.013$	$0.639\pm 0.008$	$0.712\pm 0.041$	$0.301\pm 0.056$	$0.670\pm 0.024$
yAware [ 18 ]	$0.730\pm 0.001$	$0.188\pm 0.001$	$0.625\pm 0.009$	$0.716\pm 0.004$	$0.306\pm 0.004$	$0.629\pm 0.004$
MAE [ 27 ]	$0.721\pm 0.002$	$0.152\pm 0.002$	$0.634\pm 0.012$	$0.712\pm 0.002$	$0.285\pm 0.004$	$0.674\pm 0.011$
SiamMAE [ 24 ]	$0.775\pm 0.004$	$\underline{0.220\pm 0.001}$	$0.721\pm 0.006$	$\mathbf{0.787\pm 0.001}$	$0.362\pm 0.003$	$\mathbf{0.726\pm 0.002}$

TABLE X: Linear evaluation of AD conversion prediction on pretrained models within two time windows: 1 and 3 years on the ADNI Dataset. The results are reported over repeated experiments

Model	c-index $\uparrow$
TC [ 21 ]	$\underline{0.664\pm 0.023}$
MAE [ 27 ]	$0.604\pm 0.074$
SiamMAE [ 24 ]	$0.649\pm 0.053$
CropMAE [ 22 ]	$0.603\pm 0.061$
RSP [ 32 ]	$0.599\pm 0.146$
STAMP	$\mathbf{0.747\pm 0.071}$

TABLE XI: Linear evaluation of wet-AMD risk model. The models are pretrained with the HARBOR dataset. The results are reported in Concordance Index with bootstrapping

Explore similar work

AD-DAE: Alzheimer's Disease Progression Modeling with Unpaired Longitudinal MRI using Diffusion Auto-Encoders
Nov 8, 2025Ayantika Das, Arunima Sarkar, Keerthi Ram +1Alzheimer Readmission Prediction
Department of Electrical Engineering, Indian Institute of Technology Madras (IITM), Chennai 600036, India · Sudha Gopalakrishnan Brain Centre (SGBC), IITM, Chennai 600036, India
CLIMB: Controllable Longitudinal Brain Image Generation using Mamba-based Latent Diffusion Model and Gaussian-aligned Autoencoder
Apr 17, 2026Duy-Phuong Dao, Muhammad Taqiyuddin, Jahae Kim +4Identity-Conditioned Latent Diffusion Models Latent Diffusion Model
Department of Artificial Intelligence Convergence, Chonnam National University, South Korea · Department of Nuclear Medicine, Chonnam National University, South Korea · University of South Australia, Australia +2
Self-supervised Pre-training Helps Retinal Disease Progression Modelling Most When Data Is Scarce
Sep 14, 2026Ifeoma Veronica Nwabufo, Julius Gervelmeyer, Sarah Müller +1Self-Supervised Learning Longitudinal Imaging
Hertie Institute for AI in Brain Health, Faculty of Medicine, University of Tübingen, Germany · Tübingen AI Center, University of Tübingen, Germany

Nov 8, 2025, cs.CVJ/K move · Enter open · S save

AD-DAE: Alzheimer's Disease Progression Modeling with Unpaired Longitudinal MRI using Diffusion Auto-Encoders

Ayantika Das, Arunima Sarkar, Keerthi Ram, Mohanasankar Sivaprakasam

Department of Electrical Engineering, Indian Institute of Technology Madras (IITM), Chennai 600036, India · Sudha Gopalakrishnan Brain Centre (SGBC), IITM, Chennai 600036, India

Alzheimer Readmission Prediction Magnetic Resonance Imaging

Open paper AskPDF arXiv

Generative modeling frameworks have emerged as an effective approach to capture high-dimensional image distributions from large datasets without requiring domain-specific knowledge, a capability essential for disease progression modeling. Recent generative approaches have attempted to capture progression by mapping images to a latent space and guiding representations to generate follow-up images from previous time points. However, these methods impose constraints on distribution learning, resulting in latent spaces with limited controllability for generating follow-up images without paired subject-specific longitudinal guidance. In order to enable controlled movements in the latent representational space and generate progression images from a previous time-point image without subject-specific guidance, we introduce a conditionable Diffusion Auto-encoder framework that forms a compact latent space capturing high-level semantics and providing means to control generation. Our approach leverages this latent space to condition and apply controlled shifts to the representations of previous time-point images by isolating progression and subject identity information for generating follow-up images. The shifts are implicitly guided by correlating with progression attributes and constraining to Alzheimer's disease specific regions, without paired longitudinal guidance. We validate the generations through image quality metrics, volumetric progression analysis, and downstream tasks in Alzheimer's disease datasets from different sources. This demonstrates the effectiveness of our approach for Alzheimer's progression modeling and longitudinal image generation.

Abstract

Explore similar work

AD-DAE: Alzheimer's Disease Progression Modeling with Unpaired Longitudinal MRI using Diffusion Auto-Encoders

CLIMB: Controllable Longitudinal Brain Image Generation using Mamba-based Latent Diffusion Model and Gaussian-aligned Autoencoder

Self-supervised Pre-training Helps Retinal Disease Progression Modelling Most When Data Is Scarce