Accelerating SO(2)-Equivariant Materials Models through Low-Precision Emulation

1 INTRODUCTION

• Models typically reach 110M parameters, datasets approach 100M structures, and scaling laws indicate further model-size increases can improve performance.Higher-rank targets such as Hamiltonians also increase available training data by several orders of magnitude over scalar properties.Source: §1 · p. 1
• SO(2) layers account for more than 80% of inference steps, while reduced-precision formats can introduce rotation-dependent error and BF16 degrades UMA accuracy by 20% to 50%.This leaves symmetry-constrained models relying on FP32 or FP64 general-purpose units rather than MMUs.Source: §1 · p. 2
• FlashSO(2) uses integer Ozaki emulation to reconstruct high-precision products from exact low-precision slices, with equivariance-constrained slicing, tailored scaling, and fused kernels.It runs SO(2) layers on INT8 MMUs within a 10^-4 equivariance tolerance and improves on TF32 in accuracy, throughput, and memory footprint.Source: §1 · p. 2
• MALOQ trains 3.9× faster than FP32 at 58% lower peak memory, while UMA-M runs inference 3.3× faster at 51% lower peak memory.FlashSO(2) reaches 40× and 7× the reference SO(2)-layer throughput in MALOQ and UMA-L, respectively.Source: §1 · pp. 2–3

2 RELATED WORK

• Rotational symmetry restricts model representations to physically valid forms, while SO(2) convolutions provide a more computationally efficient message-passing engine than full Clebsch–Gordan tensor products.MLIP models tune lmax as a hyperparameter; MLH models set it from the orbital basis, such as lmax = 8 for g-orbital datasets.Source: §2 · p. 3
• Prior scaling work speeds MLIPs through software fusion and compilation or distributed inference, while MLH work targets larger systems, basis sets, batching, and graph partitioning.The core SO(2) layer remains a bottleneck, and cited kernel libraries for SO(3) CGTPs do not cover its dense per-order matrix products.Source: §2 · p. 3
• Figure 2 depicts FP32 matrices split into INT8 slices, selected slice products accumulated in INT32, and scales refined from rows and columns to blocks.The caption specifies 128-entry row blocks and 128 × 128 tiles.Source: Figure 2 · p. 4
• Integer Ozaki emulation splits operands into low-precision slices and combines their products to compute high-precision matrix products on low-precision MMUs.The integer form uses INT8 slices with exact INT32 accumulation; EmuGEMM removes slicing and reconstruction round trips with fused kernels.Source: §2 · p. 4

3.1 INTEGER EMULATION WITH BLOCK SCALING

• Emup splits FP32 operands into p INT8 slices, retains products with i + j ≤ p + 1, and exactly accumulates slice products in INT32 before weighted reconstruction.Emu1, Emu2, and Emu3 require 1, 3, and 6 INT8 MMU products, respectively; Emu2 theoretically keeps 14 bits.Source: §3.1 · pp. 4–5
• Block scaling assigns a scale to each 128-entry block along K, improving bit use for small entries relative to a single scale per row.The example row entry about 1000× smaller than its maximum retains only 4 of 14 bits with row scaling.Source: §3.1 · pp. 4–5

3.2 SO(2) LAYER PIPELINE

• Each SO(2) layer rotates edge features to the global z-axis, applies per-m SO(2) convolutions and a gated activation, then rotates messages back and aggregates them at nodes.The layer uses an l-major to m-major layout transformation before its separate linear operations across m-subspaces.Source: §3.2 · p. 5

3.3 OPTIMIZING PRECISION EMULATION FOR SO(2) LAYERS

• FlashSO(2) fuses the surrounding forward operations into three kernels, keeping intermediates in registers and passing only INT8 slices and scales between kernels.It uses block-scaled INT8 emulation for GEMMs; 1-D scaling addresses activation range and long gradient sums, while weights use 2-D scaling.Source: §3.3 · pp. 5–6, Figure 3 · p. 6
• FlashSO(2) leaves weights and other model modules in FP32, remains compatible with torch.compile and CUDA graphs, and offers MiniMem recomputation to trade speed for memory.MiniMem recomputes the SO(2) layer during the backward pass.Source: §3.3 · p. 6

4.1 EQUIVARIANCE ANALYSIS

• Emu2 stays within the 10−4 equivariance tolerance for both MALOQ and UMA, while no native low-precision format meets it.TF32 lies just outside the threshold; BF16 is two orders of magnitude farther out, as is Emu1.Source: §4.1 · p. 7

4.2 ROOFLINE MODELS OF THE SO(2) LAYER

• 40× throughput in MALOQ and 7× in UMA-L over the native SO(2) layer; Emu2 sustains 236 and 232 effective Tflop/s, respectively.It also provides 6.5× and 4.0× the native layer’s arithmetic intensity, respectively, and exceeds TF32 throughput while meeting the equivariance tolerance.Source: Figure 4 · p. 7, §4.2 · p. 8

4.3 ACCELERATING ML MATERIALS MODELS

• Emu2’s Hamiltonian-element MAE is within 0.2% of FP32 on 32 held-out conformers, while TF32’s is 21.6% higher than FP32.The FP32 and Emu2 training-loss curves coincide after 2,400 epochs; BF16 and Emu1 end at 4.8× and 13× Emu2’s final loss, respectively.Source: §4.3 · p. 8, Figure 5 · p. 9
• Emu2 accelerates MALOQ training steps by up to 3.93× and lowers peak training memory by 58% at the largest input fitting in FP32.MiniMem trades some speed for 78% lower peak memory while remaining 3.33× faster than the FP32 reference.Source: §4.3 · p. 8
• Emu2 accelerates UMA-M training steps by 2.6×; the SO(2) layer accounts for 79% to 91% of each step.Source: §4.3 · p. 9
• Emu2 runs UMA-M inference 3.33× faster than FP32 at 65.5k edges, while MiniMem fits 25.6k atoms (410k edges) in 78 GiB.FP32 inference exhausts the GH200’s 95 GiB at 5.3k atoms.Source: §4.3 · p. 9

• Emu2 matches FP32 accuracy, runs 1.1× faster than TF32, and uses half its memory, improving all three axes over TF32.Source: §5 · p. 9
• An up to 81% lower effective memory footprint enables a 1.7× larger cutoff for the same structure size.Source: §5 · p. 9
• Cross-domain pretraining could leverage over 100 TB of electronic-structure intermediates from roughly 50 GB of energies-and-forces data at 40× reference throughput.Source: §5 · pp. 9–10