Joint Diffusion and Rotation Invariance: Sheaves That Learn to Lie
Published:
Authors: Ferran Hernandez Caralt, Guillermo Bernárdez Gil, Iulia Duta, Pietro Liò, Eduard Alarcón Cot
Venue: ICML 2024 Workshop on Geometry-grounded Representation Learning and Generative Modeling · arXiv:2407.20597
Universal approximation is the problem, not the solution
NSD proves that with enough capacity and diverse enough features, \(\Phi = \mathrm{MLP}(x_u \Vert x_v)\) can learn any sheaf on the graph. That is presented as a strength, and in one sense it is.
But this paper’s Remark 2.9 turns it around: “this choice does not introduce any inductive bias to ensure that learnt maps are able to deal with heterophily and oversmoothing.” A hypothesis class that contains everything prefers nothing. The theory in NSD says the right sheaf exists; the MLP gives no reason to expect gradient descent to find it, especially on small graphs.
There is a second, harder problem. Convergence of sheaf diffusion to a non-trivial point is only guaranteed for orthogonal restriction maps — otherwise \(H^0(G;\mathcal{F})\) may be \(\{0\}\), and diffusion oversmooths just as thoroughly as a GCN, only with more parameters.
Learning to lie
The opinion-dynamics reading gives the fix. Recall the framing: \(x_u\) is node \(u\)’s private opinion, and \(\mathcal{F}_{u \trianglelefteq e}x_u\) is the public opinion it expresses on edge \(e\).
Hansen and Ghrist’s Learning to Lie dynamic evolves the maps instead of the features:
This is the dual of sheaf diffusion — the features now act as the maps. Writing \(\mathcal{F}^*\) for the matrix of transposed restriction maps, it becomes \(\frac{d}{dt}\mathcal{F}^* = -\Delta_X \mathcal{F}^*\), structurally identical to \(\dot{X} = -\Delta_{\mathcal{F}}X\) with the roles swapped. It reaches agreement by adjusting expression rather than belief — but now nobody’s private opinion can change at all, which is equally unrealistic.
Doing both at once
Joint opinion-expression diffusion runs the two together as a coupled nonlinear system:
with \(\alpha, \beta\) setting which side of the conversation moves faster. The theory that comes with it is the reason to care:
- Lemma 1. \(\Psi(X,\mathcal{F}) = X^{\top}\delta^{\top}\delta X\) is non-negative and non-increasing along solutions, reaching zero exactly when \(\mathcal{F}_{u\trianglelefteq e}x_u = \mathcal{F}_{v\trianglelefteq e}x_v\) everywhere. The joint process is gradient descent on the sheaf Dirichlet energy over both variables.
- Theorem 3.2. If one diagonal block of \(\alpha\delta_0^{\top}\delta_0 - \beta x_0x_0^{\top}\) fails to be semidefinite, the trajectory converges to \((x_\infty, \delta_\infty)\) with \(x_\infty \neq 0\).
- Corollary 3.3. For any initial conditions there exists a scalar \(k\) such that starting from \((kx_0, \delta_0)\) gives \(x_\infty \neq 0\).
Theorem 3.2 is the payoff. Ordinary sheaf diffusion buys its non-collapse guarantee by constraining the hypothesis class to orthogonal maps. Joint diffusion buys it by checking a condition on the initial state — and Corollary 3.3 says the condition can always be met by rescaling one parameter. The guarantee stops depending on the learnable parameters entirely.
The geometric picture the paper gives is the clearest way to hold this. For one edge with \(d = 1\), the equilibria of the joint system are the surface \(zx = ty\), with \(z,t\) the restriction maps and \(x,y\) the features. Sheaf diffusion moves only \(x,y\), so it explores a horizontal plane through that surface and can miss equilibria that sit just above or below. Joint diffusion moves in every direction.
The architecture
Set \(\beta > \alpha\) when the data is heterophilic, prioritising movement in the maps.
What is given up is stated plainly: orthogonal and diagonal constraints — both known to help — cannot be imposed on maps that are themselves diffusing, and NSD’s universal sheaf approximation is lost. The trade is inductive bias for generality.
RiSNN: the same bias, simpler dynamics
The second model keeps the standard SNN update and changes only the map predictor. Let \(x_e(t) = \mathcal{F}_{u\trianglelefteq e}(t)x_u(t) - \mathcal{F}_{v \trianglelefteq e}(t)x_v(t)\) — the disagreement in the edge stalk, i.e. “the conversation”. Then
The map depends only on the relationship between the private opinion and the public conversation, not on the raw features. The outer product’s entries are inner products \(\langle x_{e,i}, x_{u,j}\rangle\) — the paper reads the whole thing as attention over multiple similarities. Parameter count is \(d^4\): more than JdSNN, still independent of \(c\).
Proposition 3.5. The learned restriction maps are feature-wise rotation invariant: for orthogonal \(Q\), the maps computed from \(X\) and from \(XQ\) are identical.
The mechanism is a one-line cancellation, and it is worth seeing because it is the entire content of the claim. Under \(X' = XQ\) we have \(x'_e = x_eQ\) and \(x'^{\top}_u = Q^{\top}x_u^{\top}\), so
The rotation cancels inside the outer product before the MLP ever sees it. Induction over layers extends it, with the base case handled by initialising \(\mathcal{F}(0) = \mathrm{Id}\).
Note precisely what is invariant: only the maps. The paper is careful that RiSNN itself is not rotation invariant, though setting \(\sigma = \mathrm{id}\) and \(W_2 = \mathrm{Id}\) yields a rotation-equivariant GNN — which is what makes it a candidate for the geometric problems where equivariance is a correctness requirement rather than a convenience.
Results
Four variants are evaluated, including two deliberate simplifications: RiSNN-NoT (drop the time dependence, \(x_e(t) = x_u(t) - x_v(t)\), to stabilise gradients) and JdSNN-NoW (no learnable weights at all in the map diffusion — \(W^*_1 = W^*_2 = \mathrm{Id}\), \(\mathcal{F}(0) = \mathrm{Id}\)).
| Dataset | \(h\) | RiSNN-NoT | RiSNN | JdSNN-NoW | JdSNN | Best prior sheaf model |
|---|---|---|---|---|---|---|
| Texas | 0.11 | 87.89 | 86.84 | 87.30 | 87.37 | BC-NLSD 87.57 |
| Wisconsin | 0.21 | 88.04 | 87.84 | 88.43 | 89.22 | NSD / BC-NLSD 89.41 |
| Squirrel | 0.22 | 51.24 | 53.30 | 51.28 | 49.89 | NSD 56.34 |
| Chameleon | 0.23 | 66.58 | 65.15 | 66.45 | 66.40 | NSD 68.68 |
| Cornell | 0.30 | 82.97 | 85.95 | 84.59 | 85.41 | BC-NLSD 87.30 |
| Citeseer | 0.74 | 75.07 | 76.23 | 75.93 | 73.27 | NSD 77.14 |
| Pubmed | 0.80 | 87.91 | 88.00 | 88.09 | 88.19 | MLP-NLSD 89.60 |
| Cora | 0.81 | 85.86 | 85.27 | 84.39 | 85.43 | NSP 87.38 |
The claim is statistical equivalence at lower parameter cost, not superiority, and read that way the table mostly holds. RiSNN-NoT takes Texas outright at 87.89. Wisconsin, Cornell and Cora land within roughly a point of the best prior sheaf model.
The synthetic study is the better evidence, and its design is the part worth stealing. Rather than separating classes by adding noise to different means — which would only let you test noise robustness — features are sampled from the surfaces of different \(n\)-dimensional ellipsoids sharing a centre. Because the ellipsoids are symmetric about that centre, every class has the same expected value and none are linearly separable, so a model cannot succeed by averaging. Edges come from a Watts–Strogatz variant with a controllable inter-class connection probability, giving a heterophily dial. On these datasets:
- Noise. JdSNN variants stay well ahead of standard SNNs up to 70% Gaussian feature noise. RiSNN, by contrast, is more sensitive to noise than a standard SNN — an internally inconsistent result the paper reports rather than suppresses.
- Heterophily. As inter-class edges increase, standard SNN accuracy falls significantly faster than JdSNN’s. This is the cleanest confirmation that the bias is real.
- Connectivity. SNNs improve sharply with more edges, since the map-predicting MLP gets more supervision. The proposed variants lead at low connectivity — the same data-versus-bias trade, isolated.
Cost
The limitations appendix is unusually candid. Backpropagation is harder: using the restriction maps to build the normalisation matrix \(D\) can produce exploding or undefined gradients, so \(D\) must be detached from the backward pass. That is a practical detail with theoretical consequences — the trained object is no longer exactly the model that was analysed. And time complexity is \(O(n(c^2+d^3) + md(c+d^3))\) against a diagonal SNN’s \(O(nc^2 + mc)\), so fewer parameters does not mean faster. Large \(d\) or large \(m\) will hurt.
✅ Key Takeaways
- \(\mathrm{MLP}(x_u \Vert x_v)\) is a universal sheaf approximator and therefore carries no bias toward the sheaves that handle heterophily. This paper trades that generality for structure.
- JdSNN diffuses features and restriction maps jointly. Lemma 1 makes it gradient descent on the sheaf Dirichlet energy in both variables; Theorem 3.2 gives non-trivial convergence from a condition on the initial state, not from constraining maps to be orthogonal.
- JdSNN uses \(2d^2\) parameters per layer instead of \(\ge 2d^2c\) — independent of the feature count, which is what makes federated and small-data settings viable.
- RiSNN predicts maps from \(x_e x_u^{\top}\); rotations cancel as \(QQ^{\top} = I\) inside the outer product, so the maps are provably invariant to feature-space rotation. The maps, not the model.
- Benchmarks are statistically equivalent to prior sheaf models with fewer parameters, except on Squirrel and Chameleon — the two largest graphs, where an unconstrained MLP has enough data to win.
- The ellipsoid-surface synthetic benchmark forces non-trivial aggregation by construction: all classes share an expected value, so averaging cannot separate them.
- Costs: \(D\) must be detached from backpropagation to avoid undefined gradients, and the models are slower per epoch than a diagonal SNN despite the smaller parameter count.
References
- Hernandez Caralt, F., Bernárdez Gil, G., Duta, I., Liò, P., & Alarcón Cot, E. (2024). Joint Diffusion Processes as an Inductive Bias in Sheaf Neural Networks. ICML 2024 Workshop on Geometry-grounded Representation Learning and Generative Modeling.
- Hansen, J., & Ghrist, R. (2021). Opinion Dynamics on Discourse Sheaves. SIAM Journal on Applied Mathematics, 81(5), 2033–2060.
- Bodnar, C., Di Giovanni, F., Chamberlain, B. P., Liò, P., & Bronstein, M. (2022). Neural Sheaf Diffusion. NeurIPS 2022.
- Barbero, F., Bodnar, C., Sáez de Ocáriz Borde, H., Bronstein, M., Veličković, P., & Liò, P. (2022). Sheaf Neural Networks with Connection Laplacians. ICML 2022 TAG-ML Workshop.
- Duta, I., Cassarà, G., Silvestri, F., & Liò, P. (2023). Sheaf Hypergraph Networks. Advances in Neural Information Processing Systems 36.
- Zaghen, O. (2024). Nonlinear Sheaf Diffusion in Graph Neural Networks. arXiv:2403.00337.
- Watts, D. J., & Strogatz, S. H. (1998). Collective Dynamics of ‘Small-World’ Networks. Nature, 393(6684), 440–442.
- Platonov, O., Kuznedelev, D., Babenko, A., & Prokhorenkova, L. (2023). Characterizing Graph Datasets for Node Classification: Homophily-Heterophily Dichotomy and Beyond. NeurIPS 2023.