Sheaf Neural Networks on SPD Manifolds: Second-Order Geometric Representation Learning

Graph neural networks face two fundamental challenges rooted in the linear structure of Euclidean vector spaces: (1) Current architectures represent geometry through vectors (directions, gradients), yet many tasks require matrix-valued representations that capture relationships between directions—such as how atomic orientations covary in a molecule. These second-order representations are naturally captured by points on the symmetric positive definite matrices (SPD) manifold; (2) Standard message passing applies shared transformations across edges. Sheaf neural networks address this via edge-specific transformations, but existing formulations remain confined to vector spaces and therefore cannot propagate matrix-valued features. We address both challenges by developing the first sheaf neural network operates natively on the SPD manifold. Our key insight is that the SPD manifold admits a Lie group structure, enabling well-posed analogs of sheaf operators without projecting to Euclidean space. Theoretically, we prove that SPD-valued sheaves are strictly more expressive than Euclidean sheaves: they admit consistent configurations (global sections) that vector-valued sheaves cannot represent, directly translating to richer learned representations. Empirically, our sheaf convolution transforms effectively rank-1 directional inputs into full-rank matrices encoding local geometric structure. Our dual-stream architecture achieves SOTA on 6/7 MoleculeNet benchmarks, with the sheaf framework providing consistent depth robustness.

Whether a drug works often depends not on which atoms a molecule contains, but on its three-dimensional shape: how its bonds are angled, and how rigid or flexible different regions are. Today's AI models for molecules describe geometry using arrows that point in directions. But capturing how directions relate to one another, the molecule's overall shape, needs something richer than an arrow. We taught an AI to represent each part of a molecule not as a single arrow, but as a compact "shape" that records how directions vary together. To pass these shapes between connected atoms, we extended a mathematical tool called a sheaf, previously limited to arrow-like data, so that it works directly in the curved space where these shapes naturally live, letting every chemical bond transform information in its own way. Our model reaches top performance on six of seven standard molecular benchmarks. Strikingly, it begins with plain directional inputs and automatically grows richer shape information layer by layer. It also stays accurate when stacked very deep—keeping 97% of its performance at 32 layers, where most comparable networks collapse.