기술 요약: 카츠마르츠 선형 어텐션 (KLA)

문제 제기

장기 컨텍스트 언어 모델링은 트랜스포머의 표준 Softmax 어텐션이 갖는 2 차 계산 및 메모리 비용에 의해 저해받고 있습니다. 선형 시간 순환 모델은 컨텍스트를 고정된 크기의 상태로 압축함으로써 이를 해결하지만, 상태 업데이트 규칙의 설계가 결정적인 병목 현상이 됩니다. 가이트드 델타넷 (GDN) 과 같은 기존 접근 방식은 목표 값과 현재 상태 예측 간의 잔차를 '델타 규칙'을 통해 기록하여 기존 연관성을 단순히 누적하는 것이 아니라 편집합니다.

그러나 GDN 은 망각과 업데이트를 균형 있게 조절하기 위해 학습된 스칼라 계수에 의존합니다. 저자들은 이 계수가 이론적으로 유도된 단계 크기가 아닌 경험적 설계 선택이라고 주장합니다. 부적절한 계수는 비최적 업데이트 크기를 초래할 수 있습니다: 큰 노름 (norm) 의 키에 대해서는 과도하게 보정하고, 작은 노름의 키에 대해서는 보정이 부족하여 상태 손상과 비최적 검색 성능을 초래합니다.

방법론: 카츠마르츠 선형 어텐션

본 논문은 학습된 스칼라를 카츠마르츠 투영 방법에 기반한 이론적으로 유도된 단계 크기로 대체하는 GDN 의 변형인 **카츠마르츠 선형 어텐션 (KLA)**을 제안합니다.

이론적 유도

저자들은 키 - 값 쓰기 작업을 제약 최적화 문제로 설정합니다. 토큰 $t$에 대해, 감쇠된 이전 상태 $\tilde{S}_{t-1}$에 가장 가까운 상태 $S_t$를 찾아 선형 제약 조건 $S_t^T k_t = v_t$ (여기서 $k_t$는 키, $v_t$는 값) 를 만족시키는 것입니다.

• 최적화 목적 함수: $S^T k_t = v_t$를 제약 조건으로 하여 $\min_S \frac{1}{2}\|S - \tilde{S}_{t-1}\|_F^2$
• 해: 최적 해는 제약 조건 초평면 위로의 $\tilde{S}_{t-1}$의 직교 투영입니다. 이는 $\beta_t = 1/\|k_t\|_2^2$의 단계 크기 (계수) 를 산출합니다.

KLA 업데이트 규칙

KLA 는 이 정확한 투영을 학습 가능한 완화 게이트 $\eta_t$와 수치적 안정화자 $\epsilon$을 포함하는 실용적인 업데이트 규칙으로 일반화합니다:
$$\beta_t = \frac{\eta_t}{\|k_t\|_2^2 + \epsilon}$$
상태 업데이트는 다음과 같습니다:
$$S_t = \tilde{S}_{t-1} + \beta_t k_t e_t^T$$
여기서 $e_t = v_t - \tilde{S}_{t-1}^T k_t$는 잔차 오차입니다.

주요 특성

이 논문은 이 계수가 완화되지 않은 경우 ($\eta_t=1, \epsilon=0$) 에 세 가지 바람직한 특성을 동시에 만족함을 입증합니다:

1. 직교 투영: 제약 조건을 만족시키기 위해 상태에 가하는 최소 노름 변화를 수행합니다.
2. 정확한 선 검색: 업데이트 방향을 따라 토큰당 손실의 정확한 최소화자입니다.
3. 정규화된 경사 하강 단계: 정규화된 확률적 경사 하강 (SGD) 단계로 작용합니다.

구현

KLA 는 GDN 을 위한 **즉시 교체 (drop-in replacement)**로 설계되었습니다. 다음을 유지합니다:

• 상태 형태 ($S \in \mathbb{R}^{d_k \times d_v}$).
• 게이트 메커니즘 (감쇠 및 업데이트 게이트).
• 선형 순환 구조.
• 조각별 병렬 학습 알고리즘.

유일한 구조적 변화는 조각별 솔버 내에서 대각 계수 행렬 $B = \text{Diag}(\beta_1, \dots, \beta_C)$의 계산입니다. 새로운 하드웨어 커널이나 상태 레이아웃이 필요하지 않습니다.

주요 기여

1. 자기 정규화 델타 규칙: 카츠마르츠 투영에서 유도된 KLA 계수로, 토큰당 손실에 대해 직교 투영, 정확한 선 검색 최소화자, 그리고 정규화된 경사 하강 단계임을 증명합니다.
2. 즉시 교체 순환 알고리즘: KLA 가 GDN 순환 및 조각별 솔버를 유지하며 대각 계수 행렬의 변경만 필요로 함을 입증합니다.
3. 경험적 이득: 아키텍처 확장 없이 GDN 및 기타 선형 시간 베이스라인보다 낮은 퍼플렉시티와 더 나은 상태 추적을 달성함을 검증합니다.
4. 유지된 효율성: KLA 가 $O(n)$ 시간/메모리 복잡성을 유지하고 GDN 의 프리필 지연 시간과 일치하면서 디코드 처리량을 개선함을 확인합니다.

실험 결과

실험은 SlimPajama-672B 데이터셋에서 0.4B 파라미터 규모로 1B 토큰 예산을 사용하여 수행되었습니다.

언어 모델링

• 검증 퍼플렉시티 (PPL): KLA 는 8.09 PPL을 달성하여 GDN(8.50), DeltaNet(8.26), GLA(8.42), Longhorn(8.79), Mamba2(9.41) 을 능가했습니다.
• 장기 컨텍스트 외삽: 장기 컨텍스트 미세 조정 없이도 KLA 는 1K 에서 65K 토큰에 이르는 컨텍스트 길이 전반에 걸쳐 가장 낮은 퍼플렉시티를 유지했으며, GDN 대비 일관되게 3.7~4.1 PPL 의 우위를 보였습니다.

합성 작업

• 싱글 니들 인 어 헤이스택 (S-NIAH): KLA 는 테스트된 모든 컨텍스트 길이 (1K~8K) 에서 100% 정확도에 도달한 반면, GDN 은 약간 하락했고 Mamba2 는 크게 실패했습니다.
• 멀티 쿼리 연관 회상 (MQAR): KLA 는 GDN 대비 7.03 포인트만큼 8 배 길이 외삽 정확도를 향상시켰습니다 (73.84% 대 66.81%).
• 스택 및 회문: KLA 는 상태 추적 및 반전 작업에서 경쟁력 있거나 우수한 성능을 보였습니다.

효율성

• 프리필 지연 시간: KLA 는 GDN 의 프리필 지연 시간 (예: 131K 컨텍스트에서 약 1220ms) 과 일치하여, 카츠마르츠 계수 계산이 지배적인 조각별 커널 비용에 영향을 미치지 않음을 확인했습니다.
• 디코드 처리량: 32K 컨텍스트에서 KLA 는 초당 3.54 토큰을 달성하여 GDN(초당 1.65 토큰) 보다 2.1 배 높고 Mamba2 보다 1.7 배 높았습니다.

제거 실험

저자들은 정규화 전략, 게이트 메커니즘, 상태 확장을 제거했습니다. 결과는 키 노름 정규화가 안정성에 결정적임을 나타냈습니다. 이를 제거하거나 학습된 스칼라만 사용할 경우 프리트레이닝 안정성 (PPL) 과 합성 작업 성능이 크게 저하되었습니다.

중요성 및 주장

본 논문은 키 노름 정규화된 카츠마르츠 계수가 델타 규칙 시퀀스 모델의 '1 차 설계 축'이라고 주장합니다. 쓰기 작업을 제약 투영으로 취급함으로써 KLA 는 경험적으로 학습된 스칼라에 내재된 비최적 업데이트 크기 문제를 해결합니다.

이러한 개선이 순환 상태, 하드웨어 커널, 또는 아키텍처 복잡성을 변경하지 않고 달성된다는 점이 중요합니다. 저자들은 상태 용량이나 게이트 구조가 아니라 계수 자체가 델타 규칙 쓰기가 국소 회귀 문제를 성공적으로 해결하는지 여부를 결정하는 주요 요인이라고 주장합니다. 이는 향후 선형 시간 모델이 이론적으로 근거된 단계 크기만 채택함으로써 정확도, 외삽, 디코드 효율성에서 상당한 향상을 달성할 수 있음을 시사합니다.