REViT: Roto-reflection Equivariant Convolutional Vision Transformer

기술 요약: REViT – 회전-반사 등변 합성곱 비전 트랜스포머 (Roto-reflection Equivariant Convolutional Vision Transformer)

문제 정의
등변 신경망(Equivariant Neural Networks)은 회전 및 반사 대칭이 필요한 작업(예: 분자 분석, 의료 병리, 로보틱스)에서 효과적임이 입증되었으나, 비전 트랜스포머(ViT)에 이를 적용하는 것은 여전히 도전적인 과제입니다. 전통적인 ViT는 공간 구조를 유지하기 위해 절대적 또는 상대적 위치 인코딩(Positional Encoding)에 크게 의존하는데, 이는 군 등변성(Group Equivariance)을 구현하는 것을 복잡하게 만듭니다. 구체적으로, 표준 위치 인코딩은 고정된 좌표에 고유한 벡터를 할당하기 때문에 회전이나 반사와 같은 변환 하에서 등변성을 깨뜨립니다. 상대적 위치 인코딩(RPE)을 사용하는 Group Self-Attention (G-SA)과 같은 기존의 등변 ViT 시도들은 복잡한 어텐션 메커즘 수정이 필요하며 상당한 계산 오버헤드를 초래합니다. 또한, 기존의 회전-반사 등변성에 관한 연구는 주로 합성곱 신경망(G-CNN)에 집중되어 있어, 이러한 작업을 위해 트랜스포머의 전역적 문맥 모델링 능력을 활용하려는 시도에는 공백이 존재합니다.

방법론
저자들은 명시적인 위치 인코딩이 필요 없는 이산 회전-반사 군 등변 아키텍처인 REViT(Roto-reflection Equivariant Convolutional Vision Transformer)를 제안합니다. 핵심 방법론은 두 가지 주요 혁신을 포함합니다:

1. 합성곱 패치 임베딩 및 리프팅(Convolutional Patch Embedding and Lifting):
   선형 패치 임베딩 대신, REViT는 2D 입력 이미지 $f: \mathbb{Z}^2 \to \mathbb{R}^C$를 이산 군 공간 상의 특징 맵 $F: \mathbb{Z}^2 \times C_N \to \mathbb{R}^{C'}$로 변환하는 **리프팅 레이어(Lifting Layer)**를 사용합니다. 이는 입력이 이산 회전 군(예: $p4, p8$)의 원소들에 대응하는 회전 필터들과 합성됨으로써 달성됩니다. 이를 통해 데이터는 유클리드 공간에서 위치-방향 공간($\mathbb{R}^2 \times S^1$)으로 리프팅되며, 학습 가능한 위치 파라미터 없이도 공간 및 방향 정보를 본질적으로 인코딩합니다.

2. 그룹 합성곱 셀프 어텐션 (Group Convolutional Self-Attention, G-CSA):
   표준 셀프 어텐션 메커니즘은 G-CSA로 대체됩니다. 이 모듈에서 쿼리(Query), 키(Key), 값(Value) 투영은 선형 투영 대신 군 등변 합성곱을 사용하여 계산됩니다. 어텐션 메커즘은 공간 차원과 군 차원(방향) 모두에서 작동합니다. 어텐션 가중치는 리프팅된 특징 공간 전반에 걸쳐 쿼리와 키의 내적을 기반으로 계산되어, 연산이 군 작용(회전 및 반사)에 대해 등변성을 유지하도록 보장합니다. 합성곱 투영과 합성곱에 내재된 상대적 공간 관계를 활용함으로써, REViT는 모든 어텐션 헤드마다 상대적 위치 인코딩을 계산해야 하는 per-layer 복잡성 없이 이동 및 회전 등변성을 달 수 있습니다.

주요 기여

• 이산 군 합성곱 셀프 어텐션 (G-CSA): 본 논문은 위치 인코딩에 의존하지 않고 트랜스포머에서 회전-반사($E(2,N)$) 등변성을 달성하기 위한 메커니즘으로서 G-CSA를 도입했습니다.
• REViT 구현: 저자들은 $p4$-CSA 및 $p8$-CSA와 같은 다양한 이산 회전-이동($pN$) 및 회전-반사($pNm$) 군 설정에 대해 G-CSA를 활용하여 REViT 블록을 구현했습니다.
• 확장성: 본 연구는 REViT가 ImageNet-1k와 같은 대규모 데이터셋으로 확장될 수 있음을 보여주며, 다운스트림 등변 작업의 잠재적인 백본(Backbone) 역할을 수행함을 입증합니다.

실험 결과
저자들은 Rotated MNIST, PatchCamelyon, CIFAR-10의 세 가지 데이터셋에서 REViT를 평가하였으며, 이를 ImageNet-1k로 확장했습니다.

• G-SA with RPE 대비 성능: Rotated MNIST 및 PatchCamelyon에서 REViT 모델(특히 $p4$-CSA 및 $p8$-CSA)은 상대적 위치 인코딩이 강화된 Group Self-Attention 모델(G-SA+RPE)과 대등하거나 더 나은 성능을 보였습니다. 결정적으로, REViT는 훨씬 적은 파라미터와 낮은 계산 복잡성(더 적은 곱셈-덧셈 연산 및 감소된 메모리 풋프린트)으로 이를 달 수 있었습니다.
• G-CNN 대비 성능: REViT는 모든 데이터셋과 대칭 군에 대해 대응하는 그룹 등변 CNN(G-CNN) 베이스라인을 일관되게 능가했습니다. 특히 PatchCamelyon에서 REViT는 더 높은 차수의 대칭을 가진 G-CNN 변형 모델들보다 더 적은 파라미터를 사용하면서도 더 높은 정확도를 달성했습니다.
• ImageNet 스케일링: ImageNet-1k에서 $p4m$-등변 REViT는 79.27%의 Top-1 정확도를 기록하여, 증강(Augmentation)을 적용하여 학습된 바닐라 ViT(72.08%)와 등변 ResNet 베이스라인(77.37%)을 모두 앞질렀습니다.
• 등변성 오차: 모델들은 리프팅 및 G-CSA 레이어를 거친 후 거의 제로에 가까운 등변성 오차를 보였으며, 고차 군(예: $p8$)에서의 미세한 오차는 회전 중 발생하는 보간(Interpolation) 아티팩트에 기인했습니다.

의의 및 주장
본 논문은 REViT가 기존 방법들에 비해 비전 트랜스포머에서 이산 회전-반사 등변성을 달 achieve하는 더 단순하고 효율적인 경로를 제공한다고 주장합니다. 명시적인 위치 인코딩에 대한 의존성을 제거하고 합성곱 셀프 어텐션을 활용함으로써, 저자들은 엄격한 대칭 제약을 준수하면서도 트랜스포머의 전역적 문맥 모델링 강점을 유지할 수 있음을 입증했습니다. 결과는 G-CSA가 모델 크기나 복잡성을 크게 증가시키지 않고도 그룹 등변 분류 성능을 개선할 수 있음을 시사합니다. 저자들은 REViT를 방향성이 중요한 작업(예: 의료 영상 및 방향성 객체 탐지)을 위한 실행 가능한 확장 가능 백본으로 포지셔닝하는 한편, 계산 비용이 대칭 군의 카디널리티(Cardinality)에 따라 여전히 스케일링된다는 점을 언급하며 이는 지연 시간에 민감한 애플리케이션에는 도전 과제가 될 수 있음을 인정했습니다.