# Stochastic Optimal Control for Continuous-Time fMRI Representation Learning

arXiv: https://arxiv.org/abs/2502.04892 · HTML: https://mlsanity.com/paper/2502.04892/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 87 references that are in the map.
- First posted on arXiv 7 Feb 2025.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Captum: A unified and generic model interpretability library for PyTorch](https://mlsanity.com/paper/2009.07896/) (2020)
- [Revisiting Feature Prediction for Learning Visual Representations from Video](https://mlsanity.com/paper/2404.08471/) (2024)
- [Neural Controlled Differential Equations for Irregular Time Series](https://mlsanity.com/paper/2005.08926/) (2020)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [GRU-ODE-Bayes: Continuous modeling of sporadically-observed time series](https://mlsanity.com/paper/1905.12374/) (2019)
- [Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit](https://mlsanity.com/paper/1905.09883/) (2019)
- [On the Challenges and Perspectives of Foundation Models for Medical Image Analysis](https://mlsanity.com/paper/2306.05705/) (2023)
- [A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning](https://mlsanity.com/paper/1710.05741/) (2017)
- [Modeling Irregular Time Series with Continuous Recurrent Units](https://mlsanity.com/paper/2111.11344/) (2021)
- [Trivializations for Gradient-Based Optimization on Manifolds](https://mlsanity.com/paper/1909.09501/) (2019)
- [Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking](https://mlsanity.com/paper/2409.19407/) (2024)
- [BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning](https://mlsanity.com/paper/2403.01433/) (2024)
- [Deriving reproducible biomarkers from multi-site resting-state data: An Autism-based example](https://mlsanity.com/paper/1611.06066/) (2016)
- [Neural Continuous-Discrete State Space Models for Irregularly-Sampled Time Series](https://mlsanity.com/paper/2301.11308/) (2023)
- [Adaptive importance sampling for control and inference](https://mlsanity.com/paper/1505.01874/) (2015)
- [SDE Matching: Scalable and Simulation-Free Training of Latent Stochastic Differential Equations](https://mlsanity.com/paper/2502.02472/) (2025)
- [Identifying Latent Stochastic Differential Equations](https://mlsanity.com/paper/2007.06075/) (2020)
- [Latent SDEs on Homogeneous Spaces](https://mlsanity.com/paper/2306.16248/) (2023)
- [Amortized Control of Continuous State Space Feynman-Kac Model for Irregular Time Series](https://mlsanity.com/paper/2410.05602/) (2024)

## Built on by (most cited first, 1 of 1)

- [Mapping minds not averages: a scalable subject-specific manifold learning framework for neuroimaging data](https://mlsanity.com/paper/2505.00196/) (2025)

## Abstract (arXiv)

Learning robust representations from functional magnetic resonance imaging (fMRI) is fundamentally challenged by the temporal irregularity and noise inherent in data from heterogeneous sources. Existing self-supervised learning (SSL) methods often discard critical temporal information by discretizing or averaging fMRI signals. To address this, we introduce a novel framework that reframes SSL as a Stochastic Optimal Control (SOC) problem. Our approach models brain activity as continuous-time latent dynamics, learning a robust representation of brain dynamics by optimizing a control policy that is agnostic to the temporal irregularity. This SOC framework naturally unifies masked autoencoding (MAE) and joint-embedding prediction (JEPA) to extract compact, control-derived representations. Furthermore, a simulation-free inference strategy ensures computational efficiency and scalability for large-scale fMRI datasets. Our model demonstrates state-of-the-art performance across diverse downstream applications, highlighting the potential of the SOC-based continuous-time representation learning framework.