Low-Rank Clebsch--Gordan Coupling for Equivariant Spherical Forecasting
Organizations: Princeton University
Abstract
Spherical-harmonic neural operators provide an efficient representation for global forecasting, but their discretized pointwise nonlinearities can break rotational equivariance through aliasing. Clebsch-Gordan (CG) tensor products provide exactly equivariant interactions between spherical-harmonic representations, but their computational cost grows rapidly with spectral bandwidth and has limited their use at forecasting scale. We introduce LCG-SFNO, which combines low-rank CG interactions within a controlled low-degree band with coefficient-space residual propagation and an exactly projected quadratic nonlinearity, preserving SO(3) equivariance throughout the learned operator. To make the CG interaction practical, we further develop SSHCG, a sparse GPU implementation with a factorized alternative for wider coupled bands. On spherical shallow-water dynamics and the ERA5-based WeatherBench benchmark, LCG-SFNO reduces rotational error to numerical precision (more than three orders of magnitude below SFNO on shallow water and five orders of magnitude below it on WeatherBench), while matching SFNO on shallow-water forecasting and substantially improving forecast skill relative to SFNO on WeatherBench. Matched ablations show that CG coupling accelerates convergence, while SSHCG keeps the coupling to about 10 percent of a training step at the operating bandwidth, making exact equivariant coupling practical at forecasting scale.
Figures & tables
| Shallow water | WeatherBench | |||
| Model | 1 day | 3 days | 5 days | |
| ClimaX | ||||
| FourCastNet | ||||
| SHNet (GSNO) | ||||
| SFNO | ||||
| LCG-SFNO (ours) | ||||
| Model | ACC | MSE | Epochs | Equivariance |
|---|---|---|---|---|
| A SFNO reference | never | |||
| B + full-band backbone | 39 | |||
| C + CG coupling | 13 | |||
| D + oversampled nonlinear pathway | 21 |
| Band | Triples | Dense | Fused | Tiled | Grid route |
|---|---|---|---|---|---|
| 4/8 | 85 | 3.2 | 3.2 | 3.2 | 3.2 |
| 8/16 | 489 | 4.5 | 4.2 | 3.6 | 3.6 |
| 12/24 | 1,469 | 16.6 | 6.5 | 4.6 | 4.1 |
| 16/31 | 3,280 | 53.2 | 13.3 | 7.2 | 4.7 |
| 24/31 | 9,456 | 325 | 51.9 | 17.4 | 5.1 |
| 31/31 | 16,400 | 855 † | 152 | 35.6 | 5.3 |
| Model | Params | 1 h | 5 h | 10 h |
|---|---|---|---|---|
| ClimaX | 1.96M | 421 0 | 474 5 | 535 10 |
| FourCastNet | 1.66M | 285 2 | 687 40 | 1309 110 |
| SHNet (GSNO) | 0.87M | 142 6 | 391 16 | 625 189 |
| SFNO | 0.28M | 58.0 1.2 | 121 6 | 219 22 |
| LCG-SFNO (ours) | 0.27M | 57.4 1.4 | 119 7 | 211 17 |
| 1 day | 3 days | 5 days | |||||
|---|---|---|---|---|---|---|---|
| Model | Params | ACC | MSE | ACC | MSE | ACC | MSE |
| ClimaX | 5.40M | 93.2 0.3 | 76.8 0.5 | 78.3 0.4 | 242.4 2.5 | 53.1 0.8 | 435.0 4.5 |
| FourCastNet | 5.60M | 92.8 0.1 | 82.8 0.3 | 69.3 0.1 | 325.1 0.9 | 50.5 0.1 | 440.1 0.1 |
| SHNet (GSNO) | 4.03M | 92.4 0.0 | 92.7 0.1 | 74.2 0.1 | 281.6 0.9 | 54.9 0.1 | 420.1 1.1 |
| SFNO | 5.26M | 84.0 0.0 | 187.7 0.2 | 65.4 0.2 | 350.8 1.4 | 48.6 0.1 | 451.1 0.6 |
| LCG-SFNO (ours) | 5.24M | 93.3 0.1 | 82.1 1.4 | 73.5 0.1 | 286.5 0.9 | 51.7 0.2 | 430.4 1.4 |
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
| Task | Degrees (L/R/O) | e3nn | SSHCG | cuEquivariance | |
|---|---|---|---|---|---|
| MACE-large | 128 | 0–2 / 0–3 / 0–3 | 2.02 | 1.04 | 2.30 |
| MACE-medium | 128 | 0–1 / 0–3 / 0–2 | 0.774 | 0.345 | 0.399 |
| NequIP-LiPS | 32 | 0–2 / 0–2 / 0–2 | 1.58 | 0.665 | 1.32 |
| NequIP-aspirin | 64 | 0–1 / 0–1 / 0–1 | 0.516 | 0.175 | 0.225 |
| NequIP-toluene | 64 | 0–2 / 0–2 / 0–2 | 2.18 | 0.756 | 2.79 |
| NequIP-benzene | 64 | 0–3 / 0–3 / 0–3 | 7.94 | 2.83 | 17.3 |
| Workload | e3nn | SSHCG | cuEquivariance | |
|---|---|---|---|---|
| S1 | 4 | 1.47 | 0.0738 | 7.59 |
| S1 | 8 | 7.47 | 0.226 | 180 |
| S1 | 12 | 27.3 | 0.832 | TO |
| S2 | 8 | 1.17 | 0.119 | 13.6 |
| S2 | 16 | 4.55 | 0.284 | 129 |
| S2 | 24 | 8.74 | 0.472 | TO |
| Benchmark | Model | Params | Architecture (epochs run) |
|---|---|---|---|
| Shallow water | ClimaX | 1.96M | patch 4, width 128, depth 6, 8 heads |
| FourCastNet | 1.66M | patch 4, width 128, depth 8, 8 AFNO blocks | |
| SHNet (GSNO) | 0.87M | width 8, scale factor 4 | |
| SFNO | 0.28M | width 28, 4 blocks, , scale factor 2 | |
| LCG-SFNO | 0.27M | width 25, 4 blocks, , CG band 12/24 rank 4, quadratic activation, oversampling 1.5 | |
| WeatherBench | ClimaX | 5.40M | patch 4, width 256, depth 6, 8 heads, decoder depth 2 |
| Model | ACC | MSE | Epochs | Equivariance | |
|---|---|---|---|---|---|
| A | SFNO, benchmark configuration | 84.0 0.0 | 187.7 0.2 | never | |
| B | + full spectral band, two-layer codec, regularization | 93.4 0.1 | 81.5 0.6 | 39 | |
| C | + CG coupling, activations on the data grid | 93.7 0.0 | 77.9 0.1 | 13 | |
| D | + dealiased grid, band-limited residual stream | 93.4 | 80.7 | 21 |
| Blocks | Codec | Equivariance | Step |
|---|---|---|---|
| 1 h | 5 h | 10 h | |||||||
|---|---|---|---|---|---|---|---|---|---|
| Model | vorticity | divergence | vorticity | divergence | vorticity | divergence | |||
| ClimaX | 5.2 0.0 | 259 1.0 | 1000 0.0 | 4.4 0.0 | 417 15 | 1000 0.1 | 3.8 0.0 | 600 31 | 1000 0.2 |
| FourCastNet | 3.0 0.0 | 101 3.1 | 750 4.0 | 6.3 0.8 | 635 70 | 1418 125 | 6.9 1.1 | 1927 398 | 1992 245 |
| SHNet (GSNO) | 1.4 0.1 | 147 6.4 | 279 16 | 3.9 0.8 | 419 46 | 750 19 | 7.5 7.2 | 783 314 | 1085 250 |
| SFNO | 0.6 0.0 | 116 1.4 | 58 4.2 | 2.5 0.2 | 191 1.8 | 169 16 | 4.7 1.7 | 311 17 | 340 49 |
| LCG-SFNO (ours) | 0.6 0.0 | 115 0.5 | 57 4.1 | 2.5 0.2 | 189 3.1 | 166 18 | 4.5 2.1 | 289 6.8 | 338 56 |
| Model | 2T | 10U | 10V | U600 | V600 | T600 | Average |
|---|---|---|---|---|---|---|---|
| One day | |||||||
| ClimaX | 94.3 | 91.4 | 91.5 | 93.3 | 92.2 | 96.7 | 93.2 |
| FourCastNet | 94.9 | 90.8 | 90.8 | 92.6 | 91.3 | 96.5 | 92.8 |
| SHNet (GSNO) | 97.1 | 89.5 | 89.4 | 91.7 | 90.2 | 96.5 | 92.4 |
| SFNO | 95.3 | 76.5 | 74.9 | 83.8 | 80.1 | 93.3 | 84.0 |
| LCG-SFNO | 97.4 | 90.6 | 90.3 | 92.9 | 91.6 | 97.1 | 93.3 |
| Model | 2T | 10U | 10V | U600 | V600 | T600 | Average |
| One day | |||||||
| ClimaX | 11.8 | 89.9 | 136 | 68.5 | 144 | 9.87 | 76.8 |
| FourCastNet | 9.66 | 96.0 | 146 | 75.3 | 160 | 10.5 | 82.8 |
| SHNet (GSNO) | 6.21 | 109 | 167 | 84.5 | 179 | 10.7 | 92.7 |
| SFNO | 9.93 | 228 | 367 | 158 | 344 | 19.8 | 188 |
| LCG-SFNO | 5.55 | 97.3 | 154 | 72.1 | 155 | 8.86 | 82.1 |
Equivariant Neural Networks
Equivariant neural networks encode transformations such as rotations, translations, and permutations so model outputs change consistently with their inputs.
See all 162 papers