기술 요약: 칼만 선형 어텐션 (Kalman Linear Attention, KLA)

문제 정의

Mamba 및 게이티드 선형 어텐션(GLA)과 같은 현재의 상태 공간 모델(SSM)은 트랜스포머의 이차 복잡도 병목 현상을 해결하기 위해 선형 복잡도와 병렬 처리가 가능한 대안을 제공한다. 그러나 이러한 모델들은 선형 또는 아핀(affine) 은닉 상태 업데이트에 의존하며, 이는 소프트맥스 정규화가 유도하는 비선형 상호작관에 비해 표현력이 제한적이다. 또한, 기존의 선형 RNN 및 SSM은 상태 불확실성을 나타내는 명시적인 메커니즘이 부족하며, 은닉 상태를 결정론적인 것으로 취급한다. 본 논문은 선형 업데이트의 효율성과 비선형 필터링의 표현력/강건성 사이의 간극을 다루며, 정확한 칼만 필터 업데이트를 효율적인 시퀀스 믹서로 구현할 수 있는지에 대한 질문을 던진다.

방법론: 칼만 선형 어텐션 (KLA)

저자들은 시퀀스 모델링을 정확한 베이지안 필터링으로 공식화한 확률적 시퀀스 믹서인 **칼만 선형 어텐션(KLA)**을 제안한다. 핵심 혁신은 칼만 필터를 정보 형태(information form)(정밀도 및 자연 파라미터 사용)로 재매개변수화하여 병렬 계산을 가능하게 하는 데 있다.

1. 확률적 공식화

KLA는 시퀀스 믹싱을 선형-가우시안 상태 공간 모델에서의 추론으로 모델링한다:

• 역학 사전 확률(Dynamics Prior): 잠재 상태의 진화는 정확하게 이산화된 오른슈타인-우울렌벡(Ornstein-Uhlenbeck, OU) 프로세스로 모델링된다. 이는 표준 SSM의 결정론적 역학과 구별되는, 상태 전이의 불확실성을 포착하는 프로세스 노이즈 항($\bar{p}_t$)을 도입한다.
• 관측 모델: 입력 토큰은 잠재 상태의 노이즈가 섞인 관측치로 취급되며, 관측 연산자($k_t$)와 토큰 증거에 대한 신뢰도를 나타내는 값 정밀도($\Lambda^v_t$)에 의해 매개변수화된다.
• 리드아웃(Readout): 출력은 쿼리 연산자($q_t$)를 통해 투영된 잠재 상태의 사후 평균으로부터 유도된다.

2. 뫼비우스 변환을 통한 병렬화

고전적인 칼만 필터링은 본질적으로 순차적이다. KLA는 필터의 정보 형태를 활용하여 이를 극복한다:

• 정밀도 재귀(Precision Recursion): 사후 정밀도($\lambda_t$)의 업데이트는 이전 정밀도($\lambda_{t-1}$)의 뫼비우스(Möbius, 분수-선형) 변환임이 입증되었다.
  $$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$$
  결정적으로, 뫼비우스 변환은 결합 법칙을 따른다. 이를 통해 정밀도 업데이트를 $O(T)$ 작업량과 $O(\log T)$ 깊이로 **병렬 프리픽스 스캔(parallel prefix scans, associative scans)**을 통해 계산할 수 있으며, 이는 Mamba와 같은 현대적 SSM의 병렬 효율성과 일치한다.
• 평균 재귀(Mean Recursion): 정밀도 경로가 병렬로 계산되면, 사후 평균(정보 평균) 업데이트는 역시 병렬화 가능한 아핀 변환을 따른다.

3. 불확실성으로부터의 비선형 게이팅

게이팅이 흔히 선형이거나 토큰 의존적이지만 결정론적인 선형 SSM과 달리, KLA의 게이팅은 정밀도의 비율로부터 자연스럽게 발생한다. "망각 게이트(forget gate)"와 "쓰기 게이트(write gate)"는 이력에 의존하며 비선형적이고, 사전 상태와 새로운 관측치 사이의 상대적 불확실성에 의해 구동된다. 이를 통해 모델은 모든 토큰에 대해 명시적인 학습된 게이팅 파라미터를 요구하지 않고도, 신뢰도에 따라 정보를 선택적으로 수용할 수 있다.

4. 아키텍처

KLA는 표준 SSM 또는 어텐션 레이어의 드롭인 대체제로 설계되었다. Mamba와 유사한 융합된 MLP 아키텍처를 따르며, 1D 인과적 컨볼루션과 SiLU 활성화 함수를 포함한다. 저자들은 고대역폭 메모리(HBM)에서 확장된 상태를 실체화하지 않도록 뫼비우스 스캔을 위한 효율적인 CUDA 및 Triton 커널을 구현하였다.

주요 기여

1. 결합적 재매개변수화(Associative Reparameterization): 저자들은 정보 형태의 대각 선형-가우시안 칼만 필터링이 뫼비우스 구조를 가지며, 이를 통해 베이지안 필터링을 위한 정확한 병렬 프리픽스 스캔이 가능함을 입증하였다.
2. 불확실성 기반 비선형성: KLA는 정밀도 비율에서 유도된 비선형 게이팅을 도입하여, 선형/아핀 SSM 업데이트보다 엄격하게 더 높은 표현력을 제공하면서도 선형 시간 스캔 구조를 유지한다.
3. KLA 레이어: 명시적인 믿음 상태(belief-state) 불확실성을 생성하고 심층 네트워크에 쌓을 수 있는 실용적인 드롭인 시퀀스 믹서를 제공한다.
4. 스케일링 및 검증: 저자들은 십억 토큰 규모에서 훈련된 최초의 스택형 베이지안 필터링 프리미티브를 제시하여, 안정성과 성능을 검증하였다.

실험 결과

저자들은 합성 태스크, 상태 추적 벤치마크, 대규모 언어 모델링을 통해 KLA를 평가하였다.

• 합성 태스크 (MAD-Lab): KLA는 불확실성 하에서의 선택이 필요한 태스크, 특히 Fuzzy Recall과 Selective Copy에서 GDN, GLA, Mamba, mLSTM을 능가한다. 특히 Compression 태스크에서 최첨단(SOTA) 결과를 달성하였으며, 이는 정보 형태의 필터링이 최적의 컨텍스트 압축기로 작용함을 시사한다.
• 상태 추적 (A5 Permutation): NC1-완비(NC1-complete) 상태 추적 벤치마크인 A5 순열-합성 태스크에서, KLA는 상수 깊이(1~2 블록)로 태스크를 해결한다. 반면, 선형 SSM 및 어텐션 기반 모델은 시퀀스 길이에 따라 증가하는 깊이를 필요로 한다. 이는 KLA의 분수-선형 업데이트가 선형 재귀 모델의 표현력 한계를 깨뜨리면서도 병렬 훈련이 가능함을 보여준다.
• 장기 컨텍스트 연상 회상 (MQAR): KLA는 시퀀스 길이 2048인 Multi-Query Associative Recall 태스크에서 거의 완벽한 정확도(>95%)를 달술하며, Mamba를 능가하고 이 설정에서 실패하는 GLA를 크게 앞선다.
• 대규모 사전 학습:
  • 단독 모델: 45M 및 180M 파라미터에서, 단독 KLA는 제로샷 상식 벤치마크(예: HellaSwag, PIQA)에서 GPT, Mamba, GDN과 대등하거나 개선된 성능을 보인다.
  • 하이브리드: GPT의 마지막 어텐션 레이어를 KLA로 교체한 모델(GPT+KLA)은 하이브리드 모델 중 가장 강력한 성능을 보였으며, GPT+Mamba 및 GPT+GDN을 능가한다.
• 어블레이션(Ablation): 저자들은 오른슈타인-우울렌크(OU) 사전 확률이 안정적인 딥 스태킹에 결정적임을 보여준다. 평균 회귀적 확률 역학이 없다면, KLA의 딥 스택 훈련은 불안정해진다.

의의 및 주장

본 논문은 KLA가 선형 상태 공간 모델의 효율성과 비선형 어텐션 메커니즘의 표현력 사이의 간극을 성공적으로 메운다고 주장한다. 시퀀스 믹싱을 정확한 베이지안 필터링으로 구성함으로써, KLA는 다음을 제공한다:

• 명시적 불확실성: 결정론적 믹서와 달리, KLA는 정보 게이팅을 직접 제어하는 명시적인 믿음 상태 불확실성을 보유한다.
• 엄격히 높은 표현력: 뫼비우스 업데이트를 통해 KLA는 선형 SSM이 할 수 없는 상태 추적 태스크(A5 등)를 병렬 훈련 능력을 희생하지 않고 해결할 수 있다.
• 확장성: 본 연구는 정확한 베이지안 필터링 프리미티브를 스택할 수 있고 십억 토큰 규모에서 훈련할 수 있음을 입려하여, 확률적 필터링이 본질적으로 순차적이거나 대규모 언어 모델링에 계산적으로 불가능하다는 통념에 도전한다.

저자들은 KLA를 단순한 이론적 호기기구가 아니라, 긴 컨텍스트 처리에 필요한 계산 효율성을 유지하면서 현대적 언어 모델의 강건성과 상태 추적 능력을 향상시키는 실용적인 드롭인 프리미티브로 포지셔닝한다.