Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

BrainDyn: Sheaves Meet Neural ODEs for Brain Dynamics

11 minute read

Published:

TL;DR: Standard graph message passing assumes every node speaks the same language — aggregation is a weighted sum in one shared feature space. For brain data that assumption is wrong: distinct regions encode information in different representational geometries. BrainDyn gives each edge a pair of learnable restriction maps, so signals are transformed before they are compared, and embeds the resulting sheaf Laplacian inside a neural ODE so activity evolves in continuous time. It is, as far as the authors know, the first combination of cellular sheaves with neural ODEs.
Paper: BrainDyn: A Sheaf Neural ODE for Generative Brain Dynamics
Authors: Siddharth Viswanath, Panayiotis Ketonis, Chen Liu (Yale), Michael Perlmutter (Boise State), Dhananjay Bhaskar (Wisconsin–Madison), Smita Krishnaswamy (Yale)
Preprint: arXiv:2605.19324, May 2026  ·  code

Why a sheaf, for brains specifically

Most applications of sheaf neural networks are motivated by heterophily: neighbouring nodes have different labels, so averaging them is destructive, and restriction maps let you align before you compare. BrainDyn’s motivation is different and, I think, more physically grounded.

The argument runs: a graph neural network on parcellated brain regions treats every region’s feature vector as living in one common space, because that is what summation presupposes. But visual cortex and prefrontal cortex do not encode information in the same coordinates. Communication between them is a transformation, not a transfer. A restriction map \(\rho_{i \to e_{ij}}\) is precisely a model of that transformation, and it is learned per connection.

That reframes the sheaf not as a heterophily fix but as a claim about neural coding. Whether the claim is true is an empirical question the paper does not settle, but it is a better reason to reach for a sheaf than “our benchmark has low homophily”.

The architecture in three stages

Stage 1 — stalks are memories, not features

The move I find most interesting: the stalk at each node is not the region’s current activity but an LSTM summary of its recent history.

\[ h_i(t) = \mathrm{LSTM}\big(x_i(t_p : t_q)\big), \qquad h_i(t) \in \mathbb{R}^d, \]

stacked into \(H(t) = [h_1^{\top}(t), \dots, h_N^{\top}(t)]^{\top} \in \mathbb{R}^{N \times d}\).

So the sheaf lives over dynamical states, not over instantaneous signals. A stalk holds oscillatory phase, recent trend, and delay structure. Restriction maps then align those dynamical states across a connection. This is a genuinely different object from the node-feature sheaves in the rest of this book, and the ablation supports it carrying weight: removing the LSTM degrades fMRI MSE from 0.66 to 1.60, the largest single drop reported.

Stage 2 — attention-gated restriction maps

Each edge gets two learnable maps \(\rho_{i \to e_{ij}}, \rho_{j \to e_{ij}} \in \mathbb{R}^{d \times d}\), and the transported stalks are \(h_{i \to e_{ij}} = \rho_{i \to e_{ij}} h_i\). Scalar attention coefficients then modulate each side:

\[ \alpha_i = \sigma\big(a^{\top} h_{i \to e_{ij}}\big), \qquad \delta_{ij} = \alpha_i h_{i \to e_{ij}} - \alpha_j h_{j \to e_{ij}}, \]

with \(a \in \mathbb{R}^d\) shared across edges and \(\sigma\) the sigmoid. The Laplacian action pulls the discrepancy back to the node:

\[ (L_{\mathcal{F}}H)_i = \sum_{\{j \,:\, e_{ij} \in E\}} \rho_{i \to e_{ij}}^{\top}\, \delta_{ij}. \]

Message passing is then \(L\) rounds of \(H^{(l)} = (I - L_{\mathcal{F}})H^{(l-1)}\).

The gating breaks the Laplacian structure, and the paper does not flag it. The appeal of \(L_{\mathcal{F}} = \delta^{\top}\delta\) is entirely spectral: it is symmetric, positive semi-definite, and its kernel is the space of global sections. That is what every convergence and oversmoothing result about sheaf diffusion relies on. Insert \(\alpha_i\) inside \(\delta\) but pull back with the unmodulated \(\rho_{i \to e_{ij}}^{\top}\), as equation (7) does, and the operator is no longer \(\delta^{\top}\delta\) for any coboundary \(\delta\). Fold the gate into the map — define \(\tilde{\rho}_i = \alpha_i \rho_i\) — and \(\delta_{ij} = \tilde{\rho}_i h_i - \tilde{\rho}_j h_j\) is a genuine coboundary, but then the pullback ought to be \(\tilde{\rho}_i^{\top}\), not \(\rho_i^{\top}\). It is off by exactly the factor \(\alpha_i\). The text asserts that "globally, the sheaf Laplacian corresponds to the operator \(L_{\mathcal{F}} = \delta^{\top}\delta\)" immediately after defining the gated version. Both statements cannot hold. Nothing in the empirical results depends on this — but a reader coming from the theory chapters should not assume the spectral guarantees transfer.

A smaller slip in the same region: the paper writes the Laplacian blocks as \(L_{\mathcal{F}}(i,i) = \sum \rho^{\top}\rho\, x_i\) and \(L_{\mathcal{F}}(i,j) = -\rho_i^{\top}\rho_j\, x_i\). A block of an operator should not contain the signal it acts on; the \(x_i\) belongs in the action, which the very next equation writes correctly.

Stage 3 — continuous time

\[ \frac{dx}{dt} = f_\theta\big(h^{(L)}_i\big) = \mathrm{MLP}_\theta\big(h^{(L)}_i\big), \qquad \hat{x}_i(t) = \mathrm{ODESolve}\big(f_\theta, x_i(t_q), [t_q, t_r]\big). \]
A question about equation (9). As written, the vector field depends only on \(h^{(L)}_i\) — computed from the context window before integration begins — and not on the state \(x\) being integrated. A vector field that is constant in \(x\) integrates to a straight line: \(x(t) = x(t_q) + (t - t_q) f_\theta(h^{(L)}_i)\). That cannot be what produces the reported forecasts. The strongest internal evidence is the solver. The experimental setup specifies fixed-step fourth-order Runge–Kutta. On a state-independent field RK4 returns exactly the Euler answer at four times the cost, so nobody would choose it. The field must be state-dependent in the implementation; equation (9) is missing an argument.

Training combines reconstruction with two regularisers:

\[ \mathcal{L} = \mathcal{L}_{\mathrm{MSE}} + \lambda_1 \underbrace{\sum_{e_{ij} \in E} \lVert \delta_{ij} \rVert_1}_{\text{sparsity}} + \lambda_2 \underbrace{\sum_{e_{ij} \in E \cup E_{\mathcal{P}}} \Big\lvert\, \lVert \delta_{ij} \rVert_2 - \mathbf{1}_{\{e_{ij} \in E_{\mathcal{P}}\}} \Big\rvert}_{\text{prior agreement}} . \]

The prior graph \(\mathcal{P}\) comes from Granger causality computed on the context window only — never the forecast window, which is the right call for avoiding leakage. The second term pushes discrepancy magnitude towards \(1\) on prior edges and towards \(0\) elsewhere: a soft instruction to route information along anatomically or causally plausible connections.

Results

Three modalities: resting-state fMRI (PNC, 1188 subjects, Schaefer 400-region parcellation, TR = 3 s), scalp EEG (TUSZ, 315 patients, 19 channels at 200 Hz), and simulated spiking activity from NEST.

 fMRI MSEfMRI DTWEEG MSEEEG DTW
CNN-LSTM0.890.360.550.30
BIOT1.990.551.160.47
EvolveGCN1.000.310.650.29
ODEBRAIN0.850.300.470.27
BrainDyn0.660.250.440.23
— ablation: no sheaf1.260.400.760.34
— ablation: no LSTM1.600.431.000.41

The cross-modality point is the strongest one here. Most brain-dynamics models are tuned to either fMRI or EEG and degrade on the other; BrainDyn leads on both. And both ablations land below several baselines, which is the useful kind of ablation — neither component is carrying the model alone.

The choice of metric is also well argued. Appendix F shows a forecast that shares a linear trend with the observed signal but reproduces none of its oscillatory structure, scoring PCC 0.93 and SCC 0.94 while normalised DTW correctly reports 0.49. Correlation rewards a shared drift; DTW does not. More papers should include this figure.

The perturbation result is not separated from the baseline. On the NEST out-of-distribution test — train on unperturbed activity, evaluate on windows where one neuron has been silenced — BrainDyn reports MSE \(0.671 \pm 0.038\) against ODEBRAIN's \(0.702 \pm 0.019\). Those intervals are \([0.633, 0.709]\) and \([0.683, 0.721]\). They overlap across more than half of BrainDyn's range, and BrainDyn's standard deviation is twice the baseline's. The stated improvements — 4% MSE, 2% MAE, 1% DTW — are arithmetically correct on the means, and the means are genuinely the best in the table. But "BrainDyn achieves the lowest MSE and DTW distance among all methods" invites a stronger reading than five folds with these spreads can support. The forecasting results in Table 1 are a different matter: 0.66 against 0.85 with standard deviations of 0.01 and 0.11 is a real separation.

Two smaller things a careful reader will trip over. The main text says the models predict a 30-timepoint forecast from a 10-timepoint context; the appendix says each 40-timepoint window is 30 context and 10 horizon. Those are opposite. The EEG durations settle it in favour of neither: at 200 Hz, 30 timepoints is 0.15 s and 10 is 0.05 s, but the main text pairs the 30-point forecast with 0.05 s and the 10-point context with 0.15 s — the two figures are swapped. The fMRI durations (90 s and 30 s at TR = 3 s) are both correct, which makes the EEG pairing look like a transcription slip rather than a different convention.

Where this sits in the book

BrainDyn connects two threads that have run separately: the sheaf machinery of this book, and graph neural ODEs from Book II. The combination is natural in hindsight — sheaf diffusion is already written as a dynamical system, and NSD’s own theory is a statement about \(t \to \infty\) — so making \(t\) genuinely continuous rather than counting layers is the obvious next question.

What it does not do is inherit the theory. Once attention gates sit inside the coboundary and the vector field is an MLP, the operator is a learned interaction matrix that happens to be built out of restriction maps. That is not a criticism of the results; it is a note that the phrase “sheaf Laplacian” is doing architectural work here rather than spectral work. Compare DNSD, which makes the same trade far more explicitly.

The honest limitation the authors name is that this models resting-state dynamics — self-sustaining activity with little external drive. Stimulus-driven dynamics would need training regimes where many units are perturbed simultaneously, which is future work.

✅ Key Takeaways

  • The motivation is a claim about neural coding, not about heterophily: different brain regions use different representational geometries, so inter-region communication should be a learned transformation rather than a weighted sum.
  • Stalks hold LSTM summaries of recent history rather than instantaneous activity, so the sheaf is defined over dynamical states. Removing the LSTM is the single most damaging ablation (fMRI MSE 0.66 → 1.60).
  • Forecasting results are strong and cross-modal: best on both fMRI and EEG, where baselines are typically tuned to one.
  • The attention gates sit inside the coboundary but the pullback uses the ungated map, so the operator is not \(\delta^{\top}\delta\) — the spectral guarantees of sheaf theory do not transfer, and the paper's claim that they do is inconsistent with its own equation.
  • Equation (9)'s vector field has no dependence on the integrated state; taken literally it would integrate to a straight line. The use of RK4 implies the implementation differs from the written equation.
  • The out-of-distribution perturbation gains are within overlapping error bars. The in-distribution forecasting gains are not.
  • Appendix F's demonstration that correlation is inflated by a shared trend, while DTW is not, is worth borrowing for any trajectory-forecasting paper.

References

You May Also Enjoy

The Standard Library Modules Worth Knowing Before You pip install

8 minute read

Published:

Python ships with a counter, a queue, a memoiser, a CLI parser, a test framework’s smaller cousin, and a JSON-safe date library already installed — most of what people reach for a package to solve is one import away.

Files and Context Managers: Why with Is Not Optional

9 minute read

Published:

A file object that falls out of scope without with does get closed eventually — but ‘eventually’ means whenever the garbage collector gets around to it, which is not a promise any program handling more than a handful of files can live with.

Errors and Exceptions: EAFP, the Hierarchy, and Reading a Traceback

9 minute read

Published:

Exceptions in Python are not exceptional. They are a normal control-flow mechanism that the language leans on so heavily that the idiomatic style is to try the operation and handle the failure, rather than check first — which is faster, shorter, and free of race conditions.