Stochastic Siamese MAE Pretraining for Longitudinal Medical Images

TL;DR

STAMP introduces a stochastic Siamese MAE framework that effectively captures temporal uncertainty in longitudinal medical images, improving disease progression prediction over existing methods.

Contribution

It presents a novel stochastic approach to temporal encoding in MAE, reframing the loss as a conditional variational inference for better disease evolution modeling.

Findings

Outperforms existing temporal MAE methods on OCT and MRI datasets.

Pretrained ViT models achieve higher accuracy in disease progression prediction.

Effectively models non-deterministic disease dynamics.

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.