Subgroups of $U(d)$ Induce Natural RNN and Transformer Architectures

1. 문제 제기 (Problem)

현대 시퀀스 모델 (RNN, Transformer 등) 은 업데이트 파라미터화, 안정화 레이어, 메모리 메커니즘, 읽기 (readout) 관례 등 수많은 상호작용하는 아키텍처 선택 사항들로 구성되어 있습니다. 기존 연구들은 주로 유클리드 공간 ($\mathbb{R}^n$) 에 숨겨진 상태 (hidden state) 를 두고, 안정성을 위해 연산자 (transition operator) 만을 직교 (orthogonal) 또는 유니타리 (unitary) 로 제한하는 방식을 취했습니다 (예: Orthogonal RNN).

이 논문은 이러한 접근법의 한계를 지적하며, 숨겨진 상태와 토큰 표현 자체를 $U(d)$ (d 차원 유니타리 군) 의 닫힌 부분군 (closed subgroup) 에 직접 위치시키는 것이 더 근본적이고 자연스러운 접근법일 수 있다고 주장합니다. 즉, 상태 공간 자체를 군 (Group) 구조로 정의함으로써 동역학의 유계성 (boundedness) 과 기하학적 일관성을 보장하고자 합니다.

2. 방법론 (Methodology)

2.1. 기본 프레임워크 및 공리

저자는 $U(d)$ 의 닫힌 리 부분군 (closed Lie subgroup) $G$와 그 리 대수 (Lie algebra) $\mathfrak{g}$를 기반으로 한 최소한의 공리 체계를 제시합니다.

• 인과성 (Causality): $H_{t+1}$은 $(H_t, x_t)$에만 의존합니다.
• 군 닫힘 (Group closure): $H_t \in G$이면 $H_{t+1} \in G$여야 합니다.
• 국소 업데이트 (Local update): 각 단계는 리 대수 $\mathfrak{g}$의 접선 (tangent) 요소에 의해 생성됩니다.

구체적으로, 상태 업데이트는 지수 사상 (exponential map) 을 사용하여 다음과 같이 정의됩니다:
$$H_{t+1} = H_t \exp(\Phi_\theta(H_t, M_{x_t}))$$
여기서 $\Phi_\theta$는 $H_t$와 토큰 표현 $M_{x_t}$로부터 리 대수 $\mathfrak{g}$의 원소를 생성하는 매핑입니다. 이 방식은 상태가 항상 군 $G$ 내에 머무르도록 보장합니다.

2.2. 공통 템플릿 (Shared Skeleton)

이 프레임워크는 RNN 과 Transformer 모두에 적용 가능한 공통 템플릿을 제공합니다. 부분군의 선택 ($G$) 은 상태 공간, 접선 공간 ($\mathfrak{g}$), 그리고 투사/업데이트 연산자로만 작용하며, 외부 모델 인터페이스는 변경되지 않습니다.

• RNN 템플릿: 순차적으로 접선 업데이트를 계산하고 지수 사상을 통해 군으로 매핑합니다.
• Transformer 템플릿:
  • 주의 (Attention): 토큰 간 유사도 점수를 계산할 때, 유클리드 내적이 아닌 **군 내적 (matrix inner product)**인 $\Re \text{tr}((H_i)^* H_j)$를 사용하여 좌표 불변성 (coordinate-invariant) 을 확보합니다.
  • 업데이트: 주의 (attention) 단계와 그라운딩 (grounding, 토큰 임베딩 반영) 단계를 모두 접선 공간에서의 업데이트로 처리한 후 지수 사상을 적용합니다.

2.3. 접선 혼합 (Tangent Mixing)

원래의 기하학적 구조를 유지하는 항등 매핑 ($\Gamma = \text{Id}$) 외에도, **선형 접선 혼합 (Linear Tangent Mixing)**을 도입했습니다. 이는 접선 좌표계에서 학습 가능한 선형 매핑을 적용하여, 예측 (predictive) 채널과 메모리 (memory) 채널 간의 용량을 재분배하고 업데이트를 재가중치하는 방식입니다. 이는 유한한 예산 (parameter budget) 내에서 성능을 향상시키는 핵심 요소로 작용합니다.

2.4. 실험 설정: 직교군 $O(d)$

논문은 구체적인 실험을 위해 직교군 $O(d)$를 선택했습니다.

• 상태 및 토큰 표현: $O(d)$의 원소로 매개변수화됨 ($M_v = \exp(\text{skew}(B_v))$).
• 접선 투사: $\Pi_g(Y) = \frac{1}{2}(Y - Y^\top)$ (반대칭 행렬로 투사).
• 읽기 헤드 (Readout): 군 고유의 유사도 기반 프로토타입 헤드를 사용 ($\ell_v(H) = \tau \Re \text{tr}(H^* P_v) + b_v$).

3. 주요 기여 (Key Contributions)

1. 통일된 프레임워크: $U(d)$의 닫힌 부분군을 기반으로 한 RNN 과 Transformer 의 공통 템플릿을 제안했습니다. 부분군 선택이 상태 공간, 접선 투사, 업데이트 맵을 결정하는 '드롭인 (drop-in)' 구성 요소 역할을 합니다.
2. 상태 공간의 군화: 기존 연구가 연산자만 제한했던 것과 달리, 숨겨진 상태와 토큰 표현 자체를 군 원소로 직접 정의하여 본질적으로 유계된 동역학을 확보했습니다.
3. 구체적 실증 ($O(d)$): 직교군 $O(d)$를 인스턴스화한 OSM (Orthogonal State Model) RNN 과 OSMFormer를 제안하고, 파라미터 수를 맞춘 (parameter-matched) 조건에서 기존 모델 대비 우수한 성능을 입증했습니다.
4. 선형 혼합의 효과: 접선 공간에서의 학습 가능한 선형 혼합이 $O(d)$ 설정에서 유한한 파라미터 예산 하에서 최고의 성능을 제공함을 보였습니다.

4. 실험 결과 (Results)

Tiny Shakespeare 와 Penn Treebank (PTB) 데이터셋에서 문자 단위 언어 모델링 (Character-level LM) 을 수행하여 비트 퍼 캐릭터 (BPC, 낮을수록 좋음) 를 평가했습니다.

• Tiny Shakespeare (약 100K 파라미터):
  • OSMFormer (USIM, 선형 혼합): 검증 BPC 2.254 / 테스트 BPC 2.464
  • 기존 Transformer (ALiBi): 검증 BPC 2.259 / 테스트 BPC 2.583
  • 결과: OSMFormer 는 파라미터 수가 적음 (91K vs 104K) 에도 불구하고 더 낮은 BPC 를 기록했습니다.
• Penn Treebank (약 100K 파라미터):
  • OSMFormer: 검증 BPC 1.665 / 테스트 BPC 1.614
  • 기존 Transformer: 검증 BPC 1.716 / 테스트 BPC 1.664
  • 결과: OSMFormer 가 두 지표 모두에서 우위를 점했습니다.
• 확장성 (Scaling): 100K 에서 500K 파라미터까지 스케일링 실험에서 OSMFormer 는 모든 예산 구간에서 베이스라인 Transformer 보다 일관되게 낮은 BPC 를 기록했습니다.
• RNN 비교: OSM-RNN (선형 혼합) 이 LSTM 보다 Tiny Shakespeare 에서 더 좋은 테스트 BPC (2.539 vs 2.594) 를 보였습니다.
• 옵티마이저 견고성: 다양한 학습률 및 클리핑 설정에서 OSMFormer 는 베이스라인보다 더 넓은 최적 검증 범위 내에서 안정적으로 학습되었습니다.

5. 의의 및 결론 (Significance & Conclusion)

이 논문은 시퀀스 모델링을 위해 기하학적 구조 (군 이론) 를 모델의 핵심 설계 원칙으로 통합하는 새로운 패러다임을 제시합니다.

• 방법론적 혁신: 상태 공간을 유클리드 공간이 아닌 리 군 (Lie Group) 으로 정의함으로써, 안정성 (bounded dynamics) 과 기하학적 일관성을 자연스럽게 확보할 수 있음을 보였습니다.
• 실용성: $O(d)$ 인스턴스화를 통해 기존 모델과 동등하거나 더 나은 성능을 달성했으며, 특히 선형 접선 혼합이라는 간단한 확장이 성능 향상의 열쇠가 됨을 입증했습니다.
• 미래 방향: 이 프레임워크는 $SU(d)$, $T^k$(토러스) 등 다른 부분군으로 확장 가능하며, 해석 가능성 (접선 공간에서의 업데이트 추적), 안정성 분석, 그리고 장기 기억 (long-term memory) 연구에 대한 새로운 축을 제공합니다.

결론적으로, 이 연구는 "부분군 선택 = 상태 공간 및 업데이트 규칙 선택"이라는 직관적인 설계 원칙을 통해 RNN 과 Transformer 를 재정의하고, $O(d)$ 기반 모델이 파라미터 효율성과 성능 면에서 강력한 대안이 될 수 있음을 증명했습니다.