# Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking

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

## Key facts

- Cited by 65 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 62 references that are in the map.
- First posted on arXiv 28 Sep 2024.

## Builds on (most cited first, 20 of 20)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Convolutional Sequence to Sequence Learning](https://mlsanity.com/paper/1705.03122/) (2017)
- [Brain Network Transformer](https://mlsanity.com/paper/2210.06681/) (2022)
- [BrainGB: A Benchmark for Brain Network Analysis with Graph Neural Networks](https://mlsanity.com/paper/2204.07054/) (2022)
- [Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data](https://mlsanity.com/paper/2206.11417/) (2022)
- [SwiFT: Swin 4D fMRI Transformer](https://mlsanity.com/paper/2307.05916/) (2023)
- [BrainMass: Advancing Brain Network Analysis for Diagnosis with Large-scale Self-Supervised Learning](https://mlsanity.com/paper/2403.01433/) (2024)
- [Beyond the Snapshot: Brain Tokenized Graph Transformer for Longitudinal Brain Functional Connectome Embedding](https://mlsanity.com/paper/2307.00858/) (2023)
- [SupReMix: Supervised Contrastive Learning for Medical Imaging Regression with Mixup](https://mlsanity.com/paper/2309.16633/) (2023)
- [Prompt Your Brain: Scaffold Prompt Tuning for Efficient Adaptation of fMRI Pre-trained Model](https://mlsanity.com/paper/2408.10567/) (2024)

## Built on by (most cited first, 30 of 65)

- [LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels](https://mlsanity.com/paper/2603.19312/) (2026)
- [Brain Foundation Models: A Survey on Advancements in Neural Signal Processing and Brain Discovery](https://mlsanity.com/paper/2503.00580/) (2025)
- [NeuralBench: A Unifying Framework to Benchmark NeuroAI Models](https://mlsanity.com/paper/2605.08495/) (2026)
- [Towards a general-purpose foundation model for fMRI analysis](https://mlsanity.com/paper/2506.11167/) (2025)
- [3D Foundation Model for Generalizable Disease Detection in Head Computed Tomography](https://mlsanity.com/paper/2502.02779/) (2025)
- [Brain-DiT: A Universal Multi-state fMRI Foundation Model with Metadata-Conditioned Pretraining](https://mlsanity.com/paper/2604.12683/) (2026)
- [Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models](https://mlsanity.com/paper/2605.31111/) (2026)
- [SLIM-Brain: A Data- and Training-Efficient Foundation Model for fMRI Data Analysis](https://mlsanity.com/paper/2512.21881/) (2025)
- [From Images to Signals: Are Large Vision Models Useful for Time Series Analysis?](https://mlsanity.com/paper/2505.24030/) (2025)
- [ECHO: Toward Contextual Seq2Seq Paradigms in Large EEG Models](https://mlsanity.com/paper/2509.22556/) (2025)
- [fMRI-LM: Towards a Universal Foundation Model for Language-Aligned fMRI Understanding](https://mlsanity.com/paper/2511.21760/) (2025)
- [BaRISTA: Brain Scale Informed Spatiotemporal Representation of Human Intracranial Neural Activity](https://mlsanity.com/paper/2512.12135/) (2025)
- [CGM-JEPA: Learning Consistent Continuous Glucose Monitor Representations via Predictive Self-Supervised Pretraining](https://mlsanity.com/paper/2605.00933/) (2026)
- [FlexiBrain: Resolution-Agnostic Voxel-Level Encoding for Native fMRI](https://mlsanity.com/paper/2606.11500/) (2026)
- [Voxel-Level Brain States Prediction Using Swin Transformer](https://mlsanity.com/paper/2506.11455/) (2025)
- [Scaling Vision Transformers for Functional MRI with Flat Maps](https://mlsanity.com/paper/2510.13768/) (2025)
- [BrainWorld: A Structural-Prior-Conditioned Generative Model for Whole-Brain 4D fMRI Dynamics](https://mlsanity.com/paper/2606.17742/) (2026)
- [Large Connectome Model: An fMRI Foundation Model of Brain Connectomes Empowered by Brain-Environment Interaction in Multitask Learning Landscape](https://mlsanity.com/paper/2510.18910/) (2025)
- [CrossJEPA: Cross-Modal Joint-Embedding Predictive Architecture for Efficient 3D Representation Learning from 2D Images](https://mlsanity.com/paper/2511.18424/) (2025)
- [Brain-Semantoks: Learning Semantic Tokens of Brain Dynamics with a Self-Distilled Foundation Model](https://mlsanity.com/paper/2512.11582/) (2025)
- [Medical world models: representing medical states, modelling clinical dynamics and guiding intervention policies](https://mlsanity.com/paper/2606.16721/) (2026)
- [A Brain Graph Foundation Model: Pre-Training and Prompt-Tuning across Broad Atlases and Disorders](https://mlsanity.com/paper/2506.02044/) (2025)
- [Self-Distillation of Hidden Layers for Self-Supervised Representation Learning](https://mlsanity.com/paper/2603.15553/) (2026)
- [CytoNet: A Foundation Model for the Human Cerebral Cortex at Cellular Resolution](https://mlsanity.com/paper/2511.01870/) (2025)
- [Brain-OF: An Omnifunctional Foundation Model for fMRI, EEG and MEG](https://mlsanity.com/paper/2602.23410/) (2026)
- [Stochastic Optimal Control for Continuous-Time fMRI Representation Learning](https://mlsanity.com/paper/2502.04892/) (2025)
- [BrainStratify: Coarse-to-Fine Disentanglement of Intracranial Neural Dynamics](https://mlsanity.com/paper/2505.20480/) (2025)
- [Bridging Brain with Foundation Models through Self-Supervised Learning](https://mlsanity.com/paper/2506.16009/) (2025)
- [Introducing Multimodal Paradigm for Learning Sleep Staging PSG via General-Purpose Model](https://mlsanity.com/paper/2509.22810/) (2025)
- [Region-Aware Reconstruction Strategy for Pre-training fMRI Foundation Model](https://mlsanity.com/paper/2511.00443/) (2025)

## Abstract (arXiv)

We introduce Brain-JEPA, a brain dynamics foundation model with the Joint-Embedding Predictive Architecture (JEPA). This pioneering model achieves state-of-the-art performance in demographic prediction, disease diagnosis/prognosis, and trait prediction through fine-tuning. Furthermore, it excels in off-the-shelf evaluations (e.g., linear probing) and demonstrates superior generalizability across different ethnic groups, surpassing the previous large model for brain activity significantly. Brain-JEPA incorporates two innovative techniques: Brain Gradient Positioning and Spatiotemporal Masking. Brain Gradient Positioning introduces a functional coordinate system for brain functional parcellation, enhancing the positional encoding of different Regions of Interest (ROIs). Spatiotemporal Masking, tailored to the unique characteristics of fMRI data, addresses the challenge of heterogeneous time-series patches. These methodologies enhance model performance and advance our understanding of the neural circuits underlying cognition. Overall, Brain-JEPA is paving the way to address pivotal questions of building brain functional coordinate system and masking brain activity at the AI-neuroscience interface, and setting a potentially new paradigm in brain activity analysis through downstream adaptation.