Flow Matching with In-Context Priors for Out-of-Distribution Brain Dynamics
For neuroscientists, this enables counterfactual in-silico experiments to design and evaluate novel cognitive tasks before empirical validation, addressing the bottleneck of limited task coverage in fMRI data.
This work introduces a per-timestep conditioned diffusion transformer for generating realistic fMRI brain dynamics during unseen cognitive tasks, enabling zero-shot generalization to novel tasks. The model recovers region-specific activation patterns from language alone and improves with spatial priors, achieving the first generative model of whole-cortex fMRI for unseen tasks.
Flow matching and diffusion models enable conditional generation across domains ranging from images to proteins, with recent extensions to out-of-distribution contexts. Yet generative models of neural time series have largely remained restricted to categorical conditioning, precluding compositional and zero-shot generalization. In this work, we propose a per-timestep conditioned diffusion transformer for generating realistic fMRI brain dynamics during unseen cognitive tasks by injecting both compositional language and optional spatial priors in-context. Such zero-shot generation could enable counterfactual neuroscience by supporting in-silico design and evaluation of novel cognitive experiments before empirical validation. Leveraging this model, we evaluate across hundreds of held-out task conditions and characterize predictive performance in relation to the training manifold. From language alone, the model recovers region-specific recruitment across tasks and held-out spatial activation patterns. Spatial priors, when available, complement the text pathway by anchoring generation in regions of task space where language alone degrades, while retaining the compositional structure needed for counterfactual task specification. To our knowledge this is the first generative model of whole-cortex fMRI dynamics for unseen cognitive tasks, advancing counterfactual neuroscience and data-driven experimental design.