Sheaf GNNs for Molecular Property Prediction

Why Molecular Graphs Are Heterophilic

In a molecular graph, adjacent atoms (bonded atoms) are typically of different element types — carbon bonds to oxygen, oxygen bonds to hydrogen. The features of adjacent atoms (atomic number, electronegativity, orbital structure) are systematically different.

Standard GCN on molecular graphs averages these dissimilar features — an operation that conflates chemically distinct information. This is exactly the heterophily problem: adjacent nodes (atoms) should contribute different information, not be averaged toward the same representation.

Sheaf motivation: Different bond types (C-C, C-O, C-N, C-H) encode different chemical relationships. A sheaf with bond-type-specific restriction maps can represent the relational geometry of each bond type independently — analogous to R-GCN for knowledge graphs but applied to molecular chemistry.

From 2D Connectivity to 3D Geometry

Molecular GNNs have progressively incorporated more geometric information:

Each level adds richer geometric information but at higher computational cost.

Sheaf Maps as Bond-Angle Encodings

Consider three bonded atoms: i — j — k (atom j is bonded to both i and k). The bond angle θ_{ijk} = ∠(r_i − r_j, r_k − r_j) is the key geometric quantity DimeNet uses.

In the sheaf framework: the restriction map F_{i▷e_{ij}} encodes the “orientation” of atom i relative to the bond e_{ij}. The bond angle is encoded in the composition of maps:

When restriction maps are orthogonal, the holonomy around the path i–j–k is a rotation by the bond angle. NSD with orthogonal maps implicitly learns bond-angle-like geometric information — without explicitly computing angles.

Formal claim: For a sheaf with orthogonal maps trained on molecular data, the learned maps F_{v▷e} encode the 3D geometric relationship between atom v and bond e — in the same information-theoretic sense as DimeNet’s angle features.

Multi-Relational Molecular Sheaves

Different bond types (single, double, triple, aromatic) have different geometric and electronic properties. A multi-relational molecular sheaf assigns different restriction maps per bond type:

The multi-relational Sheaf Laplacian:

This separates the contribution of different bond types to the diffusion — single bonds propagate information differently from double bonds.

Equivariant Molecular Sheaves

For molecular property prediction with 3D coordinates, E(n) equivariance is required: predictions must be invariant to rotation, translation, and reflection of the molecule.

A sheaf with O(3)-valued restriction maps achieves this: maps O_{v▷e} ∈ O(3) encode the 3D orientation of atom v relative to bond e. The Sheaf Laplacian is gauge-equivariant under O(3) — applying a global rotation to all atoms corresponds to a gauge transformation.

This is the connection Laplacian approach applied to molecular graphs — the sheaf GNN becomes a gauge-equivariant model for 3D molecules.

Comparison with EGNN: EGNN achieves E(n) equivariance via distance-only messages (no directional information). Orthogonal sheaf maps add directional information (encoded in the maps) while maintaining equivariance — potentially capturing more geometric detail than EGNN.

Experimental Setup: QM9 Benchmark

QM9 is the standard molecular property benchmark:

• 133,885 small organic molecules (≤9 heavy atoms)
• 12 quantum chemistry targets: U₀ (internal energy), HOMO/LUMO gap, dipole moment, polarisability, …
• MAE (mean absolute error) is the standard metric

A molecular sheaf GNN for QM9:

1. Node features: atom type (one-hot), degree, aromaticity
2. Edge features: bond type (one-hot), interatomic distance

3. Sheaf predictor: MLP(h_u, h_v, bond_type,

   r_u − r_v

   ) → restriction maps

4. 3-layer sheaf diffusion with d=4 orthogonal maps
5. Graph readout: mean pooling over node stalks
6. Final MLP for property prediction

Expected Performance and Limitations

Based on the theoretical analysis:

Expected benefit:

• Heterophilic bonds (C-O, N-H): sheaf maps correctly represent the dissimilar atom features
• Geometric information: orthogonal maps encode bond angles implicitly
• Multi-relational structure: different maps per bond type capture bond chemistry

Expected limitations:

• 3D position encoding: sheaf maps encode relative orientations, not absolute positions → equivariance but limited chirality resolution
• Long-range effects: with K=3 sheaf layers, only 3-hop interactions are captured → large molecules with long-range electronic effects require more layers or global attention
• No explicit angle features: unlike DimeNet which explicitly computes θ_{ijk}, the sheaf must learn this from data — may require more training examples

Comparison with SchNet: SchNet uses distance-based filter functions, not restriction maps. For predicting energy (U₀): SchNet ≈ 14 meV MAE; a sheaf GNN with orthogonal maps (without explicit 3D position) would likely achieve ≈ 20–30 meV — better than 2D-only MPNNs but not competing with full 3D equivariant models.

Drug Discovery Applications

Beyond QM9, molecular sheaf GNNs are applicable to:

• ADMET prediction: molecular heterophily (different atom types in a drug molecule) benefits from sheaf maps; OGB-molhiv, OGB-molpcba benchmarks
• Protein-ligand binding: the protein-ligand interface is a heterophilic bipartite graph — sheaf maps encode the complementarity between protein residues and ligand atoms
• Retrosynthesis: reaction graphs have heterophilic structure (reagents of different types interacting) — sheaf maps encode the chemical compatibility constraints

References

• Gilmer, J., Schütt, K. T., Ramsundar, B., Ramakrishnan, R., Bronskill, M., Gomes, C., & Dahl, G. E. (2017). Neural Message Passing for Quantum Chemistry. ICML 2017 (MPNN: unified molecular GNN framework — the 2D baseline that sheaf GNNs extend).
• Klicpera, J., Groß, J., & Günnemann, S. (2020). Directional Message Passing for Molecular Graphs. ICLR 2020 (DimeNet: explicit bond angles — the geometric information sheaf restriction maps can implicitly encode).
• Satorras, V. G., Hoogeboom, E., & Welling, M. (2021). E(n) Equivariant Graph Neural Networks. ICML 2021 (EGNN: the equivariant baseline that orthogonal sheaf maps can extend with richer geometric structure).