PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariant Sheaf Neural Networks: Learning Geometric Transport on Graphs

Alessio Borgi, Mario Severino, Fabrizio Silvestri, Pietro Liò

arXiv:2608.28853v1cs.LGcs.AI

TL;DR

Efficient first-order equivariant architectures have limited ways to transform vector information across graph edges. ESNN addresses this by learning structured matrix-valued transport with exact Euclidean equivariance, and its evaluations report improved dynamics, symmetry-breaking direction recovery, selected mesh gains, and robustness to unseen rotations.

  • Problem

    First-order equivariant architectures are limited in how vector information is transformed across graphs, while generic matrix-valued sheaf maps do not respect geometric vector transformation laws.

  • Method

    ESNN keeps scalar and vector features first-order and factorizes learned edge transport into O(n)-covariant spatial actions and invariant channel mixing.

  • Results

    Across particle dynamics, mesh simulation, point-cloud classification, and molecular property prediction, ESNN improves selected tasks, recovers symmetry-breaking directions, and remains robust to unseen rotations.

  • Takeaways & Limitations

    Learning geometric transport across edges provides a complementary route to more expressive equivariant message passing without requiring higher-order representations.

  • Takeaways & Limitations

    The architecture assumes non-self interactions using relative displacement satisfy r_ij ≠ 0, with invariant edge attributes and suitable graph construction assumptions.

Abstract

from arXiv · show

Equivariant graph neural networks provide a principled way to model geometric systems, but efficient first-order architectures remain limited in how vector information can be transformed as it moves across a graph. We introduce \textsc{ESNN}, an Equivariant Sheaf Neural Network that enriches this interaction by learning directed, matrix-valued transport between neighboring vector features while preserving exact Euclidean equivariance. Rather than increasing the order of the representation, ESNN keeps scalar and vector features first-order and places the additional geometric flexibility in the edge transport itself. We characterize this transport theoretically, showing that when relative displacement is the only covariant geometric input, every linear $O(n)$-equivariant map decomposes into independent radial and tangential components, while learned covariant features enable richer feature-conditioned transformations. We also introduce controlled symmetry relaxation for systems with a preferred ambient direction, which may be prescribed or inferred from data while recovering full $E(n)$-equivariance when the directional pathway is inactive. Across particle dynamics, mesh-based simulation, point-cloud classification, and molecular property prediction, ESNN improves dynamics prediction, recovers the gravity axis when symmetry is broken, yields substantial gains on selected mesh tasks and long-horizon rollouts, and remains robust to unseen rotations. These results show that learning how geometric information is transported across edges offers a complementary route to expressive equivariant message passing without requiring higher-order representations.

1 INTRODUCTION

ESNN addresses the limited geometric transformation capacity of efficient first-order equivariant architectures by learning matrix-valued edge transport while preserving exact E(n)-equivariance. It further supports controlled symmetry relaxation when systems contain a preferred direction.

  • 1 INTRODUCTION: First-order representations can gain richer geometric interactions by enriching edge transformations rather than introducing higher-order features.This approach complements architectures that increase representation richness through higher-order features, spherical harmonics, or tensor products.
  • 1 INTRODUCTION: ESNN learns matrix-valued transport between neighboring vector features while preserving exact E(n)-equivariance.Its transport separates an O(n)-covariant spatial action from invariant channel mixing and remains within first-order scalar and vector representations.
  • 1 INTRODUCTION: When relative displacement is the only covariant input, independent radial and tangential actions characterize the complete class of linear O(n)-equivariant transports.The decomposition separately transforms vector components parallel and orthogonal to the relative displacement.
  • 1 INTRODUCTION: Controlled symmetry relaxation accommodates preferred ambient directions that may be prescribed or learned from data.A learnable coefficient controls the directional pathway: zero recovers full E(n)-equivariance, while activation retains the subgroup preserving the preferred direction.
  • 1 INTRODUCTION: The evaluation spans particle dynamics, mesh-based physical simulation, point-cloud classification, and molecular property prediction.These experiments test richer transport under full symmetry, recovery of symmetry-breaking directions, and transfer across geometric domains.

2 RELATED WORK AND BACKGROUND

The background situates ESNN between efficient low-order equivariant architectures and sheaf-based graph learning. It retains the sheaf view of local feature spaces while directly parameterizing equivariant transport between neighboring nodes.

  • 2 RELATED WORK AND BACKGROUND: Equivariant architectures differ in whether they retain simple scalar–vector representations or use higher-order features to capture angular structure.Cartesian models use invariant scalars and first-order vectors, whereas steerable models use higher-order features, spherical harmonics, and tensor products.
  • 2 RELATED WORK AND BACKGROUND: Relative displacements remove translations, leaving O(n) transformation laws as the remaining geometric requirement for operations built from relative geometry.Under an orthogonal transformation Q, the relative displacement transforms as r_ij → Qr_ij.
  • 2 RELATED WORK AND BACKGROUND: ESNN represents each node with invariant scalar channels and covariant vector channels attached to a local feature space.The vector component transforms geometrically, while scalar features are propagated through invariant message-passing operations.
  • 2 RELATED WORK AND BACKGROUND: Cellular sheaves assign local vector spaces to nodes and edges and connect them through linear restriction maps.These maps express node features in a common edge space and motivate transport formulations based on local feature spaces.
  • 2 RELATED WORK AND BACKGROUND: ESNN directly parameterizes equivariant node-to-node transport instead of learning separate incidence maps through an intermediate edge stalk.This retains the sheaf perspective while enforcing rotation- and reflection-compatible transport for geometric vector features.

3 EQUIVARIANT SPATIAL TRANSPORT

ESNN decomposes edge transport into covariant spatial operators and invariant channel operators, allowing vector messages to be transformed geometrically before aggregation. The transport can adapt to local edge context while retaining equivariance.

  • 3 EQUIVARIANT SPATIAL TRANSPORT: ESNN separates spatial transformation on R^n from channel mixing on R^c_v using left–right transport actions.Spatial operators act on the Euclidean dimension, while channel operators act on vector channels that transform trivially under O(n).
  • 3 EQUIVARIANT SPATIAL TRANSPORT: The local edge context C_ij contains relative displacement, optional covariant endpoint features, and invariant features derived from them.These quantities condition the transport for the directed interaction j → i.
  • 3 EQUIVARIANT SPATIAL TRANSPORT: Each transport component combines an O(n)-covariant spatial operator with an invariant channel operator.The spatial operator acts on the spatial axis, while the channel operator mixes vector channels using invariant edge features.
  • 3 EQUIVARIANT SPATIAL TRANSPORT: The two transformation laws are sufficient to make the resulting transport equivariant.This construction allows vector features to be transformed before aggregation without breaking equivariance.
  • 3 EQUIVARIANT SPATIAL TRANSPORT: For fixed edge context, transport is linear in the propagated vector feature but can depend nonlinearly on local node, edge, and geometric information.This enables interaction-specific adaptation while retaining the equivariance guarantee.
  • 3 EQUIVARIANT SPATIAL TRANSPORT: The transport family includes multiple spatial-operator choices, with the unified formulation activating subsets such as identity, radial, tangential, and skew components.Learned vector features provide additional covariant inputs beyond relative displacement for richer spatial operators.

4 EQUIVARIANT TRANSPORT MAPS

ESNN defines equivariant edge transports that keep scalar and vector features first-order while separating geometric spatial actions from invariant channel mixing. Its radial–tangential family is complete for displacement-conditioned linear O(n)-equivariant transport, while learned covariant features enable richer operators.

  • Transport families: ESNN instantiates four spatial transport families: identity, diagonal isotropic scaling, feature-conditioned rotation, and radial–tangential transformation.All preserve invariant channel actions under O(n).
  • Transport families: Orthogonal Transport constructs an edge-dependent spatial rotation from current vector features while retaining equivariance through conjugate transformation of its generator and matrix exponential.The complete edge map also includes learned channel mixing and gating, so it is not generally orthogonal in the full feature space.
  • Transport families: Radial–tangential transport independently transforms vector components parallel and orthogonal to the relative displacement.Its projection action can be evaluated without materializing dense n × n matrices, with spatial cost O(ncv).
  • Equivariance and completeness: When relative displacement is the only covariant input, unrestricted radial and tangential channel transformations span the complete class of linear O(n)-equivariant transports.This completeness does not extend unchanged when learned covariant vector features or SO(n)-specific orientation information are available.
  • Equivariance and completeness: The unified formulation extends displacement-only transport by allowing spatial operators to depend on learned vector features while preserving the same equivariance principle.It therefore lies outside the setting characterized by the displacement-only completeness theorem.

5 THE ESNN ARCHITECTURE

The ESNN layer builds invariant edge contexts, transports vector messages directionally, aggregates scalar and vector information, and updates representations within their respective transformation laws. Under stated graph, edge, and nonzero-displacement assumptions, the resulting layer is E(n)-equivariant.

  • Layer pipeline: ESNN constructs invariant edge contexts from distances, radial features, optional invariant edge attributes, and vector norms before parameterizing messages and transport.Self-information instead uses an explicit identity self-loop because no directional edge description is needed.
  • Layer pipeline: Vector messages use equivariant transport maps, while scalar features follow a separate invariant pathway that can include invariant vector inner products.This lets geometric information influence both representation types without changing their transformation laws.
  • Layer pipeline: Directed messages are aggregated through normalized transport diffusion, with optional invariant attention weights preserving O(n)-equivariance.Different edge orientations may carry different maps, so the general operator need not be self-adjoint.
  • Feature updates: Channel mixing and radial vector nonlinearities update vector features covariantly, while scalar features use ordinary linear maps and nonlinearities.Channel mixing acts only on feature channels, so it does not alter the spatial transformation type.
  • Coordinate and dynamical updates: The layer optionally performs an EGNN-style coordinate update from invariant features, yielding covariant displacements for evolving geometries and dynamical-system states.The full layer is E(n)-equivariant when topology and edge attributes satisfy the stated invariance assumptions and non-self displacements are nonzero.

6 CONTROLLED SYMMETRY RELAXATION

ESNN relaxes full Euclidean symmetry when external structure selects a preferred direction, while retaining translation equivariance and the subgroup preserving that direction. The relaxation is controlled by a coefficient initialized at zero, so the fully equivariant regime remains recoverable.

  • Directional conditioning: A preferred direction can enter the invariant edge context through the orientation-dependent scalar ⟨rij, g⟩, allowing transport coefficients to depend on edge alignment.The spatial transport itself remains unchanged; only its scalar conditioning is augmented.
  • Directional conditioning: The direction may be prescribed or learned globally, while its relaxation coefficient λ is initialized at zero.At λ = 0, the directional contribution vanishes exactly and the original E(n)-equivariant model is recovered.
  • Resulting symmetry: When λ ≠ 0, the layer preserves translations and orthogonal transformations in the stabilizer Og(n) of the preferred direction.The guaranteed symmetry group is Eg(n) = Og(n) ⋉ Rn.
  • Resulting symmetry: Orthogonal transformations outside Og(n) are not enforced by the architecture, whereas λ = 0 restores full E(n)-equivariance.Thus the model interpolates between full symmetry and the subgroup compatible with the fixed direction.

7 EXPERIMENTS

The experiments test ESNN across particle dynamics, mesh simulation, and point-cloud classification, including settings with full symmetry, broken symmetry, directional fields, and unseen rotations. ESNN recovers gravity direction, improves selected long-horizon mesh rollouts, and remains robust to rotation shifts, though mesh gains are not uniform.

  • Experimental design: The study evaluates ESNN on charged and gravity-augmented N-body dynamics, three MeshGraphNets benchmarks, ModelNet40, and QM9.The experiments progressively test transport enrichment, symmetry relaxation, mesh dynamics, rotation generalization, and molecular-property prediction.
  • Particle dynamics: Gravity-relaxed ESNN models reduce MSE from approximately 0.10–0.13 to around 0.02 and learned models achieve Ag = 1 in every reported case.The learned setting closely matches the fixed-direction setting while receiving no prior gravity-axis information.
  • Rotation generalization: ModelNet40 ESNN variants maintain approximately 84–86% accuracy, with ΔOOD at most 1.1 percentage points and 0.2 points for ESNN-Ortho.Orientation-sensitive PointNet and DGCNN baselines lose more than 50 percentage points under arbitrary test-time rotations.

8 CONCLUSIONS AND LIMITATIONS

ESNN enriches first-order equivariant message passing through structured matrix-valued edge transport rather than higher-order representations. Its scope is limited to scalar and first-order vector features, and its completeness result applies specifically to displacement-conditioned linear transport.

  • Conclusion: ESNN learns structured matrix-valued transport between neighboring vector features while retaining a first-order Cartesian representation.Additional geometric flexibility is placed in edge transport instead of increasing representation order.
  • Theory: With relative displacement as the only covariant input, linear O(n)-equivariant transport reduces to independent radial and tangential actions.Learned covariant vector features allow richer feature-conditioned transformations.
  • Scope and results: Controlled symmetry relaxation recovers full E(n)-equivariance when the directional pathway is inactive and supports preferred-direction systems when active.The experiments report improvements in particle dynamics, selected mesh rollouts, and robustness to unseen rotations.
  • Limitations: The current formulation is limited to scalar and first-order vector features, while completeness is specific to displacement-conditioned linear transport.More general covariant inputs admit a broader class of spatial operators.

Supplementary Material

The supplementary material situates ESNN at the intersection of equivariant message passing and sheaf learning, and develops its algebraic transport foundations. It distinguishes continuous Euclidean equivariance on geometric fibers from other symmetry structures and formalizes separable spatial–channel transport.

  • Related Work: Order equivariance acts on combinatorial indexing structures, whereas ESNN constrains continuous Euclidean actions on geometric feature fibers.These are distinct but compatible symmetry structures.
  • Related Work: ESNN combines low-order Euclidean-equivariant features with the matrix-valued transport perspective of sheaf-based graph learning.Its direct node-to-node transport differs from classical incidence-map constructions through intermediate edge stalks.
  • Algebraic foundations: The appendix derives separable spatial–channel transport as a finite sum of left spatial actions and right channel actions.The spatial dimension carries the O(n) action, while channels transform invariantly.
  • Tensor-product structure: Figure 4 depicts each transport term acting independently on spatial and channel axes.Spatial matrices act by left multiplication, while channel matrices act by right multiplication.
  • Equivariance constraints: ESNN enforces equivariance by making spatial operators O(n)-covariant and channel maps invariant.The decomposition alone is algebraic; equivariance additionally constrains how operators depend on local edge context.

B.2 COMPLETENESS UNDER DISPLACEMENT-ONLY CONDITIONING

Under displacement-only covariant conditioning, every linear O(n)-equivariant transport separates into radial and tangential spatial components, with channel maps depending only on invariant quantities. Feature-conditioned operators extend beyond this complete displacement-only class, while general ESNN transport also supports directed interactions.

  • Completeness theorem: Theorem 4.3 characterizes every linear O(n)-equivariant transport with nonzero displacement as a radial–tangential decomposition.The result assumes relative displacement is the only covariant geometric input and allows additional O(n)-invariant scalar conditioning.
  • Completeness theorem: Invariant conditioning may change the radial and tangential channel maps through quantities such as ∥r∥, but cannot introduce additional spatial operators.The resulting form is complete under the stated displacement-only assumptions.
  • Scope and limitations: The completeness claim applies to O(n), not arbitrary SO(n)-equivariant kernels, because orientation-sensitive operators can change sign under reflections.In three dimensions, the cross-product operator is SO(3)-equivariant but acquires a reflection-dependent sign.
  • Feature-conditioned transport: Learned covariant vector features enlarge the local geometric context, enabling feature-conditioned spatial operators beyond the radial–tangential form.These transports can depend nonlinearly on evolving features while remaining linear in the transported vector once the context is fixed.
  • Sheaf interpretation: ESNN learns directed neighboring-node transport directly rather than obtaining it indirectly by composing cellular-sheaf restriction maps through an edge stalk.The general operator permits independent directed edge maps; self-adjoint and connection-style cases impose additional constraints.

C.4 AMBIENT EQUIVARIANCE AND LOCAL GAUGE TRANSFORMATIONS

ESNN enforces a common ambient Euclidean transformation on coordinates and vector features, rather than independent changes of local node or edge bases. Its ambient transport therefore need not be locally gauge-equivariant, although a gauge-equivariant extension is identified as future work.

  • Ambient equivariance: ESNN transforms coordinates and all vector features under one common Q ∈ O(n), with spatial transport changing by conjugation.This is ambient E(n)-equivariance in a shared Cartesian frame.
  • Local gauge transformations: Independent node-wise basis changes generally do not preserve the feature-derived Orthogonal Transport used by ESNN.When Qi = Qj = Q, the construction reduces to the common ambient covariance required by ESNN.
  • Scope and extension: The orthogonal, adjoint-consistent specialization can realize a classical connection-sheaf transport, but the general architecture does not claim arbitrary local gauge covariance.A locally gauge-equivariant ESNN would require transport generators built from gauge-covariant quantities.

D.1 PROOF OF PROPOSITION 3.2: O(n)-EQUIVARIANCE OF THE TRANSPORT MAP

The transport map is O(n)-equivariant when spatial operators transform covariantly and channel operators remain invariant. This applies to the general construction and to the Orthogonal and Radial–Tangential transport families.

  • General transport: Proposition 3.2 establishes O(n)-equivariance by combining covariant spatial operators with O(n)-invariant channel operators.Evaluating the transport on transformed vector features yields the required transformed output.
  • Orthogonal Transport: Orthogonal Transport is O(n)-equivariant because its cross-feature matrix transforms by conjugation, while its Frobenius normalization and gate are invariant.The resulting normalized rotation satisfies the required covariance law.
  • Radial–Tangential Transport: Radial–Tangential Transport is O(n)-equivariant because its radial projector transforms by conjugation and its coefficients are invariant.The construction applies for nonzero relative displacement.
  • Completeness connection: The displacement-only completeness proof identifies the two spatial projectors P∥ and P⊥, each paired with an arbitrary invariant vector-channel endomorphism.This follows from the stabilizer of a nonzero displacement and its commutant.

D.5 PROOF OF THEOREM 5.1: E(n)-EQUIVARIANCE OF THE ESNN LAYER

The ESNN layer is E(n)-equivariant when graph construction, edge attributes, and transport coefficients satisfy the stated invariance conditions. The proof applies across vector, scalar, diffusion, attention, residual, and coordinate updates.

  • Theorem conditions: Theorem 5.1 gives E(n)-equivariance when graph topology is invariant, edge attributes are O(n)-invariant, and non-self interactions use nonzero displacements.The transport maps must also satisfy Proposition 3.2’s covariance conditions.
  • Aggregation: Invariant diffusion or attention coefficients preserve equivariance because they do not change under the Euclidean action.The same covariance argument applies when normalized diffusion is replaced by receiver-normalized attention.
  • Layer updates: The residual vector update remains covariant because channel mixing acts only on the multiplicity dimension, while scalar updates use invariant quantities.The coordinate update also transforms correctly because its increments are covariant.
  • Symmetry relaxation: The dynamical extension preserves equivariance for translations and the subgroup O_g(n) that fixes a preferred direction, but does not enforce equivariance outside that subgroup when λ ≠ 0.When λ = 0, the directional conditioning vanishes and full E(n)-equivariance is recovered.

E.3 MODELNET40 POINT-CLOUD CLASSIFICATION

ModelNet40 evaluates graph-based point-cloud classification under matched and out-of-distribution rotation protocols, using an invariant graph-level readout. The broader benchmarks cover mesh simulation and invariant molecular-property prediction with scalar and vector geometric features.

  • ModelNet40 construction: Each CAD model is sampled as a point cloud and converted into a local k-nearest-neighbor graph for 40-way classification.The graph uses Euclidean geometric interactions, while node outputs are aggregated by an invariant readout.
  • Evaluation: An invariant graph-level readout is designed to keep predicted classes unchanged under global input rotations.Test classification accuracy is reported for each rotation protocol.
  • Rotation protocol: The z/z, z/SO(3), and SO(3)/SO(3) protocols distinguish matched rotations, unseen arbitrary orientations, and arbitrary rotations during both training and testing.The z/SO(3) setting specifically measures out-of-distribution rotation generalization.
  • Additional benchmarks: Mesh simulation compares ESNN transport families across CYLINDERFLOW, DEFORMINGPLATE, and AIRFOIL, spanning fluid and structural dynamics on unstructured meshes.The domains differ in physical regime, state representation, and boundary geometry.
  • Additional benchmarks: The mesh models combine relative displacements and invariant geometric quantities with covariant vector channels and invariant scalar channels.Rollout evaluation reports RMSE for one-step prediction, a 50-step rollout, and the complete trajectory.
  • Additional benchmarks: QM9 evaluates invariant molecular-property prediction from three-dimensional molecular graphs using mean absolute error across twelve quantum-chemical targets.The benchmark uses separate models for each target and standard graph-level invariant outputs.
Loading 2608.28853v1…