ReViT: Rotational-equivariant Vision Transformers for Neural PDE Solvers

Motivation

• Known: Group-equivariant CNN, steerable network, tensor field network 등은 discrete 혹은 continuous 회전군에 대한 exact equivariance를 제공하지만 그래프/message-passing 기반이거나 고정 커널로 인해 해상도 증가에 따른 파라미터·연산 비용이 크게 늘어난다. 기존 equivariant Transformer(SE(3)-Transformer, LieTransformer 등)는 point cloud/그래프 데이터에 특화되어 있고, vision 기반 equivariant Transformer는 스칼라 입력 분류 작업에만 국한되어 있다.
• Gap: Grid로 이산화된 물리 데이터에서 스칼라와 벡터 필드를 모두 자연스럽게 다루면서 회전 equivariance를 하드 제약으로 보장하는 Transformer 아키텍처는 현재까지 부재하다. 표준 ViT는 tokenization, positional encoding, self-attention 각 단계에서 회전 대칭을 위반하는 구조적 장벽(C1, C2 등)을 가진다.
• Why: 물리 법칙은 회전에 대해 불변/등변이므로, PDE를 위한 foundation model이 이 대칭성을 위반하면 물리적으로 일관되지 않은 예측을 생성할 위험이 있으며, 이는 계산유체역학 등에서 신뢰 가능한 AI를 구축하는 데 핵심적인 요소로 지적된다.
• Approach: strain rate tensor에서 영감을 받은 rotation-invariant local canonical basis를 patch 단위로 계산해 표준 self-attention을 그대로 사용하면서도 대칭 위반 없이 동작하도록 하고, 이를 Swin 스타일 계층적 backbone과 precomputed reference basis pyramid로 확장하여 multi-scale에서도 equivariance를 유지한다.

Achievement

1. ReViT 아키텍처: grid 기반 스칼라·벡터 필드에 대해 native하게 회전 equivariant한 최초의 Transformer 아키텍처를 제안했다.
2. Architecture-agnostic invariant embedding: 물리 기반 canonical basis로부터 patch별 rotation-invariant local coordinate system을 계산하여, 표준 multi-head self-attention을 그대로 활용 가능한 invariant token을 생성했다.
3. Hierarchical Reference Basis Pyramid: Swin에서 영감을 받은 계층적 backbone에 precomputed reference basis pyramid를 결합해 다중 해상도에서도 엄격한 equivariance를 유지했다.
4. 광범위한 검증: rotated MNIST, 2D advection, 2D Kolmogorov flow, 3D turbulent channel flow, 3D magnetohydrodynamics(MHD) 등 다양한 벤치마크에서 state-of-the-art 대비 최대 65%의 MSE 감소와 강한 일반화 성능을 입증했다.

How

• 물리 필드에 대한 group action을 정식화: 스칼라 필드는 [Lgϕ](x) = ϕ(R(g)^{-1}x), 벡터 필드는 [Lgu](x) = R(g)u(R(g)^{-1}x)로 정의.
• 표준 self-attention의 permutation equivariance는 스칼라 필드에는 충분하지만 벡터 필드에는 불충분함을 이론적으로 증명하고, 이를 유발하는 3가지 메커니즘(C1: Tokenization Barrier, C2: Loss of Spatial Equivariance, 그리고 추가 요인)을 규명.
• Tokenization Barrier(C1)를 해결하기 위해, 각 patch 내부에서 strain rate tensor에 기반한 rotation-invariant local canonical basis를 계산하여 patch 내 벡터/스칼라 정보를 invariant coordinate로 투영한 뒤 이를 토큰화.
• Positional encoding 문제(C2)를 해결하기 위해 회전에 따라 함께 변환되는 relative/equivariant positional encoding 방식을 도입.
• Swin-style 계층적 encoder-decoder 구조에서 각 해상도 단계마다 미리 계산된 reference basis pyramid를 활용해 patch merging/splitting 등 multi-scale 연산에서도 equivariance가 깨지지 않도록 설계.
• discrete group C4(2D), chiral octahedral group O(3D)에 대해서는 exact equivariance를, 연속 SO(d) 회전에 대해서는 bounded approximate equivariance를 이론적으로 보장.

Originality

• 기존 equivariant Transformer 연구가 point cloud/그래프(unstructured data)나 스칼라 분류 작업에 국한되어 있던 것과 달리, grid 구조의 inductive bias를 유지하면서 스칼라와 벡터 필드를 모두 native하게 지원하는 최초의 회전 equivariant ViT를 제시.
• 유체역학의 strain rate tensor 개념을 차용해 patch 단위 canonical basis를 구성함으로써, 별도의 group-lifted attention이나 quadratic complexity를 요구하지 않고도 표준 optimized self-attention을 그대로 재사용할 수 있게 한 아이디어가 독창적이다.
• Swin-style 계층 구조에 precomputed reference basis pyramid를 결합해 multi-scale에서도 equivariance를 엄밀히 보존한 설계는 기존 group CNN/equivariant Transformer가 다루지 못한 확장성 문제를 해결한다.

Limitation & Further Study

• 연속 SO(d) 회전에 대해서는 exact가 아닌 bounded approximate equivariance만 보장되므로, 임의의 연속 회전 각도에서 성능 저하 가능성이 존재할 수 있다.
• discrete group(C4, O)에 한정된 exact equivariance이므로 더 일반적인 연속군이나 다른 물리적 대칭(예: 스케일링, 반사 등)으로의 확장 여부는 추가 검증이 필요하다.
• 평가가 특정 PDE 벤치마크(MHD, 난류 채널 유동 등)에 집중되어 있어, 더 넓은 범위의 물리 시스템이나 실제 관측 데이터에 대한 일반화 가능성은 추가 연구가 필요하다.
• 계층적 reference basis pyramid의 precompute 비용 및 매우 고해상도 3D grid에서의 실제 연산/메모리 효율에 대한 상세한 분석이 본문 발췌에서는 제한적으로 보인다.

Evaluation

총평: 물리 시뮬레이션을 위한 foundation model 구축에 있어 회전 대칭성을 하드 제약으로 보장하는 최초의 grid 기반 Vision Transformer를 제시했다는 점에서 이론적·실용적 기여도가 크며, 다양한 PDE 벤치마크에서의 큰 성능 향상이 이를 뒷받침한다.