ReViT: Rotational-equivariant Vision Transformers for Neural PDE Solvers

Summary

ReViT proposes a Vision Transformer that is rotationally equivariant on grid-based physical fields — an under-explored corner of the equivariance literature, which has focused mostly on point clouds and molecular graphs. The core move is local canonicalization: from each patch derive an $\mathrm{SO}(d)$ frame $\mathbf{B}_i$ from macroscopic reference fields (mean velocity and vorticity in 3D; image gradients on RotMNIST) via a light analytical construction (AOV) or a strain-tensor eigendecomposition (BOA), project vector features into that frame to obtain invariant tokens, run standard multi-head self-attention, and lift back through $\mathbf{B}_i$ at the decoder. A hierarchical Swin-style backbone is made compatible via a "reference basis pyramid" that globalises features before pooling/interpolation and re-localises afterward, plus a rebased relative positional encoding that keeps directional anisotropy while remaining invariant to global rotations. The design cleanly separates equivariance from feature learning and is architecture-agnostic — a genuine conceptual contribution beyond canonicalization work that has mostly lived in the graph/point-cloud world.

The empirical range is strong (RotMNIST, 2D advection vs graph baselines, 2D Kolmogorov flow, 3D MHD, 3D turbulent channel flow), and the exact-equivariance backbone reproduces cleanly: the released ReViT code recovers a rotational-equivariance defect at float64 noise ($\sim 10^{-10}$) across the chiral octahedral group $O$, and ReViT-side MSE on MHD and TCF replicates within a few percent. The main methodological limitation is that the "exact equivariance" story is more fragile than the paper acknowledges: the non-equivariant fallback for collinear/quiescent patches, the underspecified frame of the decoder's canonical query grid, and the reliance on a bespoke basis-aware self-attention for RotMNIST all undercut the strictness claim. The 3D continuous-rotation story is inferred rather than tested, several headline percentages do not reconcile with the appendix tables, and the release ships no competing baselines — every "ReViT beats X" ranking currently depends on paper-supplied numbers. The idea is a real contribution, but the paper trades away rigor at the very seams where the equivariance guarantee matters most.

Strengths

• Conceptual contribution. Bringing local canonicalization into grid-based Vision Transformers, and doing it in a way that keeps the standard optimized multi-head self-attention intact, is a genuinely useful move. The construction is architecture-agnostic and the reference basis pyramid is an elegant way to make merging/expansion compatible with per-patch bases.
• Clear diagnostic of why standard ViTs fail. Section 3.1's separation into the tokenization barrier, positional-encoding failure, and Schur-lemma representational mismatch usefully organises the equivariance problem for vector fields on grids and would be citable in its own right.
• Exact equivariance holds up under execution. Running the shipped scripts/test_equivariance.py on the released ReViT-3D checkpoint yields a rotational-equivariance defect of $\sim 9.7 \times 10^{-11}$ (worst relative error $\sim 2 \times 10^{-10}$) over all 24 elements of $O$; the per-angle rollout MSE on MHD (0.0079) and TCF (0.00179) is identical to $\sim 6$ significant figures across every rotation, exactly as perfect equivariance predicts. The invariant-embedding claim (MSE $= 0$ for rebased features) similarly reproduces at $\sim 10^{-15}$.
• Breadth of empirical evaluation. RotMNIST, 2D advection, 2D KF, 3D MHD, and 3D TCF span classification, scalar rollout, and vector rollout, plus a comparison surface that includes group-equivariant ViTs, GNN-equivariant baselines, L-GATr, and modern PDE transformers. The KF orthogonal-pair analysis (Figure 5) is a particularly clean lens on equivariance vs augmentation.
• Theoretical framing of the approximate-vs-exact story. Proposition G.1 and the $O$ / $SO(d)$ separation give the reader a clear vocabulary for what is guaranteed and what is not, and the KF sweep at $\pi/12$ increments actively investigates that gap rather than hiding behind grid-aligned rotations.
• Compute-favourable design. Avoiding the group-lifting that inflates GSA/GE-ViT memory (1.81 GB vs 95.5 GB reported for $R_{12}$) is a real practical advantage, and — with the caveats on capacity control below — the 3D results with a $\sim 3.6$M parameter model are notably compact.

Weaknesses

• Non-equivariant fallback silently breaks the exactness claim. The Axis Fallback for collinear/Beltrami patches (Appendix C.1.3) crosses $\bar{\mathbf{u}}_i$ with a globally fixed axis $\mathbf{e}_x$ (then $\mathbf{e}_y$), which does not co-rotate with the input. On any patch where the fallback triggers, $\mathbf{B}_i \mapsto \mathbf{R}\mathbf{B}_i$ fails and the exactness guarantee ceases to hold — yet the paper's summary and Proposition 4.1 never restrict the claim to non-degenerate patches. This is the biggest gap between the story ("hard-coded equivariance") and the algorithm as pseudocoded.
• Underspecified decoder equivariance. The claim that $\mathbf{z}_i^{\text{local}} = \mathrm{CrossAttn}(\mathbf{Q}_{\text{grid}}, \mathbf{h}_i^{(L)})$ is invariant requires the fixed canonical grid $\{\boldsymbol{\xi}_m\}$ to live in the patch-local frame. The text never states this; if $\mathbf{Q}_{\text{grid}}$ is in the global frame the whole equivariant-decoding argument collapses. This is a Section 4.4 detail that is doing load-bearing work and should not be implicit.
• Basis-aware self-attention is a hidden ingredient for the RotMNIST result. Table 1 uses an additional RoPE-style rebasing $\mathbf{R}_{ij} = \mathbf{B}_i^\top \mathbf{B}_j$ that is only described in Appendix F.1 and never ablated. The main-body architectural comparison therefore conflates ReViT-proper with a benchmark-specific extension.
• Fair-comparison problems in the empirical suite. The RotMNIST baselines are reported at 44.67k parameters while ReViT is 2.75M ($\sim 60\times$); L-GATr is reduced from 3.4M to 0.034M ($\sim 100\times$) before being called "suboptimal"; UNet3D is capped at 3.5M based on an unsupported convergence claim while other 3D baselines run at $\sim 10$M; and the advection comparison evaluates ReViT on a bilinearly-interpolated regular grid while GNN baselines run on the native mesh. Individually each choice is defensible; collectively they mean the accuracy contrasts are not parameter- or representation-controlled.
• 3D continuous-rotation robustness is not evaluated. All 3D results use the grid-aligned group $O$, exactly where ReViT is mathematically forced to $R^2 = 1$. The approximate-$SO(3)$ story is only tested in 2D (KF, $\pi/12$ sweep). It is not clear whether the "hard constraint" narrative survives non-grid-aligned rotations in 3D, especially given the fallback issue above.
• Statement/table inconsistencies undermine the numbers. The Schur's Lemma statement is not universally valid for the 2D vector representation over $\mathbb{R}$; the $R^2$ formula's denominator is degenerate as written; Table 1's GSA-Nets timings are non-monotonic in $|H|$ against the paper's own $O(N^2|H|^2)$ argument; the KF2D "$56.6\%$ MSE reduction" does not reconcile with Table 7's $0.213 \to 0.136$ ($\sim 36\%$); the coincidence of $11.6\%$ as both a runtime overhead and an accuracy improvement warrants a double-check; P3D appears at both $11.2$M and $11.5$M parameters; TCF improvement is quoted at both $65\%$ and $96.3\%$ without a shared baseline label. None fatal in isolation, collectively they erode trust.
• Boundary-effect mitigation is asserted, not shown. The shifted-window attention "averaging out" and hierarchical merging "absorbing" boundary inconsistencies (Section 4.5) is intuition rather than a controlled measurement; the $\epsilon_{\text{resamp}}$ bound in Appendix G is about interpolation, not patch-boundary crossings.
• Minor conceptual and wording overreach. "First Transformer natively rotationally equivariant for grid-based scalar and vector fields" ignores GSA/GE-ViT on the scalar side (which the paper itself concedes); the conclusion characterises ReViT as enforcing symmetry "as a hard constraint" without qualifying that continuous rotations are only approximately equivariant; the $30°$ example claimed as "maximising" grid mismatch contradicts Figure 17 (which peaks at $45°$ offsets); the AViT axial-attention description conflates spatial axes with channels; "same parameter count ($\sim 3.2$M vs $\sim 3.6$M)" is internally contradictory.
• Missing displayed equations in Appendix C. Multiple sentences in Appendix C.2 and C.3 end with a colon that would ordinarily introduce a displayed formula (strain tensor definition, patch averaging, eigendecomposition, projection, chirality enforcement, structure tensor construction, mean-gradient AOG basis), but no formula appears in the rendered text — if these are genuinely absent from the PDF, the mathematical specification of BOA and the scalar-field adaptations is incomplete.

Reproducibility & code

Execution was based on the real implementation at github.com/tum-pbs/revit; the paper's project-website workspace itself contains only static assets, but the linked repo does run once the gaps below are patched. Overall, ReViT's own numbers reproduce closely; the missing pieces are all on the comparative-baseline side.

• Equivariance backbone reproduces at machine precision. scripts/test_equivariance.py confirms $R^2 = 1$ across all 24 elements of $O$ for ReViT-3D (defect $\sim 9.7 \times 10^{-11}$), and the shipped ReViT-3D checkpoint yields per-angle rollout MSE $= 0.007913$ across every MHD rotation and $= 0.001789$ across every TCF rotation (spread $\sim 10^{-8}$). The BOA/AOV rebasing is likewise numerically zero to $\sim 10^{-15}$.
• ReViT's own accuracy reproduces closely; baseline halves do not. Reproduced ReViT-3D MSE tracks the paper within a few percent (MHD $0.79$ vs $0.82$; TCF $0.18$ vs $0.21$, in $10^{-2}$ units), and using the paper's AViT baseline the MHD-headline arithmetic recomputes to $\sim 64\%$. But no baseline (AFNO, P3D, UNet3D, Swin3D, AViT, PDEArena UNet, GSA-Nets, GE-ViT, PDETrans/PDETrans-Aug, L-GATr) is included in the release, so every "ReViT beats X" ranking — Table 2's leaderboard, the $R^2 = 1$ vs $\approx 0$ polar figure, the "up to $65\%$" and "$63\%$ over AViT" headlines, the PDEArena-UNet comparison, and the KF augmentation studies — is only checkable on ReViT's side.
• RotMNIST comparison is not turnkey. The repository has no training entry point for classification; Train_rollout.py handles only rollout tasks, and a standalone classification driver had to be written around the shipped EqVisionTransformerClassification module. That from-scratch driver reached $97.75\%$ (paper $98.26\%$; $\sim 0.5$ pp gap consistent with training-recipe drift), which drops ReViT just below GE-ViT ($98.01\%$) and GSA-Nets $R_{12}$ ($97.97\%$) if the paper's baseline numbers are trusted. The memory/speed comparison ($\sim 4\times$, $\sim 53\times$) could not be exercised at all — no baseline code, no profiling harness.
• KF2D training data is only shipped for angle 0. The KF2D trainer hard-requires pre-rotated data files ({base}_{split}_V_{angle}.npy) at every augmentation angle, but the HuggingFace release provides only angle 0; without a data-prep script (written in this run using the repo's own rotate_vector_field), the trainer fails out of the box. This blocks Table 3 / Figure 5 for a naive reproducer.
• 2D Advection data is unreleased. The config helper points to adv_*_uniform_grid.pkl files absent from HuggingFace (README marks them "coming soon", confirmed via the HF API), so no ADV_2D number in Section 5.3 — the $\sim 37.2\%$ MSE reduction over PDETrans, the $\sim 11.6\%$ compute overhead, the stability curve, and the graph-vs-transformer comparison — can be exercised end-to-end.
• KF2D augmentation policy is opaque. "Random augmentation angle" leaves undefined the distribution, sampling frequency, and effective dataset size; Table 7's $\times k$ column is likewise ambiguous. Since PDETrans-Aug is the main non-equivariant comparison, none of the KF-augmentation claims are reproducible without this policy in code or a table.
• KF equivariance-error metric is not written down. The "relative difference in MSE" that drives the $+0.0\%$, $+162.8\%$, and $130.8\%$ headline numbers in Section 5.4 is defined only in prose; a reader cannot reproduce or interpret the sign convention.
• Per-seed variance is under-documented. Table 2 lists ReViT std as $\pm 0.00$ on both MHD and TCF; the reproduced per-angle MSE is genuinely constant to $\sim 6$ significant figures on a single checkpoint, but the paper does not say whether $\pm 0.00$ is rounding or a genuine near-machine-precision seed spread. Releasing per-seed values would resolve this.
• ReViT-3D configuration is not derivable from the text. "Halve the depth configuration $[2,4,8,4,2]$" alone does not land on the reported $\sim 3.6$M parameters — widths, decoder capacity, and heads matter. A tabulated per-stage configuration is needed for anyone reimplementing.
• 3D efficiency and architectural-variant sweeps are un-runnable. Table 8 (3D benchmark efficiency) and Table 10 (Tri/LQ $\times$ att/mean $\times$ $M \in \{2,8\}$) rely on timing/memory/GFLOPs for ReViT and every baseline; the release ships no profiling harness and no baseline checkpoints, so none of those cells can currently be reproduced.

Recommended Changes

Essential

• Fix the fallback so exactness is not silently violated. Replace the fixed-global-axis Axis Fallback (Appendix C.1.3) with an equivariant tie-break (e.g., a strain eigenvector, or a rotation-equivariant projection derived from another co-rotating field). If this is not possible, restrict Proposition 4.1 and every "exact $O$-equivariance" statement to the non-degenerate patch set and report the fraction of patches that trigger the fallback on each benchmark.
• State the frame of the decoder's canonical query grid. In Section 4.4, add a sentence (and, if needed, an equation) making it explicit that $\{\boldsymbol{\xi}_m\}$ and its Fourier features live in the patch-local frame; the equivariance of $\mathbf{u}'_i = \mathbf{B}_i \mathbf{z}^{\text{local}}_i$ hinges on this and it is currently implicit.
• Ablate and disclose the basis-aware self-attention on RotMNIST. Report Table 1 accuracy both with and without the RoPE-style $\mathbf{R}_{ij} = \mathbf{B}_i^\top \mathbf{B}_j$ mechanism, and either introduce it in the main-body architecture or clearly label it a RotMNIST-specific extension.
• Fix the $R^2$ metric equation. Rewrite the denominator so it is not $\sum_n |g \cdot f(\mathbf{u}_n) - g \cdot f(\mathbf{u}_n)|^2$ — use $\sum_n |g \cdot f(\mathbf{u}_n) - \overline{g \cdot f(\mathbf{u}_n)}|^2$ (or state the intended definition explicitly) in both Section 4.5 and Appendix H.
• Reconcile the tables/numbers that don't match. In particular: KF2D "$56.6\%$ MSE reduction" vs Table 7's $0.213 \to 0.136$; TCF "$65\%$" vs "$96.3\%$"; P3D $11.2$M vs $11.5$M; whether the $11.6\%$ 2D coincidence is genuine. Please either fix or state the baseline explicitly for each percentage.
• Level the empirical comparisons. Report parameter-matched RotMNIST (either scale baselines up or ReViT down); run L-GATr at a competitive parameter budget on Adv (or explicitly caveat the reported L-GATr is $\sim 100\times$ smaller); run UNet3D at $\sim 10$M with evidence for the "convergence difficulties" claim, or explain why $3.5$M is fair; and quantify the bilinear-interpolation error floor on the Adv regular-grid pipeline so the reader can separate representation from architecture.
• Add a 3D continuous-rotation evaluation. Test MHD (which is periodic) at non-grid-aligned angles with periodic domain expansion analogous to KF, and report either an isotropy-averaged metric or a $\pi/12$ sweep. Without this, the "hard constraint" narrative in 3D rests entirely on grid-aligned rotations where equivariance is mathematically forced.
• Release the baseline models and the missing scientific tooling. Add (i) ReViT baseline implementations/checkpoints (AFNO, P3D, UNet3D, Swin3D, AViT, GSA-Nets, GE-ViT, PDETrans/PDETrans-Aug, L-GATr, PDEArena UNet) needed for Tables 1, 2, and Figure 3; (ii) a train_classification.py (or equivalent) so the RotMNIST result is reachable without an ad-hoc driver; (iii) the ADV_2D .pkl inputs (or a deterministic preparation script from the REMUS mesh data); (iv) a data-prep script that produces the pre-rotated KF2D {base}_{split}_V_{angle}.npy files the trainer requires; and (v) an inference/train-time profiling harness for Tables 8 and 10.

Suggested

• Document the PDETrans-Aug policy and Table 7 x-factor. Tabulate the angle distribution, sampling frequency, periodic-padding treatment, and effective dataset size for each augmentation variant.
• Release per-seed metric values behind Table 2's $\pm 0.00$ std. If the near-zero spread is genuine, that is independently notable and worth foregrounding rather than hiding behind rounding.
• Tabulate the full ReViT-3D configuration. Per-stage depths, embed dims, heads, and decoder settings, so the $\sim 3.6$M parameter count is independently derivable.
• Qualify or drop the "first for grid-based scalar and vector fields" claim. Restrict to vector fields (which is the real contribution) or to "the first grid-based Transformer that is equivariant while natively handling both scalar and vector inputs."
• Rewrite the Schur's Lemma paragraph. State that $\mathbf{W} = \lambda \mathbf{I}$ is the real-type case ($d=3$); note that the 2D vector representation of $\mathrm{SO}(2)$ is of complex type with a two-parameter commutant $a\mathbf{I} + b\mathbf{J}$. The broader "unconstrained $\mathbf{W}$ breaks equivariance" argument is unaffected.
• Correct the absolute-PE equation. Replace $\mathbf{P}(\mathbf{R}(g)x) = \pi_g \mathbf{P}(x)$ with the intended non-equality, since the equality as written implies the opposite of the sentence's conclusion.
• Support or downgrade the boundary-effect claim. Replace "averaging out" and "absorbing" with a controlled measurement (e.g., equivariance defect as a function of hierarchical depth or window shift), or explicitly label the paragraph as intuition.
• Soften the "proving" and "hard constraint" language. Replace "proving that hard-coded equivariance is effective for grid-based neural PDE solvers" with "demonstrating on KF that..." and qualify the conclusion's "hard constraint" to reflect the exact-$O$ / approximate-$SO(d)$ split.
• Clarify the intra-patch set-attention aggregator. State explicitly that the local set-attention omits positional encoding, and clarify whether it is a distinct block from the main W-MSA.
• State the 3D handedness convention explicitly. Write $\mathbf{b}_2 = \mathbf{b}_3 \times \mathbf{b}_1$ in the main body — the sign convention is load-bearing.
• Write down the KF equivariance-error metric formula in Section 5.4 so the $+0.0\%$, $+162.8\%$, and $130.8\%$ percentages can be interpreted and reproduced.
• Fix the smaller nits. Correct the $30°$ example in Section E.2 (Figure 17 peaks at $45°$ offsets); fix the AViT parenthetical "(channels)" to "(spatial axes)"; adjust "same parameter count ($\sim 3.2$M vs $\sim 3.6$M)" to "comparable"; either state $\mathbf{b}_2 = \mathbf{b}_3 \times \mathbf{b}_1$ explicitly in the main body or move Algorithm 1 up into Section C.1.4 where the text refers to it; reconcile the BOA singularity criteria in Section C.2.3 with what Algorithm 1 actually checks; verify the displayed equations in Appendix C.2/C.3 render in the compiled PDF; cite the actual originators of invariant graph networks (Maron et al.) and Tensor Field Networks (Thomas et al.); and exclude AFNO from the "$R^2 \geq 0.6$" statement given its $\pm 1.00$ std.
• Report the fallback activation fraction and a sensitivity analysis to $\epsilon$. The $\sim 8\%$ magnetic-vs-velocity difference is drawn from a single MHD regime; documenting how the degenerate-patch fraction varies with $M_A$, $M_S$ (and other flow regimes) strengthens the generality argument and directly interacts with the exactness fix above.