기술 요약: MuonSSM – 시퀀스 모델링을 위한 직교화된 상태 공간 모델 (Orthogonalizing State Space Models for Sequence Modeling)

문제 정의

상태 공간 모델(SSM)은 긴 시퀀스 모델링을 위해 어텐션 메커니즘의 효율적인 선형 시간 대안으로 부상했습니다. 그러나 기존의 SSM(예: Mamba, DeltaNet)은 확장된 호라이즌(horizon)에서 종종 불안정성과 메모리 퇴화 문제를 겪습니다. 본 논문은 현재의 SSM이 메모리 상태에 대해 **1차, 랭크-1 업데이트(first-order, rank-one updates)**에 의존한다는 점을 지적합니다. 이러한 업데이트는 세 가지 주요 문제를 야기합니다:

1. 스펙트럼 비등방성(Spectral Anisotropy): 메모리 행렬의 특이값(singular values)이 매우 불균일해져, 조건수가 나쁜(ill-conditioned) 최적화 지형을 형성합니다.
2. 그래디언트 퇴화(Gradient Degradation): 입력 의존적 전이 행렬을 통한 반복적인 수축(contraction)으로 인해, 긴 시퀀스에서 그래디언트 소실 현상이 발생합니다.
3. 메모리 간섭(Memory Interference): 시간적 관성(temporal inertia)의 부족과 불균형한 업데이트 기하학으로 인해, 새로운 업데이트가 이전 정보를 덮어쓸 수 있습니다.

최근의 변형 모델들이 게이팅이나 정규화를 통해 이를 완화하려고 시도하지만, 이들은 축적된 업데이트의 기하학적 구조를 명시적으로 제어하거나 장거리 신호 전파를 개선하기 위한 시간적 관성을 도입하지는 못합니다.

방법론: MuonSSM

저자들은 전이 행렬(recurrent transition matrix) 자체를 수정하는 대신, 메모리 업데이트의 기하학적 구조를 명시적으로 제어함으로써 SSM 학습을 안정화하는 프레임워크인 MuonSSM을 제안합니다. 이 접근 방식은 효율적인 병렬 스캔에 필요한 아핀 구조(affine structure)를 유지하면서, 두 가지 핵심 구성 요소를 통해 표준 SSM 역학을 보강합니다.

1. 모멘텀 증강 역학 (Momentum-Augmented Dynamics)

MuonSSM은 타임스텝에 걸쳐 업데이트 방향을 축적하는 보조 모멘텀 행렬 $M_t$를 도입합니다. 이는 모멘텀 기반 최적화를 모사하여, 상태 업데이트에 "시간적 관성"을 제공합니다.

• 모멘텀 상태는 다음과 같이 업데이트됩니다: $M_t = \gamma M_{t-1} + \text{NS}(\tau \beta_t v_t k_t^\top)$.
• 메모리 상태 $S_t$는 이전 상태와 모멘텀 항을 결합하여 업데이트됩니다: $S_t = S_{t-1}(\alpha_t(I - \beta_t \eta k_t k_t^\top)) + M_t$.
• 이는 그래디언트 흐름을 위한 추가적인 채널을 생성하여, 장거리 크레딧 할당(credit assignment)에서의 지수적 감쇠를 완화합니다.

2. 단일 반복 뉴턴-슐츠(Newton–Schulz, NS) 정규화

스펙트럼 불균형을 해결하기 위해, 입력 주입(rank-1 행렬)은 모멘텀 상태에 더해지기 전 가벼운 단일 단계 뉴턴-슐츠 반복을 통해 정규화됩니다.

• 연산자: $\text{NS}(X) = (a + b \tilde{X}\tilde{X}^\top + c (\tilde{X}\tilde{X}^\top)^2)\tilde{X}$, 여기서 $\tilde{X}$는 프로베니우스 노름(Frobenius norm)에 의해 스케일링된 입력입니다.
• 기능: 이 연산자는 특이값을 제한(단위 근처로 유지)하며, 명시적인 SVD 없이도 입력의 랭크를 보존합니다.
• 기하학적 효과: 표준 프로베니우스 정규화와 달리, NS 연산자의 역전파 자코비안(backward-pass Jacobian)은 현재의 랭크-1 쓰기(write)에 직교하는 방향을 증폭합니다. 이는 "랭크 풍부화(rank enrichment)"를 촉진하여, 메모리 상태가 저차원 부분 공간으로 붕괴되는 것을 방지합니다.

병렬성 (Parallelizability)

결정적으로, 메모리와 모멘텀의 결합된 역학은 **블록-아핀 재귀(block-affine recurrence)**로 구성될 수 있습니다. 이를 통해 모델은 병렬 연관 스캔(parallel associative scans)을 활용할 수 있어, 재귀적 깊이를 $O(L)$에서 $O(\log L)$로 줄이고 계산 복잡도를 시퀀스 길이 $L$에 대해 선형적으로 유지할 수 있습니다.

주요 기여

1. 프레임워크 설계: 모멘텀 기반 2차 역학(second-order dynamics)과 암시적 스펙트럼 컨디셔닝을 스캔 기반 재귀에 통합한 MuonSSM을 도입했습니다.
2. 이론적 분석:
   • 증강된 역학이 병렬성을 보존함을 증명했습니다.
   • 모멘텀 경로가 소멸하지 않는 그래디언트 성분을 도입하여 장거리 전파를 안정화함을 입증했습니다.
   • NS 연산자가 특이값을 제한하고, 역방향 신호를 직교하는 방향으로 편향시켜 랭크 풍부화를 유도함을 분석했습니다.
3. 실험적 검증: 언어, 비전, 시계열 태스크에 걸친 포괄적인 실험을 통해 정확도, 견고성 및 길이 일반화 측면에서 일관된 개선을 보여주었습니다.

실험 결과

저자들은 세 가지 모달리티에 대해 Mamba, LongHorn, Gated DeltaNet 등의 베이스라인과 MuonSSM을 비교 평가했습니다.

• 언어 모델링: FineWeb-Edu 데이터셋에서 MuonSSM 변형 모델들은 더 낮은 퍼플렉시티(perplexity)와 우수한 상식 추론 벤치마크 성능을 달-성했습니다. 특히, "Needle-in-a-Haystack" 검색 태스크에서 MuonSSM은 최대 8K 컨텍스트 길이까지 높은 정확도를 유지하며, 훈련 호라이즌(2K 토큰)을 넘어서면 급격히 성능이 저하되는 베이스라인들을 크게 앞질렀습니다.
• 비전: MambaVision 프레임워크에서 Mamba 믹서를 MuonSSM으로 교체했을 때 ImageNet-1K 분류, 객체 탐지(COCO), 의미론적 분할(ADE20K)에서 일관된 성능 향상을 보였습니다. 결정적으로, MuonSSM은 손상(IN-C), 렌더링(IN-R) 및 적대적 공격(IN-A)에 대해 현저히 개선된 견고성을 보여주었으며, 이는 질감 편향(texture bias)에서 형태 불변적 특징 학습(shape-invariant feature learning)으로의 전환을 시사합니다.
• 시계열: 인간 활동 인식 태스크(Human Activity Recognition: MuWiGes, UESTC-MMEA-CL, MMAct)에서 MuonSSM 변형 모델들은 베이스라인을 능가했으며, 데이터셋의 복잡성이 증가할수록 그 격차가 커졌습니다.

절제 연구 및 분석 (Ablation & Analysis):

• 스펙트럼 컨디셔닝: 시각화를 통해 표준 SSM은 스펙트럼 붕괴(높은 조건수)를 겪는 반면, MuonSSM은 거의 등거리(near-isometric) 스펙트럼을 유지함을 확인했습니다.
• 용량 vs. 기하학: 베이스라인의 상태 차원을 두 배로 늘리는 것은 MuonSSM의 상당한 개선에 비해 미미한 이득만을 주었으며, 이는 이점이 단순한 용량 확장이 아닌 기하학적 컨디셔닝에서 기인함을 확인시켜 줍니다.
• 효율성: 단계당 상수 배수의 오버헤드가 발생하지만, MuonSSM은 수렴 속도를 약 1.3배 가속화하여 목표 성능에 도달하기 위한 총 훈련 단계를 줄여줍니다.

의의 및 주장

본 논문은 업데이트의 기하학적 컨디셔닝이 안정적이고 확장 가능한 시퀀스 모델링을 위한 원칙적인 경로임을 주장합니다. 모멘텀과 뉴턴-슐츠 정규화를 통해 메모리 업데이트의 기하학을 명시적으로 제어함으로써, MuonSSM은 병렬 스캔의 선형 시간 효율성을 희생하지 않고도 1차 SSM의 근본적인 불안정성을 해결합니다.

저자들은 이 접근 방식이 다음과 같은 효과를 가진다고 주장합니다:

• 그래디언트 전파를 개선하고 스펙트럼 증폭을 완화합니다.
• 랭크 붕괴를 방지하여 장기 호라이즌에서 메모리 표현을 풍부하게 합니다.
• 긴 컨텍스트 태스크에 대한 강력한 대안을 제공하며, 복잡한 전이 행렬 파라미터화만이 아니라 업데이트 역학을 통해서도 안정성을 달성할 수 있음을 입증합니다.

이 연구는 향후 시퀀스 모델링이 재귀 구조로부터 업데이트 메커니즘의 안정성을 분리함으로써, 다양한 모달리티에 걸쳐 안정적인 훈련을 위한 단순하면서도 효과적인 메커니즘을 얻을 수 있음을 시사합니다.