KronLSTM이 스칼라 보정으로 긴 문맥 선형 어텐션의 기억 간섭을 줄인다

무엇인가

Transformer의 softmax 어텐션은 문맥 길이에 대해 계산 비용이 제곱으로 늘고 KV 캐시는 선형으로 커진다. 이를 줄이려는 선형 어텐션은 문맥을 고정 크기 행렬 기억으로 압축하고 토큰마다 갱신한다. 문제는 기억 성장을 어떻게 통제하느냐다. 망각 게이트는 긴 문맥에서 1에 가깝게 유지해야 해 점점 제약이 커지고, xLSTM의 mLSTM은 정규화 상태와 안정화 상태를 따로 두지만 큰 네트워크에서 수치적 불안정이 보고된다. DeltaNet 계열은 델타 규칙으로 기억을 유계로 유지하지만, 저장된 값을 수정할 때 이미 저장된 다른 연관까지 흔든다는 간섭 문제가 남는다.

어떻게 동작하나

이 논문의 핵심은 크로네커 델타 규칙이다. 델타 규칙은 d_v 차원의 잔차 S_t k_t − v_t 를 빼서 기억을 고치기 때문에, 현재 기억 내용에 의존하는 방향으로 수정이 일어나고 k_t 와 겹치는 키에 저장된 값도 함께 영향을 받는다. 크로네커 델타 규칙은 대신 스칼라 하나, 즉 v_t^T S_t k_t 라는 연관 강도를 목표 c_t 쪽으로 보정하고, 그 방향을 입력 쌍만으로 결정되는 외적 v_t k_t^T 로 잡는다. 갱신식은 기존 연관을 지우는 erase 항과 새 값을 쓰는 write 항으로 분리되며, 목표를 c̃_t = c_t/β_t 로 두면 지우는 양은 β_t 가, 쓰는 양은 c_t 가 각각 전담하게 된다.

무엇과 다른가

이 규칙은 벡터화하면 a_t = k_t ⊗ v_t 를 입력으로 하는 Widrow–Hoff 델타 규칙이 되고, 전이 행렬이 일반화 하우스홀더 행렬이라 DeltaNet의 청크 병렬 학습 방식을 그대로 쓸 수 있다. 간섭 분석이 중요한데, 다른 저장 쌍 (k_i, v_i) 의 연관 변화량은 값 겹침과 키 겹침의 곱에 비례한다. 즉 둘 중 하나만 직교해도 기존 연관이 보호되는 반면, Hebbian 규칙과 델타 규칙의 간섭은 키 겹침만으로 결정된다. 용량 측면에서도 각 삼중항이 d_k·d_v 개 항목에 선형 제약 하나를 걸기 때문에, 특징이 선형 독립이면 최대 d_k·d_v 개의 스칼라 연관을 정확히 저장할 수 있다. 하나의 키가 서로 직교하는 여러 값을 담을 수도 있다.

어떻게 쓰나

KronLSTM은 이 규칙을 게이트 선형 어텐션 믹서에 넣은 구조다. 키 차원에 작용하는 대각 망각 게이트 F_t 가 기억을 오른쪽에서 감쇠시키고, 지우기 게이트 β_t 와 쓰기 게이트 c_t 는 입력의 선형 사영에 시그모이드를 씌워 헤드별 스칼라로 만든다. 망각 게이트는 시그모이드 값을 헤드별로 학습되는 양의 지수 a = exp(a_log) 로 거듭제곱해, 헤드마다 유지 시간 스케일을 크게 다르게 가져가면서도 높은 프리액티베이션 없이 조절한다. 쿼리·키·값은 짧은 인과적 depthwise 컨볼루션과 SiLU를 거쳐 L2 정규화되고, 읽기 결과는 시그모이드 출력 게이트와 헤드별 RMSNorm을 지나 출력 사영으로 섞인다. 키와 값이 정규화되어 있고 β_t 가 시그모이드이므로 β_t‖k_t‖²‖v_t‖² = β_t < 1 이 성립해 모든 갱신이 비확장적이다. mLSTM과 달리 정규화 상태도 안정화 상태도 없이 각 헤드의 기억 행렬만 상태로 들고 간다.

전제와 한계

이론 쪽에서는 i.i.d. 삼중항 스트림을 가정해 기대 반복자가 최소제곱 기억 S_LS 로 기하급수적으로 수렴하고, 마르코프 체인이 유일한 정상 분포를 가진다는 것을 보인다. 기억 노름도 유계인데, 비확장 영역 κ ≤ 1 에서 S_0 = 0 으로 시작하면 모든 시점에서 E[‖S_t‖²_F] ≤ (1 + q_β^{t/2})² · E[c²]/λ_min(Ā) ≤ 4E[c²]/λ_min(Ā) 를 만족한다. 과도 구간 계수가 4에서 1로 기하급수적으로 줄어들 뿐, 데이터 분포만으로 결정되는 상수에 처음부터 묶인다는 뜻이다.

실험은 1.5B 파라미터 모델을 FineWeb-Edu에서 Llama2 토크나이저로 100B 토큰 학습해 진행했다. 비교 대상은 Gated DeltaNet, Kimi Delta Attention(KDA), Mamba-3, Gated DeltaNet-2 이고, 순수 선형 모델과 윈도 2048 슬라이딩 윈도 어텐션을 5:1로 섞은 하이브리드 두 설정을 모두 평가했다. 상식 추론 5개 과제와 LAMBADA에서 KronLSTM은 6개 과제 평균 정확도가 가장 높고 LAMBADA 퍼플렉시티가 가장 낮았다. 순수 선형 설정에서 6개 중 5개 과제 최고 성능이며, 가장 강한 경쟁자 KDA 대비 평균 1.6 포인트 앞선다. 하이브리드에서도 6개 중 4개 과제에서 앞섰다. 2k 문맥의 실세계 리콜 스위트(SWDE, SQuAD, FDA, TriviaQA, NQ, DROP)에서도 평균 정확도가 가장 높고 3개 과제에서 1위다. RULER 니들인어헤이스택을 4k·8k·16k에서 평가한 결과, 선형 모델 중에서는 S1을 빼고 모든 길이에서 1위 또는 2위였고, 하이브리드 중에서는 모든 과제의 길이 평균과 16k 각 과제에서 최고 정확도를 냈다. MQ 과제의 리드는 문맥이 4k에서 16k로 늘 때 2.4 포인트에서 4.1 포인트로 벌어지고, MV에서는 4k에서 GDN에 3.5 포인트 뒤지지만 8k와 16k에서는 각각 1.1, 3.8 포인트 앞선다. 75B 토큰에서 분기해 나머지 25B를 32k로 이어 학습한 뒤에도 KronLSTM은 GDN-2보다 평균 정확도가 높고 퍼플렉시티가 약 4 포인트 개선되며 6개 리콜 과제 중 5개에서 앞선다. 상태 추적은 β_t = 2σ(·) ∈ (0,2) 로 확장해 parity와 S3 순열 과제에서 Gated DeltaNet과 비슷하고 Mamba-3보다 앞서지만, 모듈로 산술에서는 Mamba-3가 이긴다. 단일 H100 80GB에서 16,384 토큰 스텝 기준 처리량은 36.8~44.5K tokens/s로 GDN-2보다 11~18% 높고, 8K×2까지는 KDA의 4~8% 이내다. 가장 빠른 것은 스칼라 감쇠를 쓰는 GDN이다.

실무 관점에서 이 논문은 선형 어텐션 레이어를 Transformer 블록의 어텐션 자리에 갈아 끼우려는 팀에 유용하다. 정규화 상태 없이 기억 노름 상한이 데이터 분포로 보장된다는 점은 긴 문맥 파인튜닝에서 수치 불안정을 줄이는 근거가 되고, 청크 병렬 형식이 DeltaNet과 동일해 기존 커널 자산을 재사용할 수 있다. 다만 처리량이 GDN보다 낮고 KDA와도 격차가 있으므로, 문맥 길이와 배치 형태에 따른 실제 처리량을 먼저 재보고 도입 여부를 정하는 편이 좋다.

저자들이 인정한 약점은 S1 과제다. 노이즈 문장이 자주 반복되는 가장 인공적인 과제로, KronLSTM은 여기서 경쟁 모델에 뒤진다. 원인은 반복되는 키에 대한 잔차 누적으로 설명한다. 또한 32k로 문맥을 늘린 뒤에는 두 모델 모두 16k 대비 검색 정확도가 크게 떨어지고, 하이브리드 기준 베이스라인 선형 레이어들은 4k·8k·16k에서 각각 최대 51.8, 26.6, 12.6 정확도에 그쳐 니들을 주로 슬라이딩 윈도 안에서만 찾는다는 점도 드러난다. 저자들은 모델 확장과 완전한 softmax 어텐션 레이어를 포함한 하이브리드 학습을 향후 과제로 남긴다.