FlashSO(2) speeds symmetry-preserving layers with integer emulation
Accelerating SO(2)-Equivariant Materials Models through Low-Precision Emulation · arXiv:2610.35164
In plain words
Some materials models are designed so that turning a structure also turns their predictions in a consistent way. But the layer that handles this rotational symmetry takes much of the computing time, and ordinary low-precision math can break that consistency.
The paper introduces FlashSO(2), which breaks matrix calculations into small integer calculations and combines nearby operations to save time and memory. In tests of MALOQ and UMA outputs, its Emu 2 setting met the stated rotational-error limit; the result is limited to those tested outputs.
The research
Materials models that preserve how predictions should change under rotation spend much of their computation in a costly layer, while ordinary reduced-precision math can introduce orientation-dependent errors. Can that layer use faster integer hardware and still preserve the symmetry these models rely on?
FlashSO(2) emulates precise matrix products on fast integer units and fuses surrounding SO(2) operations to accelerate symmetry-preserving materials models.
The authors tested whether the SO(2) layer could use integer matrix hardware without losing rotational equivariance, the property that outputs transform consistently when inputs rotate.
They measured the largest output error across rotation angles for MALOQ Hamiltonians and UMA forces; Emu 2 stayed within the 10⁻⁴ tolerance for both tested targets.
They then compared the layer with its FP32 reference and measured model-level workloads. Fused Emu 2 raised layer throughput in both tested configurations, while MALOQ training and UMA-M inference also ran faster than their FP32 comparisons under specified conditions.
Together, these results show that the approach can preserve the tested symmetry while improving performance in these models.
These results come from representative MALOQ and UMA models, with most evaluations on one NVIDIA GH200 superchip. The reported model-level speedups apply to specific input sizes and settings, so they do not establish the same gains for other models, hardware, or workloads.
What the evidence shows
- The speed and symmetry problem
FlashSO(2) is a software and hardware co-designed implementation of the SO(2) layer in graph neural network materials models. These models preserve rotational equivariance, meaning their outputs change consistently with rotations of the input.
The layer accounts for much of an inference step in many such models, but native reduced precision can introduce rotation-dependent errors, limiting use of faster matrix hardware.
- How FlashSO(2) works
To address that tension, FlashSO(2) emulates matrix products with scaled INT8 slices and exact INT32 accumulation, then reconstructs the result. Its Emu 2 setting uses two slices.
Block scaling gives groups of entries their own scale, helping retain precision for small activations and avoid overflow in long sums. The implementation also fuses operations around the matrix products, keeping intermediate results in registers rather than moving them between separate kernels.
- Testing the symmetry constraint
The authors checked whether the implementation preserved rotational equivariance by measuring the largest component-wise residual as rotation angles changed. They tested MALOQ Hamiltonian outputs and UMA force outputs against a 10⁻⁴ tolerance.
Only FP32 among native formats met that limit; TF32 fell just outside it, while Emu 2 stayed within it for both tested targets. Emu 3 was indistinguishable from FP32.
- Tracing the layer speed gains
A staged layer comparison separated kernel fusion from integer emulation: the authors compared the reference FP32 layer with fused FP32, then compared fused FP32 with fused Emu 2. Against the reference, fused Emu 2 reached 40x throughput in MALOQ and 7x in UMA-L on GH200. These are layer-level results for the tested configurations, not model-wide speedups.
- Checking gains in models
The authors also measured model-level performance. In MALOQ training under torch. compile, Emu 2 made steps up to 3.93x faster at the largest input that fit in FP32 memory, and reduced peak memory by 58% at the input size where FP32 ran out of memory.
In UMA-M inference, it ran a specified fcc-copper input 3.33x faster than FP32. These results support gains beyond the isolated layer in the tested settings.
Deep
- Original paper authors
- Denghui Lu, Alexander Maeder, Alexandros Nikolaos Ziogas, Mathieu Luisier, Manasa Kaniselvan
- Submitted to arXiv
- License
- arXiv non-exclusive