Extender가 어텐션 KV 캐시를 104분의 1로 줄인다
The Extender: A Log-Structured Transformer
무엇인가
표준 Transformer에서 층 사이의 유일한 통신 채널은 잔차 스트림 h다. 모든 층의 입력이 h이고 출력도 h에 대한 업데이트이므로, 추론 시 각 어텐션 층은 자기 몫의 full-width 잔차 또는 KV 캐시를 따로 유지해야 한다. 이 캐시는 모델 폭 d_model, 깊이 L, 문맥 길이 T에 선형으로 비례해 커지고, 대형 모델에서는 차원 축소·층 재사용·양자화를 적용하기 전에 100GB를 넘길 수 있다. 논문이 겨냥하는 것은 특히 턴과 턴 사이에 계속 붙들고 있어야 하는 persistent 어텐션 메모리다. 기존 GQA, MLA, YOCO/LCKV/CLA 같은 층 재사용, 양자화는 이 footprint를 줄이지만 어텐션 메커니즘 자체를 바꾸지는 않는다.
어떻게 동작하나
Extender는 통신 채널을 둘로 분리한다. 하나는 기존 잔차 h로, 각 층이 full-width 업데이트 δ_ℓ을 더해 누적하고 다음 토큰 예측과 최종 로짓에 쓰인다. 다른 하나는 로그 구조 채널 x로, 각 층이 훨씬 작은 확장 ε_ℓ을 x 끝에 덧셈이 아니라 이어붙이기(concatenation)로 추가한다. FFN과 어텐션 쿼리 q는 h를 보지만, 어텐션의 k·v 투영은 x만 입력으로 받는다. x는 append-only이므로 모든 층을 통과한 x_*는 각 층의 kv 투영 입력을 전부 prefix로 포함하고, 따라서 층마다 다른 x_ℓ를 보면서도 단일 공유 x_*-캐시 하나로 모든 어텐션 서브층을 커버할 수 있다. 그 결과 persistent 어텐션 메모리가 2L·d_model에서 Σ|ε_ℓ|로 줄어든다.
무엇과 다른가
구체적인 구성은 다음과 같다. 각 층은 x_{ℓ-1}에서 최근 d_model개 특징을 잘라낸 윈도 s_ℓ을 읽는데, 기본값은 Last() 슬라이딩 윈도다. FFN은 δ_ℓ과 ε_ℓ 두 개의 출력을 내므로 출력 폭이 d_model + d_ε가 된다. 어텐션 쿼리는 RMSNorm(s_ℓ) + RMSNorm(h_{ℓ-1})을 입력으로 받고, 어텐션 출력은 선형 가중치를 거쳐 h에 더해진다. 로짓은 h에서만 나오므로 마지막 층은 ε을 내지 않는다. 기본 설정은 d_ε=32, ε을 내는 마지막 층 ℓ_max = L-2, 첫 층만 |ε_0|=64다. 26층·폭 1664 모델에서 Σ|ε_ℓ|은 800 특징이고 s_25에는 입력 임베딩용으로 864 특징이 남는다. d_ε=64면 Σ|ε_ℓ|=1600이 되어 후반 층의 기록이 초기 임베딩을 밀어낼 수 있다. 저자들은 32보다 작거나 64보다 큰 값은 성능이 크게 떨어진다고 보고한다. k 투영에는 Last 윈도가 항상 낫지만, v 투영은 항등 매핑이 중요한 과제가 많아 Random이 절충안이 될 수 있다고 본다.
어떻게 쓰나
추론에서는 x_*-캐시와 별도로 ephemeral KV 캐시를 둔다. 어텐션을 계산할 때 필요한 k·v를 x_*-캐시에 W_k, W_v를 곱해 그때그때 재구성하고, 이 비용을 여러 입출력 토큰에 걸쳐 분산시키면 작아진다. 따라서 턴이 끝나면 ephemeral KV 캐시를 해제할 수 있다. 재투영 비용은 총 KV 차원에 선형이므로 GQA를 쓰면 GQA가 절약하는 공간만큼 재투영도 빨라진다. 턴이 진행 중일 때의 메모리 요구량은 표준 Transformer와 동일하며, 기존 차원 축소·층 재사용·양자화를 그대로 얹을 수 있다.
전제와 한계
실험은 ClimbMix와 ProLong 데이터로 학습한 모델을 짧은 문맥 DCLM CORE와 긴 문맥 RULER에서 평가한다. 비교 대상 Reference Transformer는 Llama 스타일 MHA에 nanochat 설계를 따라 Muon 옵티마이저와 어텐션 SoftCap을 적용한 모델이다. 198M/436M/920M 세 크기에서 Chinchilla 기준 파라미터당 약 20토큰(각 4B/9B/18B 토큰)으로 처음부터 학습했고, 파라미터 수가 1% 이내로 맞춰진 조건에서 CORE 점수는 Extender와 Transformer가 비슷했다. RULER에서는 920M 모델 기준으로 Extender가 평균적으로 앞섰다. 64k 문맥에서 Extender의 persistent 어텐션 메모리는 약 5450만 특징(bf16 109MB)인 반면 MHA Reference Transformer는 56.7억 특징(bf16 11.3GB)으로 104배 차이다. 세부적으로 Extender는 어려운 multi-key 과제에서 큰 우위를 보였고, variable tracking 과제에서는 Transformer가 확실히 이겼다. ClimbMix 사전학습 직후, ProLong 학습 이전 단계에서는 2k 학습 문맥을 넘는 길이에서 Extender가 유의하게 우세했고, ProLong 이후에도 긴 문맥 과제에서 Transformer가 Extender를 앞서지 못했다.
절제 실험은 d_model=1024, L=13, d_ε=32의 약 200M Extender를 ClimbMix로 4B 토큰 학습해 수행했다. 어텐션 쿼리 입력을 h만 또는 x만으로 주면 긴 문맥 과제에서 유의하게 떨어졌고, 둘을 이어붙여 W_q를 두 배로 키운 경우에는 뚜렷한 향상이 없었다. FFN 입력에 h를 포함하면 성능이 올라가지만, h를 완전히 빼도 손실이 치명적이지는 않았다. 즉 층간 통신의 상당 부분이 h가 아니라 x로 전달된다. 다만 out-of-domain 과제에서는 FFN 입력보다 쿼리에 h가 들어가는 것이 훨씬 중요했다. 학습 안정성을 위해 잔차 쓰기가 폭주하지 않도록 1/L Σ relu(RMS(δ_ℓ) - τ)² (τ=8, 계수 3×10⁻³) 정규화 항을 각 토큰에 적용한다.
개발자 관점에서 이 논문은 다중 턴 서빙이나 장문 컨텍스트 추론에서 턴 사이에 유지되는 KV 캐시 메모리가 병목인 상황을 겨냥한다. 어텐션 메커니즘을 바꾸지 않고 어텐션이 읽는 입력만 바꾸므로, 기존 GQA·양자화·층 재사용 최적화와 배타적이지 않다. 다만 턴이 진행되는 동안의 메모리는 표준 Transformer와 같고, x_*-캐시에서 k·v를 다시 투영하는 연산이 추가된다는 점을 감안해야 한다. 또한 d_ε, ℓ_max, k/v 윈도 선택이 성능에 민감하고, 잔차 쓰기 정규화 같은 학습 레시피 조정이 필요하다.
저자들이 밝힌 한계는 명확하다. d_ε=32 구성이 199M에서 924M까지는 잘 확장되지만, 자원 제약 때문에 그보다 훨씬 큰 모델 크기는 아직 검증하지 못했다. 토큰 예산도 대체로 파라미터 수의 20배 수준에 묶여 있다. RULER 점수는 절대값 자체가 낮은데, 이는 instruction-following이나 RULER용 파인튜닝을 하지 않은 사전학습 단계 모델을 비교했기 때문이며, 논문의 목적도 점수 극대화가 아니라 두 아키텍처의 동등 조건 비교라고 밝히고 있다.