braindyn-sheaf-neural-ode

name

braindyn-sheaf-neural-ode

description

BrainDyn: A Sheaf Neural ODE framework for modeling continuous-time neural dynamics on structured brain graphs. Combines cellular sheaf theory (learnable restriction maps, sheaf Laplacian) with neural ODEs for generative brain dynamics. Supports multi-modal neural data (fMRI, EEG, spiking neural network simulations), in silico perturbation prediction, and structured latent representations. Use when: (1) modeling brain dynamics as continuous-time ODEs on brain graphs, (2) implementing sheaf neural networks for heterogeneous inter-region communication, (3) building virtual brain models for perturbation analysis, (4) forecasting neural activity across modalities (fMRI/EEG/spiking), (5) studying sheaf Laplacians as generalization of graph Laplacians for brain networks. Activation: braindyn, sheaf neural ODE, sheaf neural network, brain dynamics forecasting, virtual brain model, in silico perturbation, cellular sheaf brain, sheaf Laplacian, generative brain dynamics, continuous-time neural dynamics.

BrainDyn: A Sheaf Neural ODE for Generative Brain Dynamics

Paper: arXiv:2605.19324 | Viswanath et al. (Yale, Boise State, UW-Madison) | May 2026

Overview

BrainDyn is a sheaf neural ordinary differential equation (neural ODE) model for continuous-time dynamics on structured brain graphs. First work to combine cellular sheaves with neural ODEs.

Key insight: Standard GNNs aggregate in a shared feature space (scalar summation), causing oversmoothing and erasing heterogeneous dynamics. Cellular sheaves equip each edge with learnable restriction maps that transform node features into edge-specific shared spaces, enabling brain regions to maintain distinct representational geometries while communicating.

Architecture

Three Core Components

[LSTM] → [Restriction Maps + Attention] → [Sheaf Laplacian Message Passing] → [Neural ODE]

1. Memory-based Node Stalks (Section 4.2)

   • LSTM encodes temporal window [x(tₚ), ..., x(tq)] into hidden state hᵢ(t) ∈ ℝᵈ
   • hᵢ(t) interpreted as stalk representation for node i
   • H(t) = [h₁ᵀ(t), ..., h_Nᵀ(t)]ᵀ stacks all node stalks

2. Learnable Sheaf Restriction Maps with Attention (Section 4.3)

   • For edge eᵢⱼ: learnable ρᵢ→ₑᵢⱼ, ρⱼ→ₑᵢⱼ ∈ ℝᵈˣᵈ project stalks to edge space
   • hᵢ→ₑᵢⱼ = ρᵢ→ₑᵢⱼ hᵢ (edge-specific transformation)
   • Attention: αᵢ = σ(aᵀ hᵢ→ₑᵢⱼ) modulates contribution
   • Edge discrepancy: δᵢⱼ = αᵢ hᵢ→ₑᵢⱼ − αⱼ hⱼ→ₑᵢⱼ
   • Sheaf Laplacian: (L_ℱ H)ᵢ = ∑_j ρᵢ→ₑᵢⱼᵀ δᵢⱼ

3. Continuous-time Neural ODE (Section 4.5)

   • dx/dt = f_θ(hᵢ⁽ᴸ⁾) = MLP_θ(hᵢ⁽ᴸ⁾) where h⁽ᴸ⁾ is post-message-passing state
   • Message passing: H⁽ˡ⁾ = (I − L_ℱ) H⁽ˡ⁻¹⁾ for L rounds
   • Solved via RK4 (torchdiffeq), step size Δt=1

Prior Graph Construction (Section 4.1)

• Granger causality from input context window only (no future leakage)
• Top-k strongest edges retained
• Used to initialize sheaf structure and regularize learned interactions

Training Objective

ℒ = ℒ_MSE + λ₁ ℒ_sparse + λ₂ ℒ_prior

• ℒ_MSE: reconstruction error
• ℒ_sparse: L1 penalty on edge discrepancy signals (parsimonious communication)
• ℒ_prior: learned interactions close to prior Granger graph

Evaluation Results

Multi-modal Forecasting

Method	fMRI MSE ↓	EEG MSE ↓	NEST MSE ↓
CNN-LSTM	0.89	0.55	0.896
BIOT	1.99	1.16	1.041
EvolveGCN	1.00	0.65	1.029
ODEBRAIN	0.85	0.47	0.702
RiTINI	1.91	1.02	0.904
BrainDyn	0.66	0.44	0.671

• fMRI (PNC, 1188 subjects, 400 regions): 22% better than ODEBRAIN
• EEG (TUSZ, 315 patients, 19 channels): best across all metrics
• NEST spiking (perturbation, 10K networks): best out-of-distribution generalization

In Silico Perturbation

• Models trained on unperturbed activity, evaluated on single-neuron silencing
• BrainDyn extrapolates to perturbed OOD dynamics; both sheaf + LSTM essential

Implementation Guide

Key Equations

Sheaf Laplacian block structure:

• Diagonal: ℒ_F(i,i) = ∑_j ρᵢ→ₑᵢⱼᵀ ρᵢ→ₑᵢⱼ
• Off-diagonal: L_F(i,j) = −ρᵢ→ₑᵢⱼᵀ ρⱼ→ₑᵢⱼ

Message passing update: H⁽ˡ⁾ = (I − L_ℱ) H⁽ˡ⁻¹⁾

Datasets

• PNC fMRI: Philadelphia Neurodevelopmental Cohort, 400 Schaefer regions, TR=3s
• TUSZ EEG: Temple University Hospital Seizure Corpus, 19 channels, 200 Hz
• NEST: 100 iaf_psc_alpha neurons, small-world graph (k=8, β=0.1, 400 edges)

Hyperparameters

• AdamW, lr=1e-3, weight_decay=1e-5, ReduceLROnPlateau
• ODE solver: RK4 (torchdiffeq), Δt=1
• LSTM: per-node, sliding window encoding
• Sheaf dimension d: learnable
• Batch size: 64, 5-fold CV, single H200 GPU

Dependencies

import torch
import torch.nn as nn
from torchdiffeq import odeint  # Neural ODE solver
# Graph construction: use Granger causality (statsmodels or custom)

Complexity

Per-sample: O(NT(FD+LD²)) + O(EDM) + O(4SHV²)

• N nodes, T time steps, F features, D stalk dimension, E edges
• L message passing rounds, M sheaf dimension, V MLP width, H horizon, S RK4 steps

Biological Interpretation

The sheaf framework encodes a biologically motivated hypothesis: communication between brain regions involves transformation, not just summation. Each restriction map ρᵢ→ₑᵢⱼ represents how region i's signal is transformed before entering the communication channel with region j — analogous to synaptic transformations, receptor-specific filtering, or modulatory gating.

The sheaf Laplacian measures disagreement after alignment — the brain's tendency to reduce representational mismatches across connected regions, driving synchronization/desynchronization patterns.

Limitations

• Currently models resting dynamics only (self-sustaining, minimal sensory input)
• Not yet tested on stimulus-driven dynamics (movie-watching, audio inputs)
• Requires simultaneous perturbation of multiple neuronal units for training on sensory tasks

Activation Keywords

• braindyn
• sheaf neural ode
• sheaf neural network
• cellular sheaf brain
• sheaf Laplacian
• brain dynamics forecasting
• virtual brain model
• in silico perturbation
• continuous-time neural dynamics
• generative brain dynamics