Research map / Time series / Foundation models for time series
Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Brain-JEPA applies a joint-embedding predictive architecture to pretrain a ViT on UK Biobank fMRI, using gradient-based ROI positional embeddings and spatiotemporal block masking to beat BrainLM and other fMRI baselines on downstream tasks.
Problem
Existing fMRI foundation models (e.g., BrainLM, MAE-based) reconstruct noisy, sparse fMRI time series, perform poorly in linear probing, lack a meaningful spatial order for ROI positional encoding, and are evaluated mainly on Caucasian cohorts.
Approach
Predict latent representations of masked fMRI target blocks (JEPA) rather than reconstructing them, with positional embeddings derived from functional connectivity gradients (Brain Gradient Positioning) and a three-region spatiotemporal masking scheme (cross-ROI, cross-time, double-cross) that injects inductive bias.
Key results
- Brain-JEPA (ViT-B) reaches 0.844 Pearson correlation for age and 81.52% sex accuracy on HCP-Aging, versus 0.832 and 74.39% for BrainLM.
- On the MACC Asian NC/MCI task, Brain-JEPA achieves 65.98% accuracy versus 61.65% for BrainLM, despite pretraining only on Caucasian UK Biobank data.
- Brain-JEPA reaches the peak performance of the 300-epoch ablation with only 50-200 pretraining epochs, and ViT-L performs best among model sizes.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use Brain-JEPA instead of BrainLM | UKB held-out age prediction · UKB held-out, age | MSE: 0.612 (0.041) -> 0.501* (0.034) | Helped |
| use Brain-JEPA instead of BrainLM | HCP-Aging sex prediction · HCP-Aging, sex | accuracy (%): 74.39 (1.55) -> 81.52* (1.03) | Helped |
| use Brain-JEPA instead of BrainLM | HCP-Aging Flanker prediction · HCP-Aging, Flanker | MSE: 0.971 (.054) -> 0.972 (.038) | No clear effect |
| use Brain-JEPA instead of BrainLM | ADNI NC/MCI classification · ADNI, NC/MCI | F1 (%): 85.66 (1.27) -> 86.32 (0.54) | Helped |
| use Brain-JEPA instead of BrainLM | MACC NC/MCI classification · MACC, NC/MCI (Asian) | accuracy (%): 61.65 (3.35) -> 65.98* (2.84) | Helped |
| pretrain with more UKB data (ablation) | Brain-JEPA · HCP-Aging age | Pearson correlation (ρ): 0.659 (.043); 0.768 (.012); 0.813 (.015); 0.844 (.030) | Helped |
| replace JEPA with MAE while retaining Brain-JEPA contributions (ablation) | Brain-JEPA · HCP-Aging sex | accuracy (%): 81.52 (1.03) -> 76.36 (2.58) | Hurt |
| use 30 gradient components instead of 3 (ablation) | Brain-JEPA · HCP-Aging sex | accuracy (%): 76.96 (1.77) -> 81.52 (1.03) | Helped |
| replace anatomical spatial positions with Brain Gradient Positioning (ablation) | Brain-JEPA · UKB age and sex; HCP-Aging Neuroticism and Flanker; ADNI, MACC, OASIS-3 and CamCAN tasks | task-specific metrics | Mixed |
| replace standard multi-block sampling with Spatiotemporal Masking (ablation) | Brain-JEPA · downstream tasks | performance | Helped |
| use a larger backbone (ViT-S to ViT-B to ViT-L) (ablation) | Brain-JEPA · downstream tasks | accuracy/correlation | Helped |
Surprising findings
- Brain-JEPA with anatomical-location embeddings instead of gradient positioning underperforms BrainLM on some tasks, indicating gradient positioning is needed for certain tasks.
- Using 3-dim gradients performs substantially worse than 30-dim gradients on HCP-Aging and ADNI.
Limitations
- Only ViT-S, ViT-B and ViT-L tested; larger models such as ViT-H not explored due to compute.
- Pretraining on a single Caucasian-dominated cohort (UK Biobank); limited dataset diversity.
- Interpretation via attention is preliminary, and the authors note that gradient positioning may be needed for some tasks where BrainLM performs better.
Experimental setup
Datasets: UK Biobank HCP-Aging ADNI MACC OASIS-3 CamCAN
Benchmarks: UK Biobank 20% held-out (age, sex) HCP-Aging (age, sex, neuroticism, flanker) ADNI (NC/MCI, amyloid +/-) MACC Asian NC/MCI OASIS-3 AD conversion CamCAN depression
Metrics: MSE Pearson correlation accuracy F1
Baselines: BrainLM BrainNetCNN BrainGNN BNT BrainMass CSM SwiFT SVM/SVR
Builds on: I-JEPA BrainLM Vision Transformer FlashAttention diffusion map
Closest papers by meaning
- Brain Harmony: A Multimodal Foundation Model Unifying Morphology and Function into 1D Tokens (NeurIPS 2025 · similarity 0.84)
- Dynamical Modeling of Behaviorally Relevant Spatiotemporal Patterns in Neural Imaging Data (ICML 2025 · similarity 0.84)
- FAST-Brain: A Flow-Aligned Spatio-Temporal Surrogate Brain Model (NeurIPS 2026 · similarity 0.84)
- How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks (NeurIPS 2024 · similarity 0.84)
- Connecting Joint-Embedding Predictive Architecture with Contrastive Self-supervised Learning (NeurIPS 2024 · similarity 0.84)
- BrainLM: A foundation model for brain activity recordings (ICLR 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.