RotRNN: Modelling Long Sequences with Rotations

기술 요약: RotRNN – 회전을 이용한 긴 시퀀스 모델링

문제 정의
긴 시퀀스 모델링은 장기 의존성(long-range dependencies)을 포착하는 어려움으로 인해 머신러닝 분야에서 여전히 도전적인 영역으로 남아 있습니다. 트랜스포머(Transformers)는 이차적인 계산 스케일링 문제와 취약한 귀납적 편향(inductive biases) 문제를 겪으며, 비선형 RNN은 기울기 소실 및 폭주 문제로 어려움을 겪습니다. 최근의 상태 공간 모델(SSM)과 선형 재귀 유닛(LRU)은 최첨단 성능을 달성했습니다. 그러나 이러한 모델들은 다음과 같은 중대한 단점을 보입니다:

1. 복잡한 초기화: S4와 같은 SSM은 최적의 다항식 투영에 기반한 HiPPO 초기화와 같이, 훈련 과정 중에 완전히 이해되지 않는 복잡한 이론적 유도 과정을 필요로 합니다.
2. 구현과 이론의 불일치: LRU는 개념적으로는 더 단순하지만, 실제 구현에서는 그 이론적 동기를 온전히 반영하지 못합니다. 구체적으로, 실수 행렬의 고윳값(eigenvalues)에 대한 공액 쌍(conjugate pairs)을 강제하지 못하며, 실수 출력값을 보장하기 위한 제약 조건이 부족하고, 훈련 전반에 걸쳐 일관된 은닉 상태 크기를 보장하지 못하는 학습된 정규화 계수에 의존합니다.
3. 안정성: 기존 방식들은 서로 다른 깊이와 무작위 시드(random seeds)에 대해 견고하지 않을 수 있는 복잡한 정규화 기법이나 초기화 절차를 요구하는 경우가 많습니다.

방법론
저자들은 재귀 상태 행렬 $A$를 회전 행렬으로 파라미터화하는 선형 재귀 모델인 RotRNN(Rotational Recurrent Neural Network)을 제안합니다. 이 방법론은 세 가지 핵심 축을 기반으로 합니다:

1. 회전 행렬을 통한 파라미터화:
   제약이 있는 행렬을 직접 학습하는 대신, RotRNN은 일반적인 가중치 행렬 $M$을 학습하고 이를 왜곡 대칭 행렬(skew-symmetric matrix)의 행렬 지수 함수($A = \exp(M - M^\top)$)를 사용하여 특수 직교군 $SO(N)$으로 매핑합니다. 이를 통해$A$가 훈련 과정 내내 유효한 회전 행렬로 유지되도록 보장합니다.
   선형 재귀의 합성곱 형태(convolutional form)에 필요한 행렬 거듭제곱을 효율적으로 계산하기 위해, 저자들은 회전 행렬 $A$를 블록 대각(block-diagonal) 형태로 분해합니다: $A = P \Theta P^\top$. 여기서 $P$는 직교 행렬이며, $\Theta$는 각도 $\theta_i$로 정의된 $2\times2$ 회전 블록들로 구성됩니다. 이 구조를 통해 값비싼 밀집 행렬 연산 없이도 $A^k$를 빠르게 계산할 수 있습니다.

2. 원칙에 기반한 정규화:
   회전 행렬의 직교성(고윳값이 단위 원 위에 존재함)을 활용하여, 저자들은 훈련 전반에 걸쳐 기대 은닉 상태 노름(expected hidden state norm, $E[\|x_t\|^2] = c$)을 일정하게 유지하는 정규화 절차를 도출합니다.
   재귀식은 $x_t = \gamma A x_{t-1} + \xi B u_t$로 정의됩니다. 여기서 $\gamma$는 학습 가능한 스칼라 감쇠 인자(decay factor)이며, $\xi$는 일정한 노름 조건을 만족하도록 유도된 재조정 계수(rescaling coefficient)입니다. 무한 시퀀스 극한에 기반하여 정규화를 수행하는 LRU와 달리, RotRNN의 정규화는 모든 타임스텝에서 유효하므로, 정규화 상수가 $t$승으로 계산될 때 발생하는 초기 입력의 지수적 감소 현상을 방지합니다.

3. 멀티 헤드 아키텍처:
   단일 스칼라 감쇠 인자 $\gamma$(다양한 메모리 지평을 요구하는 작업에 잘 일반화되지 않을 수 있음)의 한계를 해결하기 위해, RotRNN은 $H$개의 독립적인 저차원 헤드를 채택합니다. 각 헤드는 고유의 회전 행렬 $A^{(h)}$, 입력 행렬 $B^{(h)}$, 그리고 감쇠 인자 $\gamma^{(h)}$를 가집니다. 출력은 결합(concatenate)된 후 선형적으로 혼합되어, 모델이 다양한 스케일에서 의존성을 포착할 수 있도록 합니다.

주요 기여

• 새로운 선형 RNN 레이어: 상태 행렬을 회전으로 파라미터화하여, 알고리즘이 훈련 과정 내내 이론적 동기에 충실하도록 보장하는 재귀 레이어를 제안했습니다.
• 효율적인 계산: 파라미터화된 회전을 블록 대각 형태로 분해함으로써 효율적인 행렬 거듭제곱을 계산하는 방법을 제시하여 빠른 선형 재귀를 가능하게 했습니다.
• 강건한 정규화: 점근적 수렴에 의존하는 이전 연구들과 달리, 일정한 기대 은iel 상태 크기를 유지하는 원칙적인 정규화 절차를 도출하여 더 강력한 안정성 보장을 제공합니다.
• 이론적 통찰: LRU의 특수한 경우(고윳값이 공액 쌍을 형성하는 경우)와 멀티 헤드 RotRNN 사이의 수학적 동등성을 도출하여, 선형 재귀 모델의 내부 메커니즘을 밝히는 데 기여했습니다.

결과
저자들은 여러 긴 시퀀스 모델링 벤치마크에서 RotRNN을 평가했습니다:

• Long Range Arena (LRA): RotRNN은 ListOps, Text (IMDB), Retrieval과 같은 이산 입력 작업에서 SOTA 급 선형 재귀 모델(S4, S5, LRU)과 경쟁력 있는 성능을 보였으며, 베이스라인 중 가장 높은 점수를 기록했습니다. 다만, 픽셀 수준의 이미지 작업(Path-X, CIFAR)에서는 일부 베이스라인보다 약간 낮은 성능을 보였습니다.
• Speech Commands: 35개 클래스 Speech Commands 분류 작업에서, RotRNN은 유사한 파라미터 수를 가진 재현된 LRU와 동일한 성능(정확도 95.2%)을 보였으며, 더 복잡한 딥 상태 공간 모델들과 비교해도 경쟁력 있는 수준을 유지했습니다.
• 안정성 분석: ListOps 데이터셋에 대한 은닉 상태 노름의 경험적 분석 결과, RotRNN은 훈련, 레이어 깊이, 무작위 시드 전반에 걸쳐 거의 일정한 크기를 유지함을 보여주었습니다. 반면, LRU는 은닉 상태 노름이 크게 변동하며 안정적인 크기로 수렴하는 데 훨씬 더 오랜 시간이 걸렸습니다.

의의 및 주장
본 논문은 RotRNN이 기존의 선형 재귀 모델들에 대해 개념적으로 단순하고 효율적이며 수학적으로 원칙적인 대안을 제공한다고 주장합니다. 이 연구의 주요 의의는 이론적 동기와 실제 구현 사이의 간극을 메우는 데 있습니다. 회전 행렬을 사용함으로써 RotRNN은 SSM의 복잡한 초기화 방식과 LRU의 구현 불일치 문제를 피할 수 있습니다. 저자들은 자신들의 모델이 "블랙박스" 식의 초기화에 의존하지 않고도 안정성을 보장하는 강건한 정규화 절차를 제공하며, 표준 벤치마크에서 경쟁력 있는 성능을 달성한다고 단언합니다. 나아가, 본 연구는 회전 기반 블록 구조의 관점에서 LRU를 바라보는 것이 이전 모델들의 선형 재귀 메커니즘을 명확히 하는 데 도움이 될 수 있음을 시사합니다.