긴 문맥을 위한 사전조건화 주의 메커니즘 LUCID

초록

LUCID는 키 간 상관관계를 RKHS에서 사전조건화하여 소프트맥스 주의의 확산 문제를 해결한다. 키‑키 유사도 행렬의 하삼각 형태 역행렬을 이용해 확률 질량을 중요한 토큰에 집중시키면서도 기존 소프트맥스의 온도와 그래디언트 흐름을 유지한다. 1B 파라미터 모델을 128K 토큰까지 학습시킨 결과, BABILong·RULER·SCROLLS·LongBench 등 장문 검색 벤치마크에서 14~18% 성능 향상을 달성했다.

상세 분석

본 논문은 트랜스포머의 핵심인 소프트맥스 기반 점곱 주의가 긴 시퀀스에서 “주의 노이즈”라 불리는 확산 현상에 취약함을 지적한다. 시퀀스가 길어질수록 키 벡터 간 내적이 커져 exp(⟨k_i,k_j⟩) 값이 크게 겹치게 되고, 이는 소프트맥스 분모를 통해 모든 토큰에 일정 수준의 확률 질량을 할당하게 만든다. 결과적으로 실제로 필요한 “needle” 토큰에 대한 집중도가 떨어진다. 온도를 낮춰 샤프한 분포를 만들면 그래디언트가 사라지는 vanishing‑gradient 문제도 발생한다.

LUCID는 이러한 근본 원인을 RKHS 관점에서 재해석한다. exp(⟨q,k⟩)=⟨ϕ(q),ϕ(k)⟩ 로 표현되는 무한 차원 특징 공간에서 키 벡터 ϕ(k) 들을 사전조건화 행렬 P = (M⊙exp(KKᵀ))⁻¹ 로 정규화한다. 여기서 M은 인과 마스크이며, P는 하삼각 행렬이므로 forward substitution 으로 효율적으로 역연산이 가능하다. 사전조건화는 키 간 상관관계를 최소화(조건수 κ≈1)하여, 쿼리가 실제로 관련된 키에만 높은 가중치를 부여하도록 만든다.

이론적으로는 LUCID가 기존 소프트맥스의 온도 파라미터를 그대로 유지하면서도 출력 단계에서 추가적인 “디컨볼루션” 역할을 수행한다는 점을 증명한다. Theorem 1에 따르면, 사전조건화 행렬이 가역적이면 소프트맥스의 Jacobian이 0이 되는 상황(온도→0)과 무관하게 ∂o/∂q 가 0이 되지 않는다. 즉, 샤프한 주의 분포를 얻으면서도 학습 가능한 그래디언트를 보존한다.

학습 측면에서는 기존 선형 목표(−vᵀSϕ(k)) 대신 quadratic 목표 ½‖Sϕ(k)−v‖² 를 사용함으로써 “erase‑then‑write” 메커니즘을 구현한다. 이는 DeltaNet의 무한 차원 확장으로, 현재 상태와 일치하면 업데이트가 멈추는 self‑regulation 특성을 제공한다. 결과적으로 키가 중복될 때 발생하는 간섭을 자연스럽게 제거한다.

효율성 측면에서 LUCID는 기존 소프트맥스와 동일한 O(N²d) 복잡도를 유지한다. 사전조건화 행렬의 대각선은 정규화된 키 벡터의 노름을 1로 맞춤으로써 1이 되며, off‑diagonal 값은 RMS 정규화에 의해 제한된다. 구현에서는 cuBLAS의 TRSM 커널을 활용해 fp32 연산을 빠르게 수행한다. 실험에서는 1 B 파라미터 모델을 2K→65K 토큰 길이로 연속 사전학습한 뒤 128K까지 확장했으며, 수치 안정성 문제가 전혀 보고되지 않았다.

실험 결과는 두 가지 주요 벤치마크에서 두드러진 향상을 보여준다. BABILong에서는 최대 18%의 정확도 상승, RULER의 multi‑needle 설정에서는 14% 향상을 기록했다. 또한, SCROLLS와 LongBench의 장문 검색·요약 태스크에서도 기존 Path Attention·DeltaNet·Differential Transformer 대비 일관된 우위를 보였다. 합성 실험에서는 두 단계(자기 복제 → 누적 평균) 학습 시, 표준 소프트맥스는 Jacobian이 급격히 감소해 두 번째 단계에 적응하지 못하는 반면, LUCID는 높은 Jacobian을 유지해 빠르게 전환한다.

요약하면, LUCID는 키‑키 상관을 RKHS에서 사전조건화함으로써 소프트맥스의 근본적인 한계를 극복하고, 긴 컨텍스트에서도 정확하고 학습 가능한 주의 메커니즘을 제공한다. 이는 장문 검색·추론·인-컨텍스트 학습 등 미래 대규모 언어 모델의 핵심 과제에 직접적인 영향을 미칠 수 있다.