AS-TIME: Time-Conditioned Multimodal Modeling of Aortic Stenosis Progression

Reference implementation of AS-TIME, a time-conditioned multimodal framework for patient-specific aortic stenosis (AS) progression prediction from a single baseline echocardiography study and the elapsed time to follow-up (Δt).

Given the multi-view echo videos and the clinical report from a baseline study at t0, together with the interval Δt to the follow-up study t1, AS-TIME predicts the AS severity class at t1.

Overview of AS-TIME. Multi-view echo videos and their report are encoded with frozen backbones and aggregated via token-level MIL. A time-conditioned gate fuses the modalities using Δt, followed by an MoE head that produces the prediction.

Model

AS-TIME (ASTIME in src/model.py) is a hierarchical, Δt-conditioned fusion model:

1. Modality-specific feature extraction — each echo video is encoded by a frozen EchoPrime backbone (512-d) and each report sentence by a frozen BioClinical-ModernBERT encoder (768-d). Both are linearly projected to a shared d_model. Δt is log-transformed, standardized, and turned into a sinusoidal time embedding (src/time_embedding.py).
2. Token-level MIL — attention-based multiple-instance learning aggregates the variable-length video and report tokens into per-modality study embeddings.
3. Δt-conditioned gated fusion — a gate conditioned on the time embedding and the video/report cosine similarity produces a video weight w ∈ [0, 1]; the fused representation is w·h_video + (1−w)·h_report.
4. Δt-guided mixture-of-experts — an MoE head, with its gate conditioned on Δt, routes the fused representation across experts to model heterogeneous progression dynamics.

The precomputed EchoPrime / BioClinical-ModernBERT embeddings are consumed directly by the dataloader, so the frozen backbones are not part of this repository.

Baselines (Table 1)

All methods are selected through the ASTimeModel(...) factory in src/model.py via modality / fusion flags:

Method	Class	Flags
AS-TIME (ours)	ASTIME	--use_video --use_report --use_time --use_tri_top
Cross-Attention	ASTimeModel_CrossAttnV2	--use_video --use_report --use_time
Concatenation	ASMLP	--use_video --use_report --use_time --use_bimodal_mlp
EchoPrime + PosEmb	ASTimeModel_CrossAttn	--use_video --use_time
BERT / Report-only	ASReportOnlyModel	--use_report
Video-only	ASVideoOnlyModel	--use_video
Positional Embedding (time-only)	ASTimeOnlyModel	--use_time

The TAMME and MedGemma baselines from the paper are external and are not included here.

Repository layout

src/
  model.py            # AS-TIME + baseline models and the ASTimeModel factory
  model_training.py   # training / evaluation loop (entry point)
  dataloader.py       # paired-study dataset over precomputed embeddings
  time_embedding.py   # sinusoidal Δt encoder
  utils.py            # seeding, label collapsing, metrics, plotting helpers
script/
  run.sh              # train AS-TIME over 5 seeds (+ baseline flag reference)
  sweep.yaml          # Weights & Biases hyperparameter sweep

Data

The dataloader expects a CSV of paired studies (one row per t0 → t1 pair) and two .npy dictionaries of precomputed embeddings. Required CSV columns:

• split — train / val / test
• delta_t — interval (years) between the two studies
• stream_id, predicted_cardiac_view_id — per-video lists for the baseline study
• report_full_0_cls — baseline report sentence embeddings
• as_label_1 — target severity label at follow-up
• exam_id_0, exam_id_1, as_label_0 — identifiers / metadata

Point the code at your data with environment variables (or the corresponding flags):

export AS_TIME_VIDEO_EMB=data/as_time_embeddings_video.npy
export AS_TIME_REPORT_EMB=data/as_time_embeddings_report_cls.npy
export CSV_PATH=data/as_time_pairs.csv

Usage

pip install -r requirements.txt

# Train AS-TIME (5 seeds, as in the paper)
bash script/run.sh

# Or a single run directly
python src/model_training.py \
    --csv_path "$CSV_PATH" --class_num 2 \
    --use_video --use_report --use_time --use_tri_top \
    --n_experts 12 --epochs 50 --batch_size 16 \
    --lr 3.8e-6 --weight_decay 1e-3 --d_model 256

Pass --use_wandb to log to Weights & Biases.

Explainability

After a forward pass, ASTIME.get_last_explainability() returns the multi-level interpretability signals highlighted in the paper: per-instance MIL attention (video_attn, report_attn), the modality gating weights (video_weight, report_weight), and the MoE expert routing probabilities (moe_weight).

Temporal ablations

The Δt ablations from the paper (Table 2) are controlled by ASTIME constructor flags:

• zero_time_signal=True — "no temporal signaling": Δt is zeroed everywhere.
• condition_moe_on_time=False — keep Δt in gated fusion but not in expert routing.

Acknowledgements

This research was supported in part by the Canadian Institutes of Health Research (CIHR) and the Natural Sciences and Engineering Research Council of Canada (NSERC), and through computational resources and services provided by Advanced Research Computing at the University of British Columbia.