ReViT: Rotational-equivariant Vision Transformers for Neural PDE Solvers

Physics obeys strict symmetries like rotational equivariance. However, the standard Transformer architectures widely used in physics foundation models do not enforce these constraints by construction. We introduce ReViT, a rotationally equivariant Vision Transformer framework for neural PDE solvers operating on grid-based physical fields that achieves exact equivariance for the discrete groups $C_4$ (2D) and the chiral octahedral group $O$ (3D), with bounded approximate $\mathrm{SO}(d)$ equivariance for continuous rotations. ReViT maps scalar and vector inputs into locally invariant representations derived from physics-based canonical bases, enabling the use of standard self-attention without symmetry violations. Built on a hierarchical Swin-style backbone with a precomputed reference basis pyramid, ReViT preserves equivariance across multi-scale operations. We evaluate ReViT on a wide range of 2D and 3D PDE benchmarks, such as Magnetohydrodynamics and Turbulent Channel Flows, demonstrating significant gains over state-of-the-art baselines. ReViT exhibits strong generalization, and reduces MSE by up to 65\% compared with the best-performing alternatives.

Many natural phenomena, like air flowing over a wing, magnetic fields in plasmas, water swirling in a channel, obey a simple rule of physics: if you rotate the entire situation, the behavior rotates along with it. Foundation models are increasingly used as fast surrogates for traditional physics simulators, but the standard Transformer architectures behind today's most powerful models do not enforce this rotational consistency. As a consequence, their predictions can change in unphysical ways when the input is rotated, leading to unphysical artifacts and reduced accuracy. We introduce ReViT, a Vision Transformer architecture that guarantees rotational equivariance on regular grids by construction. The key idea is to combine the permutation equivariance inherent to self-attention with rotational invariance obtained through local canonicalization: we re-express the data in each local patch within a coordinate frame that the physics itself picks out, so the model sees the same underlying scene regardless of orientation. We further extend this guarantee across multiple scales of detail, allowing the model to capture both fine and coarse features. Across a range of 2D/3D physics benchmarks, including turbulent flows and magnetized plasmas, ReViT substantially outperforms leading alternatives, reducing prediction error by up to 65\%, while remaining efficient enough to scale. By making physical symmetry a built-in property of the architecture, this work brings AI-based scientific simulation closer to being a trustworthy tool for research and engineering.