# Velocityformer: Broken-Symmetry-Matched Equivariant Graph Transformers for Cosmological Velocity Reconstruction

arXiv: https://arxiv.org/abs/2605.21483 · HTML: https://mlsanity.com/paper/2605.21483/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 55 references that are in the map.
- First posted on arXiv 20 May 2026.

## Builds on (most cited first, 21 of 21)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups](https://mlsanity.com/paper/1802.03690/) (2018)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Geometric and Physical Quantities Improve E(3) Equivariant Message Passing](https://mlsanity.com/paper/2110.02905/) (2021)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [DeepSphere: a graph-based spherical CNN](https://mlsanity.com/paper/2012.15000/) (2020)
- [DeepSphere: Efficient spherical Convolutional Neural Network with HEALPix sampling for cosmological applications](https://mlsanity.com/paper/1810.12186/) (2018)
- [Geometric Algebra Transformer](https://mlsanity.com/paper/2305.18415/) (2023)
- [Scattering Networks on the Sphere for Scalable and Rotationally Equivariant Spherical CNNs](https://mlsanity.com/paper/2102.02828/) (2021)
- [Neural Scaling Laws for Boosted Jet Tagging](https://mlsanity.com/paper/2602.15781/) (2026)
- [Translation and Rotation Equivariant Normalizing Flow (TRENF) for Optimal Cosmological Analysis](https://mlsanity.com/paper/2202.05282/) (2022)
- [Predicting the Thermal Sunyaev-Zel'dovich Field using Modular and Equivariant Set-Based Neural Networks](https://mlsanity.com/paper/2203.00026/) (2022)
- [CosmoBench: A Multiscale, Multiview, Multitask Cosmology Benchmark for Geometric Deep Learning](https://mlsanity.com/paper/2507.03707/) (2025)
- [Geometric deep learning for galaxy-halo connection: a case study for galaxy intrinsic alignments](https://mlsanity.com/paper/2409.18761/) (2024)

## Abstract (arXiv)

Precise measurement of the kinematic Sunyaev-Zel'dovich (kSZ) effect - a probe of the large-scale distribution of baryonic matter, a key observable for cosmological inference - requires accurate reconstruction of galaxy velocities from spectroscopic surveys. The signal-to-noise ratio (SNR) of kSZ measurements scales directly with the correlation coefficient r between reconstructed and true velocities. We introduce Velocityformer, an equivariant graph transformer architecture designed to match the specific symmetry of the observational data. While the underlying physics is equivariant with respect to translations and rotations, observational effects break this symmetry due to the preferred line-of-sight direction. Matching the model's inductive bias to the data's broken symmetry consistently improves performance across all model sizes and training volumes, with Velocityformer improving r by 35% over the standard linear theory baseline and outperforming ML baselines at every data volume. By matching the model's inductive bias to the data and conditioning on the physics-based long-wavelength solution, Velocityformer is highly data-efficient, training to high accuracy on as few as 4 low-fidelity simulations, and generalises zero-shot across input geometry, cosmological parameters, and galaxy sample. On high-fidelity simulated galaxy catalogues, this yields a 30% improvement in r over the physical baseline, directly translating to the same SNR gain on observational data.