기술 요약: 변분 선형 어텐션 (VLA)

1. 문제 제기

장문맥 시퀀스 모델링은 계산 효율성과 메모리 안정성 간의 근본적인 트레이드오프에 직면해 있습니다. 표준 트랜스포머 어텐션은 최적의 연관성 검색을 달성하지만, $O(T^2)$ 의 시간 복잡도와 $O(Td)$의 메모리 사용량으로 인해$10^5$토큰을 초과하는 시퀀스에는 비현실적입니다. 선형 어텐션은 커널 특징 맵과 재귀적 상태 업데이트를 통해 복잡도를$O(T)$로 낮춤으로써 계산 병목 현상을 해결합니다. 그러나 이는 무조건적인 가산적 누적으로 인해 내부 메모리 상태$S_t$ 가 프로베니우스 노름 (Frobenius norm) 에서 무한히 증가 ($O(T)$) 하는 치명적인 실패 모드를 도입합니다. 이 무한한 증가는 저장된 연관성 간의 점진적인 간섭을 초래하여, 문맥 길이가 증가함에 따라 검색 정확도가 저하됩니다.

DeltaNet 은 상태를 균일하게 감쇠시키는 스칼라 게이트 메커니즘을 통해 이를 완화하려 시도하지만, 방향 선택적 제어 기능이 부족합니다. 스칼라 게이트는 최근에 덮어씌워진 방향을 우선적으로 잊으면서 더 오래되고 안정적인 연관성을 보존할 수 없어, 높은 메모리 부하 하에서 급격한 정확도 붕괴를 초래합니다. 저자들은 핵심 문제가 계산 비용이 아니라 메모리 안정성, 특히 메모리 업데이트의 기하학에 있다고 주장합니다.

2. 방법론: 변분 선형 어텐션

본 논문은 선형 어텐션 메모리 업데이트를 온라인 정규화된 최소제곱 문제로 재해석한 변분 선형 어텐션 (VLA) 을 제안합니다.

2.1 핵심 공식화

무조건적인 누적 ($S_t = S_{t-1} + v_t k_t^\top$) 대신, VLA 는 페널티가 부과된 예측 오차를 최소화하는 메모리 행렬 $S_t$ 를 찾습니다:
$$S^*_t = \arg \min_S \sum_{s=1}^t \|v_s - S\hat{k}_s\|^2 + \text{tr}(S M_t S^\top)$$
여기서 $\hat{k}_s$ 는 정규화된 키 (key) 이고, $M_t$ 는 이전에 본 키들의 기하학을 인코딩하는 시간 가변 페널티 행렬입니다.

2.2 적응형 페널티 기하학

VLA 는 각 단계당 $O(d^2)$ 비용으로 Sherman-Morrison 랭크-1 공식을 사용하여 역 페널티 행렬 $A_t = M_t^{-1}$ 을 정확하게 유지합니다.

• 메커니즘: $M_t$ 는 학습된 페널티 방향의 외적 (outer products) 을 누적합니다. 결과적으로 $A_t$ 는 최근 "페널티 질량" (즉, 쓰여진 방향) 을 받은 방향으로 수축합니다.
• 효과: 최근 쓰여진 부분 공간은 더 작은 업데이트를 받고, untouched 된 부분 공간은 사용 가능하게 유지됩니다. 이는 DeltaNet 의 스칼라 게이트에 대한 행렬 값 아날로그로 작용하여, 방향 선택적 메모리 업데이트를 가능하게 합니다.

2.3 재귀적 업데이트 규칙

이 공식에서 유도된 정확한 재귀적 업데이트 규칙은 다음과 같습니다:

1. 역 페널티 업데이트: $A_t = A_{t-1} - \frac{z_t z_t^\top}{1 + u_t^\top z_t}$ (Sherman-Morrison).
2. 게이팅 벡터 계산: $\hat{\alpha}_t = \frac{A_t \hat{k}_t}{\|A_t \hat{k}_t\|}$.
3. 상태 업데이트: $S_t = S_{t-1} + (v_t - S_{t-1}\hat{k}_t) \hat{\alpha}_t^\top$.

중요하게도, 키 $\hat{k}_t$ 와 게이팅 벡터 $\hat{\alpha}_t$ 는 모두 단위 길이로 정규화됩니다. 출력은 표준 선형 어텐션과 유사한 재귀식의 비율로 계산되며, ELU+1 특징 맵을 사용합니다.

2.4 구현 효율성

Sherman-Morrison 업데이트는 상수 인자 오버헤드 (표준 선형 어텐션보다 약 $5\times$ 많은 연산) 를 도입하지만, 저자들은 병렬 Blelloch-scan 공식과 퓨즈드 Triton 커널을 구현했습니다. 이는 모든 $T$ 단계를 단일 GPU 커널 런치로 융합하여 토큰별 디스패치 오버헤드를 제거하고, 순차적 Python 구현 대비 $14\times$ 의 속도 향상을 달성합니다.

3. 주요 기여

3.1 이론적 보장

본 논문은 두 가지 안정성 속성에 대한 엄밀한 증명을 제공합니다:

1. 유계 상태 성장 (명제 1): 유계 입력 하에서 프로베니우스 노름 $\|S_t\|_F$ 는 자기 제한적입니다. 표준 선형 어텐션에서 $\|S_T\|_F = O(T)$ 인 것과 달리, VLA 의 상태는 예측 잔차가 사라짐에 따라 유한한 평탄면으로 수렴합니다.
2. 단위 야코비안 스펙트럼 노름 (명제 2): 재귀 야코비안 $J_t = \partial S_t / \partial S_{t-1} = I - \hat{\alpha}_t \hat{k}_t^\top$ 은 벡터가 단위 정규화되어 있다면 모든 시퀀스 길이와 헤드 차원에 대해 스펙트럼 노름이 정확히 1입니다. 이는 임의의 깊이에서 안정적인 그래디언트 흐름 (폭발 또는 소멸 없음) 을 보장합니다. 정규화되지 않은 경우, 스펙트럼 노름은 $d_h/\lambda_0$ 로 증가하여 그래디언트 증폭을 초래합니다 (예: $d_h=96$ 일 때 25 단계 후 $10^{32}$).

3.2 연관성 메모리 용량

VLA 는 헤드당 $d_h$ 개의 연관성을 유지하는 이론적 용량을 유지합니다. 그러나 가산적 누적과 달리, VLA 는 새로운 쓰기를 최근 사용된 부분 공간에 직교하는 방향으로 라우팅합니다. 이는 균일한 희석 없이 오래된 연관성을 보존할 수 있게 하여, 용량 한계까지 정확한 회상을 유지합니다.

3.3 효율성

제안된 Triton 퓨즈드 커널은 $O(T)$ 스케일링을 달성하며, 약 43,000 토큰에서 소프트맥스 어텐션 지연 시간을 하회하여 장문맥 추론에 실용적입니다.

4. 실험 결과

실험은 Multi-Query Associative Recall (MQAR) 작업과 복사 (Copy) 작업을 사용하여 공유된 2 층 트랜스포머 백본에서 수행되었습니다.

• 상태 노름 안정성: $T=1,000$ 에서 표준 선형 어텐션은 $\|S_T\|_F \approx 1,600$ 에 도달하는 반면, VLA 는 15 미만에 머무릅니다 ($109\times$ 감소). DeltaNet 또한 유계이지만 VLA 의 방향 선택성이 부족합니다.
• 연관성 회상 (용량): 쌍의 수 $n_{pairs} < d_h$ 인 영역 (구체적으로 $n=24, d_h=32$):
  • VLA: 1.000의 완전 일치 정확도를 유지합니다.
  • DeltaNet: 0.010으로 붕괴됩니다.
  • 표준 선형 어텐션: 0.074로 하락합니다.
  • 소프트맥스: 0.26으로 하락합니다.
• 장문맥 유지: 고정된 $n=8$ 과 $T$ 를 512 까지 증가시킬 때, VLA 는 1.000 정확도를 유지합니다. DeltaNet 은 스칼라 게이트가 쿼리가 도착하기 전에 오래된 연관성을 감쇠시키기 때문에 모든 길이에서 $\approx 0.01$ 로 붕괴됩니다.
• 과부하 영역: $n > d_h$ 일 때, VLA 는 적응형 페널티 행렬이 최근 쓰인 방향을 우선적으로 덮어써서 오래된 연관성을 부분적으로 보존하기 때문에, 가산적 선형 어텐션보다 더 우아하게 저하됩니다.

5. 중요성과 주장

본 논문은 장문맥 신뢰성은 계산적 문제일 뿐만 아니라 메모리 기하학 문제라고 주장합니다.

• 기하학적 제어: 주요 기여는 업데이트의 계산 비용을 단순히 줄이는 것보다 상태가 업데이트를 허용받는 방향을 제어하는 것이 더 중요함을 입증한 것입니다. 적응형 페널티 역행렬에 의해 지배되는 잔차 오차 업데이트로 무조건적 누적을 대체함으로써, VLA 는 수치적으로 안정적이고 상수 메모리 시퀀스 모델링을 달성합니다.
• 고전 이론과의 연결: VLA 는 선형 어텐션을 재귀적 최소제곱 (RLS) 적응 필터와 연결하여, 이전의 "빠른 가중치 (fast-weight)" 접근법에서 결여되었던 원칙적인 유도 및 안정성 이론을 제공합니다.
• 실용적 범위: 저자들은 VLA 가 $d_h$ 개의 가장 최근에 강화된 연관성에 대한 고정 크기, 수치적으로 안정된 요약을 제공하며, KV 캐시 인프라 없이 스트리밍 추론에 적합하다고 주장합니다. VLA 는 헤드당 $d_h$ 개의 연관성이라는 근본적인 용량 한계를 극복하지는 못하지만, 그 한계 내에서 우수한 안정성과 검색을 제공한다고 명시합니다.

이 연구는 VLA 가 상수 인자 오버헤드를 incur 하지만, 이론적 안정성 (단위 야코비안, 유계 상태) 과 실험적 성능 (용량까지 완벽한 회상) 의 결합이 표준 선형 어텐션이 간섭으로 인해 실패하는 장문맥 트랜스포머에 대한 견고한 해결책을 확립한다고 결론지었습니다.