Source-linked AI summary
Block Diffusion for Longitudinal EHR: Accuracy, Efficiency, and Representation Trade-offs
Ashvin Gupta, Joshua Placidi, Brendan Delaney, Alessandra Russo
TL;DR
AR models perform strongly in zero-shot EHR prediction, but token-by-token generation limits parallelism for trajectories with temporally clustered events, and BD's performance relative to strong AR models remains less understood. This paper compares block diffusion with AR on MIMIC-IV under controlled conditions. Small-block BD retains much of AR's predictive performance with faster evaluation, while frozen probes are near-identical across tasks and rollout remains weaker in most settings.
Problem
AR models show strong zero-shot EHR prediction, but their sequential generation limits parallelism, and block diffusion's performance relative to strong AR models for zero-shot clinical prediction remains less understood.
Source: §1 · pp. 1–2Method
The paper conducts a controlled comparison of AR and block diffusion on tokenised MIMIC-IV trajectories, varying block size, denoising budget, and decoding strategy.
Source: §1 · p. 2, §2.1 · p. 2Results
Small-block BD retains much of AR's predictive performance while reaching up to a 1.43× speedup; frozen probes yield near-identical predictive performance across tasks, although rollout remains weaker in most settings.
Source: §1 · p. 2, §4 · pp. 3–4Takeaways & Limitations
Block diffusion is a promising computationally attractive framework for longitudinal EHR modelling, with improved decoding important for closing the residual gap to AR zero-shot prediction.
Source: §4 · p. 4Fixed block sizes may not align with variable-length clinically meaningful event groups, motivating variable-length diffusion generation.
Source: §4 · p. 4
Abstract
from arXiv · showhide
Autoregressive (AR) models have proved effective for zero-shot prediction from longitudinal electronic health records, but impose strictly sequential generation on trajectories containing temporally clustered and interdependent clinical events. We investigate block-based masked diffusion (BD) as an alternative on MIMIC-IV under matched tokenisation, context length, model capacity, and evaluation conditions. Across block sizes, BD exhibits a clear quality--efficiency trade-off, with small-block configurations retaining much of the predictive performance of AR while enabling substantially faster inference. A representative $B{=}8$ configuration achieves approximately $1.43\times$ faster evaluation. Although generative rollout remains below AR in most settings, frozen linear probes on BD and AR representations achieve near-identical performance across clinical prediction tasks, indicating that BD learns comparably informative clinical representations. Increasing the denoising budget yields diminishing returns, while alternative decoding strategies provide limited gains. Overall, our results show that block diffusion is a viable and computationally attractive approach to longitudinal EHR modelling, with further improvements in iterative decoding offering a promising route toward closing the remaining generative performance gap.
1 Introduction
Autoregressive models have strong zero-shot EHR prediction performance, but sequential token generation limits parallelism despite the clustered timing of many clinical events. This paper compares block diffusion with AR models and examines predictive performance, efficiency, representations, and decoding alternatives.
1 Introduction
- AR models perform strongly on zero-shot clinical prediction, but their strict left-to-right generation limits parallelism for temporally clustered EHR events.In MIMIC-IV, 47.4% of recorded timestamps contain at least eight tokens.Source: §1 · p. 1
- Block diffusion denoises tokens jointly within blocks while retaining causal conditioning across previously generated blocks.Prior work describes a quality–parallelism trade-off: larger blocks allow more parallel generation but generally degrade likelihood and generation quality.Source: §1 · p. 1
- The comparison tests how block size and denoising budget affect predictive performance and efficiency, whether rollout gaps reflect representations or generation, and whether decoding alternatives recover performance.Source: §1 · p. 2
- Block diffusion achieves rollout performance close to AR while providing up to a 1.43× evaluation speedup in the primary operating regime.Source: §1 · p. 2
- Frozen linear probes achieve near-identical downstream performance across tasks for AR and block-diffusion representations, despite a larger generative-rollout gap.Source: §1 · p. 2
- Alternative decoding strategies are not sufficient to close the remaining performance gap to AR.The paper also studies how block size and denoising budget affect the quality–efficiency trade-off.Source: §1 · p. 2
2 Methods
The methods compare capacity-matched AR and block-diffusion models on tokenised MIMIC-IV trajectories. Experiments vary block size, denoising budget, and decoding strategy, and assess both rollout predictions and frozen representations.
2.1 Models and generative prediction
- AR uses causal self-attention and next-token cross-entropy, while BD applies masked diffusion within blocks and causal conditioning across blocks.BD samples tokens into masked states according to a time-dependent retention probability.Source: §2.1 · p. 2
- BD trains by predicting original tokens in a corrupted active block conditioned on clean preceding blocks, with noise levels sampled independently across blocks.Source: §2.1 · p. 2
- At inference, AR generates token by token while BD iteratively denoises each fully masked block before committing it to the KV cache; prediction scores use 20 rollouts per patient.The score is the fraction of rollouts in which the positive endpoint occurs first; AUROC, AUPRC, and wall-clock time are reported.Source: §2.1 · p. 2
2.2 Data and experiments
- The study evaluates ICU mortality, ICU readmission, ICU admission, and hospital mortality using matched tokenisation, data splits, 1024-token context, model capacity, and optimizer schedule.Source: §2.2 · pp. 2–3
- With S = B/2, the primary block-size comparison fixes the denoising budget at 0.5 denoiser evaluations per generated token; fixed-block sweeps at B = 4, 8, and 128 test additional refinement.The sweeps assess whether more reverse steps recover performance lost at larger block sizes.Source: §2.2 · p. 3
- At B = 8, frozen final-token representations from AR and BD are tested with linear classifiers across all four outcomes to assess encoded predictive information without generative decoding.Mean pooling yielded lower performance for both models.Source: §2.2 · p. 3
- At B = 8 and S = 4, the study compares ancestral random unmasking with confidence-based unmasking and ReMDM-style remasking.Confidence-based unmasking keeps the expected reveal budget; remasking permits already revealed positions to be reconsidered.Source: §2.2 · p. 3
3 Results
Block diffusion is competitive across four clinical tasks, with small blocks offering the strongest quality–efficiency trade-off. AR remains stronger in most rollout outcomes, while frozen probes show near-identical predictive performance across tasks; denoising gains saturate and decoding alternatives offer limited improvements.
3 Results
- BD is competitive across four downstream tasks, but AR is stronger on three of four outcomes; frozen linear probes achieve near-identical AUROC across tasks.Hospital mortality is the exception in rollout, where increased AUROC comes at the cost of AUPRC; the remaining gap largely emerges during rollout rather than representation learning.Source: §3 · p. 3
- B = 8, S = 4 remained approximately 0.03 AUROC below AR while achieving a 1.43× evaluation speedup.Performance generally declines at larger block sizes, and very large blocks lose the latency advantage as each denoising step becomes more costly.Source: §3 · p. 3
- Additional denoising steps initially improve predictive performance, but gains saturate beyond a moderate budget while evaluation cost continues to rise.ReMDM-style revision yields small metric-dependent changes, while confidence-based unmasking degrades AUROC.Source: §3 · p. 3
4 Discussion
Small-block diffusion retains much of AR’s predictive performance with faster evaluation, but larger blocks expose a quality–efficiency trade-off. Denoising and decoding changes have limited ability to close the rollout gap, while fixed block boundaries and other EHR-specific structures motivate future work.
4 Discussion
- Up to 1.43× speedup accompanies small-block diffusion’s retention of much of AR’s predictive performance, while latency follows a U-shaped curve as block size grows.The figure compares quality and evaluation efficiency across block sizes; larger blocks reduce sequential depth but increase per-step compute.Source: §4 · p. 3, Figure 1 · p. 4
- Denoising gains plateau, and near-identical frozen-probe performance suggests the remaining rollout gap arises largely during generation rather than representation learning.The paper reports diminishing returns from added denoising and comparable predictive information in BD and AR representations.Source: §4 · p. 4
- Confidence-based unmasking degrades performance, while ReMDM-style revision provides no meaningful gain toward closing the AR gap.These ablations suggest that revising token commitment alone is unlikely to resolve the remaining difference.Source: §4 · p. 4
- Fixed-size blocks may not align with clinically meaningful event groups, whose sizes vary across episodes, motivating variable-length diffusion generation.The discussion presents variable-length generation as a future direction.Source: §4 · p. 4
- EHR-specific diffusion could generate events between successive time markers jointly, treating co-occurring events as unordered or partially ordered sets.The discussion also suggests alternative corruption processes that distinguish corrupted clinical events from plausible clean ones.Source: §4 · p. 4
- Generative endpoint prediction is only one possible EHR use; diffusion could also support infilling, record reconstruction, and counterfactual trajectory exploration.The paper concludes that improved decoding remains important for closing the residual gap to AR zero-shot prediction.Source: §4 · p. 4
A.1 Responsible Use
The study evaluates retrospective prediction on MIMIC-IV and does not establish that block diffusion is suitable for clinical deployment.
A.1 Responsible Use
- Retrospective MIMIC-IV prediction does not establish clinical deployment suitability.The authors limit the work to research use pending external validation, subgroup fairness assessment, calibration analysis, prospective evaluation, and clinical oversight.Source: §A.1 · p. 7
- Performance may not generalise across institutions, patient populations, or care settings.The experiments are limited to MIMIC-IV.Source: §A.1 · p. 7
- Generative predictions may involve uncertainty, miscalibration, and clinically consequential errors.The authors state that outputs should not guide patient care without validation and human supervision.Source: §A.1 · p. 7
A.3 Decoding with AR vs BD
The illustration contrasts sequential AR generation with block diffusion’s joint denoising of multiple positions while preserving causal generation across blocks.
A.3 Decoding with AR vs BD
- AR generates each event sequentially, whereas BD jointly denoises multiple positions within an active block while preserving causal generation across blocks.Source: Figure 2 · p. 7
- The illustration frames temporally clustered clinical events as the context for comparing the two decoding schemes.Source: Figure 2 · p. 7
- BD retains causal generation across blocks even as it processes positions jointly within each active block.Source: Figure 2 · p. 7
A.4 MIMIC-IV tokens
The figure describes the distribution of how many nontime-gap ETHOS event tokens share a recorded timestamp in MIMIC-IV.
A.4 MIMIC-IV tokens
- The figure plots the survival distribution of nontime-gap event-token counts sharing a recorded timestamp in MIMIC-IV.Source: Figure 3 · p. 7
- The count concerns event tokens in the ETHOS representation, excluding time-gap tokens.Source: Figure 3 · p. 7
- The plotted quantity is the number of tokens sharing a timestamp, rather than the number of distinct timestamps.Source: Figure 3 · p. 7
A.5 Block diffusion objective and sampling
Block diffusion trains on corrupted active blocks conditioned on clean preceding blocks, then generates by iteratively denoising each masked block before committing it to the cache.
A.5 Block diffusion objective and sampling
- Noise is sampled independently across blocks, and the denoiser reconstructs the corrupted active block while conditioning on clean preceding blocks.The training objective averages the masked-language-modelling loss across blocks.Source: §A.5 · p. 8
- At inference, each next block starts fully masked and is denoised over reverse steps before remaining masks are resolved and the block is committed to the KV cache.Generation then proceeds to the next block.Source: §A.5 · p. 8
- For confidence-based decoding, positions with the highest sampled-token probabilities are committed within an approximately fixed reveal budget; ReMDM-style revision can remask revealed positions for reconsideration.Source: §A.5 · p. 8
- S/B is the nominal number of sequential denoiser evaluations per generated token; the primary block-size comparison fixes this budget at 0.5.This normalizes sequential decoding depth, not exact compute, because each forward pass's wall-clock cost depends on active block size.Source: §A.5 · p. 9
A.6 Hospital-mortality AUROC versus AUPRC
For hospital mortality, block diffusion's AUROC and AUPRC results diverge, while frozen probes are nearly identical. The supplied figures and table captions also describe speed, AUPRC, and decoding comparisons.
A.6 Hospital-mortality AUROC versus AUPRC
- BD (B=8, S=4) exceeds AR hospital-mortality AUROC by ∆=+0.039 but trails AUPRC by ∆=−0.102; frozen probes are near-identical (AR 0.936, BD 0.933).The score analysis links the AUROC/AUPRC divergence to generative scoring rather than representation quality.Source: §A.6 · p. 9, Table 2 · p. 9
A.7 Quality–efficiency operating points
- On ICU mortality, the B=8, S=4 operating point achieves approximately 1.43× faster evaluation while retaining much of AR's predictive performance.Figure 4 plots evaluation speedup relative to AR against AUROC.Source: Figure 4 · p. 9
A.8 AUPRC results
- Figure 5 compares rollout and frozen-probe AUPRC differences across four tasks, ICU mortality AUPRC across block sizes at S/B=0.5, and AUPRC as denoising budget increases.The caption identifies 95% bootstrap confidence intervals.Source: Figure 5 · p. 10
A.9 Decoding Interventions
- Table 3 reports ICU-mortality decoding strategies for B=8 and S=4, with medians over completed full-validation runs where available and confidence from the available full-split run.Its brackets are 95% confidence intervals over three random seeds.Source: Table 3 · p. 10