PaperlayerView exact revision on arXiv

Source-linked AI summary

Spherical Channels for Modeling Atomic Interactions

C. Lawrence Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, Brandon Wood

arXiv:2206.14331v2physics.chem-phcs.CEcs.LGphysics.comp-ph

TL;DR

DFT estimates atomic energies and forces but is computationally expensive, motivating machine-learning approximations. SCN uses spherical-function atom embeddings and edge-oriented message passing, and achieves state-of-the-art OC20 results for energy and force prediction.

  • Problem

    DFT estimates atomic energies and forces but can take hours or days to evaluate a single material, motivating faster machine-learning approximations.

    Source: §1 · p. 1

  • Method

    SCN is a graph neural network with spherical-function atom embeddings and edge-oriented messages, whose equivariance constraints can be relaxed for more expressive transformations.

    Source: §1 · p. 2, §2.3 · p. 6

  • Results

    SCN achieves state-of-the-art OC20 results across energy and force tasks, with 8%−11% improvements in force and initial-structure-to-relaxed-energy prediction.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    For catalyst discovery, machine-learning approximations may enable high-throughput screening where DFT evaluations take hours or days per material.

    Source: §1 · p. 1

    SCN is prone to overfitting on smaller datasets, and its MD17 and QM9 results are not state of the art.

    Source: §5 · p. 10

Abstract

from arXiv · show

Modeling the energy and forces of atomic systems is a fundamental problem in computational chemistry with the potential to help address many of the world's most pressing problems, including those related to energy scarcity and climate change. These calculations are traditionally performed using Density Functional Theory, which is computationally very expensive. Machine learning has the potential to dramatically improve the efficiency of these calculations from days or hours to seconds. We propose the Spherical Channel Network (SCN) to model atomic energies and forces. The SCN is a graph neural network where nodes represent atoms and edges their neighboring atoms. The atom embeddings are a set of spherical functions, called spherical channels, represented using spherical harmonics. We demonstrate, that by rotating the embeddings based on the 3D edge orientation, more information may be utilized while maintaining the rotational equivariance of the messages. While equivariance is a desirable property, we find that by relaxing this constraint in both message passing and aggregation, improved accuracy may be achieved. We demonstrate state-of-the-art results on the large-scale Open Catalyst dataset in both energy and force prediction for numerous tasks and metrics.

1 Introduction

DFT estimates atomic energies and forces but can take hours or days per material, motivating machine-learning approximations. SCN combines spherical-harmonic atom embeddings with edge-oriented messages and reports state-of-the-art results on OC20.

1 Introduction

  • DFT can estimate atomic energies and forces but takes hours or days to evaluate a single material; machine learning could enable faster, high-throughput screening.The introduction cites catalyst discovery and renewable-energy storage among relevant applications.Source: §1 · p. 1
  • Rotation equivariance can aid generalization, but may restrict model expressiveness through constrained transformations or limited interaction sizes.The introduction contrasts these constraints with non-equivariant models, which may need to learn approximate equivariance from data augmentation.Source: §1 · pp. 1–2
  • SCN represents each atom with spherical-function channels and steers embeddings by edge orientation, using rotation-invariant coefficients and approximately equivariant nonlinear aggregation.The paper also reports improved performance when equivariance constraints are relaxed and additional coefficients are used.Source: §1 · p. 2
  • On OC20, SCN achieves state-of-the-art force and initial-structure-to-relaxed-energy prediction, with 8%−11% improvements and greater sample efficiency.The reported improvements apply to force and initial structure to relaxed energy prediction.Source: §1 · p. 2

2 Approach

SCN represents each atom with spherical functions and uses edge-oriented message passing and spherical aggregation to predict energies and forces. Its design uses rotational constraints selectively, with some operations relaxing strict equivariance.

2 Approach

  • SCN takes atom identities and positions as input to a graph neural network that predicts structure energy and per-atom forces.The graph connects each atom to nearby atoms.Source: §2 · p. 2

2.1 Node Embeddings

  • Each node embedding comprises spherical channels that encode orientation information for all neighboring atoms, represented with spherical harmonics.The coefficients are initialized from atomic number for degree l = 0, with higher-degree coefficients set to zero.Source: §2.1 · p. 3

2.2 Message Passing

  • For each edge, SCN rotates the source embedding to align the edge with the z-axis, computes a message using the rotated embedding and edge information, then rotates the message back.The message network receives atomic numbers and interatomic distance as edge inputs.Source: §2.2 · pp. 3–4

2.2.1 Message Equivariance

  • Restricting message inputs and outputs to m = 0 yields rotation-equivariant messages, while including m ∈ [−1, 1] can improve performance without strictly enforcing equivariance.For m ∈ [−1, 1], four z-axis rotation samples are combined; higher m values require more taps.Source: §2.2.1 · pp. 4–5

2.3 Message Aggregation

  • After summing incoming messages, SCN applies an unconstrained pointwise nonlinearity on a spherical grid and converts the result back to spherical-harmonic coefficients.Using degrees L and L − 1 as separate resolutions allows the network to learn circular functions on the sphere; the transformation is approximately rotation-equivariant, subject to discrete sampling limitations.Source: §2.3 · pp. 5–6

2.4 Energy and Forces Estimation

  • SCN estimates energy by integrating a per-atom energy function over orientations and summing across atoms; it predicts forces either from energy gradients or from integrated directional force estimates.The gradient approach enforces energy conservation but is more computationally expensive; the direct force approach is more efficient but does not enforce energy conservation.Source: §2.4 · p. 7

3 Experiments

Experiments evaluate SCN on OC20 tasks and compare it with existing models, then test model variations on OC20 2M. SCN reports stronger force and relaxation-energy results, while ablations examine how angular information, equivariance, and scaling choices affect performance.

3.1 Implementation Details

  • Experiments cover OC20 test tasks, 2M ablations, and a separate direct-energy model comparison for IS2RE; default runs directly estimate forces and use m ∈ [−1, 1] in message passing.The OC20 dataset contains over 130M training examples.Source: §3 · p. 7, §3.1 · p. 7

3.2 OC20 Tasks

  • On OC20 2M, SCN improves S2EF force prediction by ≈14% and IS2RE energy prediction by ≈10% over GemNet-OC.Energy MAE on S2EF is comparable to GemNet-OC.Source: Table 1 · p. 8, §3.2 · p. 8
  • On All and All+MD, SCN improves S2EF force MAE by ≈9% and outperforms GemNet-OC by over 7% on IS2RE.SCN L = 6 performs similarly to GemNet-OC on S2EF energy MAE; SCN predicts energy and forces with one model, while the cited GemNet-OC comparator combines two models.Source: §3.2 · p. 8, Table 2 · p. 9

3.3 Ablation studies

  • Using m ∈ [−1, 1] improves results over m = 0, while removing message rotations or nonlinear aggregation worsens metrics; larger m ranges add cost with diminishing returns.Force results are especially affected when messages are not rotated, which the passage relates to forces’ dependence on angular information.Source: §3.3 · p. 9
  • Gradient-derived, energy-conserving forces are slightly less accurate and 10 times slower than direct force prediction in the tested smaller model.The gradient-based model uses L = 4, H = 512, and 12 layers.Source: §3.3 · p. 9
  • Increasing L has a stronger effect on force estimation, whereas increasing depth improves energy estimation.The authors suggest angular resolution may explain the force effect, while greater information travel may help predict energy from the full atomic structure.Source: §3.3 · p. 9
  • Multiple z-axis rotations with 4-taps improve force prediction, while aggregating two bands improves both energy and force prediction.Four-tap sampling reduces throughput by over 2× but increases sample efficiency enough to yield similar training times for the same accuracy.Source: §3.3 · p. 9

4 Related work

Prior atomic-property GNNs use invariant scalar features or explicitly equivariant representations, with recent equivariant models performing well on small molecular datasets. SCN uses spherical harmonics but relaxes strict equivariance to allow more expressive operations.

4 Related work

  • Invariant GNNs progress from species and pairwise distances in CGCNN[45] and SchNet[35] to triplet or quadruplet angles in DimeNet[12, 13], SphereNet[28], and GemNet[14, 15].Enumerating these higher-order interactions requires careful design to maintain efficiency.Source: §4 · pp. 9–10
  • Recent equivariant models outperform traditional GNNs on small molecular datasets, while SCN belongs to this broader class without being strictly equivariant.The cited small datasets are MD17[7] and QM9[32].Source: §4 · p. 10
  • SCN shares spherical-harmonic representations with equivariant models but applies unconstrained functions to orientation-invariant m = 0 coefficients and includes broader coefficients to relax message equivariance.SCN also projects spherical channels onto a grid for pointwise 1 × 1 convolutions and nonlinear mixing across harmonic degrees.Source: §4 · p. 10
  • Unlike approaches that typically use L = 1 or L = 2, SCN reports increased accuracy at higher degrees L = 4, 6, and 8.The passage suggests this may result from SCN’s increased expressivity.Source: §4 · p. 10

5 Discussion

The discussion concludes that SCN performs best when equivariant constraints are relaxed, while identifying computational costs and overfitting on smaller datasets as limitations. It also notes potential societal risks from chemistry advances and suggests exploring pretraining on OC20 before fine-tuning on smaller datasets.

5 Discussion

  • Enforcing energy conservation through force estimation from energy gradients is computationally expensive, potentially limiting SCN’s use in molecular dynamics.Non-energy-conserving models remain practically useful for structure relaxations and transition-state searches; scaling with L is O(n^2), which may limit L beyond 8.Source: §5 · p. 10
  • SCN may overfit smaller datasets, and its results on MD17 and QM9 are not state-of-the-art; pretraining on OC20 and fine-tuning is proposed for future study.The authors attribute overfitting risk to increased expressivity and larger model sizes.Source: §5 · p. 10
  • The discussion warns that chemistry advances can have unintended consequences, citing fertilizer overuse, carbon-intensive production, and wartime explosives, and expresses hope that OC20-focused research will serve beneficial applications.Source: §5 · p. 10
  • The spherical-harmonic GNN achieves best performance when rotation-equivariant output constraints are relaxed for more expressive nonlinear transformations, and demonstrates state-of-the-art results across numerous tasks on the large-scale OC20 dataset.Source: §5 · p. 10

Checklist

The checklist reports release and documentation of experimental resources, while noting that repeated runs were too expensive for error bars. The appendix describes rotation-equivariance measurements, which identify approximate equivariance and several sources of deviation.

Checklist

  • The checklist reports code release and training-detail and resource disclosures, but no error bars because repeated experiments were too expensive; rough run-variation estimates are provided in Table 1.Source: Checklist · p. 14

Appendix: Table of contents

  • The appendix contents list rotation-equivariance measurements, direct OC20 IS2RE results, implementation details, spherical-harmonic properties, overfitting, and model-size analyses.Source: Appendix: Table of contents · p. 15

A Empirical measurement of equivariance to rotations

  • The rotation study measures differences between predictions for original and randomly rotated inputs, identifying nonzero-order message-passing coefficients, nonlinear aggregation, output blocks, and numerical precision as possible sources.The reported comparison uses 1,000 structures and models with 12 layers and L = 6.Source: A Empirical measurement of… · p. 15
  • Using m ∈ [−1, 1] coefficients increases rotational MAD relative to equivariant m = 0 message passing, but a four-tap model is nearly identical to m = 0 without AMP.MAD remains nonzero for m = 0 because of other error sources, and AMP produces higher MAD from numerical-precision limits.Source: A Empirical measurement of… · p. 15

A.1 Message aggregation

Message aggregation uses discrete sphere sampling, so it is not strictly rotation-equivariant, although sufficiently fine sampling yields approximate equivariance in practice. Nonlinear depth increases rotational differences, and the paper uses two nonlinear layers followed by a linear layer.

A.1 Message aggregation

  • Nonlinear aggregation shows 4%–7% differences between rotated outputs, and differences are higher with two nonlinear layers than with one; the paper uses two nonlinear layers followed by a linear layer.Differences between SiLU and ReLU are minimal.Source: §A.1 · pp. 15–16
  • Discrete sphere sampling makes the aggregation not strictly equivariant, but sampling at resolution 2(L + 1) or higher leads to approximate equivariance in practice.Nonlinearities can introduce frequencies above the maximum spherical-harmonic degree L, creating a risk of aliasing; smooth activations may reduce but not eliminate it.Source: §A.1 · p. 15
  • Table 4 compares percentage rotational differences in message-aggregation outputs across activation functions and depths.Source: Table 4 · p. 16

A.2 Energy and Force Output Blocks

The energy and force output blocks are nearly equivariant and have a negligible effect on overall network equivariance. Their rotational differences are small relative to output errors and smaller than those of other blocks.

A.2 Energy and Force Output Blocks

  • The output blocks are nearly equivariant for both energy and forces, with negligible negative impact on the network’s overall equivariance.The authors suggest that integrating over the entire sphere may cause differences to cancel.Source: §A.2 · p. 16
  • The 3-layer output-block MADs are below 1% relative to energy and force MAEs and are significantly smaller than those of other blocks.Source: Table 5 · p. 17

B OC20 IS2RE Direct Results

The section distinguishes relaxation-based IS2RE from direct prediction and reports state-of-the-art energy MAE for SCN across test splits, with its largest improvements on out-of-domain splits.

B OC20 IS2RE Direct Results

  • IS2RE can predict relaxed energy either after using an S2EF model to relax atomic positions or directly from the initial structure.Source: B OC20 IS2RE Direct Results · p. 17
  • The relaxation-based approach is more accurate, while direct prediction is more efficient in training and inference and relies more on global atomic arrangement.The direct approach cannot use detailed position and angular information in the same way.Source: B OC20 IS2RE Direct Results · p. 17
  • SCN achieves state-of-the-art energy MAE across test splits among models without ancillary losses, with improvements most pronounced on out-of-domain splits.Energy within Threshold is comparable to SEGNN.Source: B OC20 IS2RE Direct Results · p. 17

C Implementation details

The implementation uses reduced spherical-channel resolution for more distant neighbors, spherical Fibonacci sampling for output blocks, and a PaiNN baseline with directly predicted forces.

C.1 Varying spherical channel resolution

  • For neighbors ranked 13–40, the implementation reduces L by 2 (minimum 4) and H by a factor of 4; it omits neighbors ranked beyond 40 or beyond 8 Å.The 12 closest neighbors use the paper’s standard settings.Source: §C.1 · p. 17

C.2 Point sampling on the sphere for output blocks

  • Output-block integrals are approximated with approximately evenly distributed spherical Fibonacci points, weighted to compensate for varying point density.The Gaussian weighting uses σ = 0.5.Source: §C.2 · p. 18

C.3 PaiNN baseline

  • The reimplemented PaiNN baseline predicts forces directly from vectorial features, breaking energy conservation but proving essential for good OC20 performance.Source: §C.3 · p. 18

D Note on spherical harmonics properties

E Overfitting on the training dataset

Training shows substantial energy overfitting but little force overfitting, suggesting that SCN fits complex functions while leaving room to improve generalization.

E Overfitting on the training dataset

  • Energy errors show significant overfitting, unlike force errors, which show very little overfitting during training.The authors suggest the difference may reflect energies being per-structure while forces provide more per-atom examples.Source: E Overfitting on the training dataset · p. 18
  • The energy overfitting persists despite the large OC20 All + MD dataset and indicates room to improve generalization through model changes or data augmentation.Similar trends occur on OC20 2M.Source: E Overfitting on the training dataset · p. 18

F Impact of model size

SCN outperforms GemNet-OC at similar model sizes, and its accuracy improves with scale; even its smallest model exceeds other reported approaches.

F Impact of model size

Related papers

Loading 2206.14331v2…