SUBTRACE: JOINT AMORTIZED INFERENCE OF DISEASE PROGRESSION SUBTYPES AND TIMELINES FROM CROSS-SECTIONAL DATA

Event-based models infer disease progression from cross-sectional biomarkers, but subtype extensions such as SuStaIn and Bayesian EBM subtyping (bebms) are computationally expensive and return only ordinal event sequences. We propose Subtrace, a simulation-trained Transformer for amortized subtype-aware disease progression inference. Given a cohort, Subtrace predicts subtype-specific continuous biomarker timelines, patient subtype assignments, and patient stages in a single forward pass. It does so by learning per-patient, per-biomarker abnormality evidence and fitting subtype timelines to that evidence at the predicted disease stage. Across nine synthetic experiments, Subtrace improves ordering and staging over bebms and SuStaIn. Most importantly, Subtrace scales favorably with dimensionality: with more biomarkers, ordering improves and normalized stage error decreases. Applied to ADNI and NACC, independently trained models recover similar Alzheimer's disease subtype patterns consistent with prior research. These results support simulation-trained amortized inference as a practical alternative to per-dataset MCMC for subtype-aware disease progression modeling. Code is available at https://gitlab.com/4peerreview/iclr2027_subtrace.