WhiteMatter가 모든 층의 KV 재사용으로 캐시를 절반으로 줄인다

무엇인가

Transformer 디코더는 각 층이 자기 깊이에서 만든 키·값(KV)만 보고 과거 토큰에 어텐션한다. 이 제약 덕분에 학습과 프롬프트 처리는 병렬화하기 쉽지만, 다른 깊이에서 이미 계산해 둔 표현을 재사용하지 못한다. 저자들은 추론 비용이 학습 비용을 넘어서고, 특히 에이전트와 긴 추론 트레이스 생성에서 디코딩이 추론 시간을 지배하는 상황을 문제로 든다. 기존 시도인 Feedback Transformer는 과거 토큰의 모든 층 상태를 하나의 가중합으로 압축해 공유 KV를 만들고, LCKV는 가장 깊은 층 상태로 공유 KV를 만든다. 저자들은 두 방식 모두 여러 층의 상태를 하나의 층 폭 표현으로 압축하는 공유 병목을 갖는다고 지적하며, 층마다 필요한 소스 깊이가 다르다는 가설을 세운다.

어떻게 동작하나

제안 방법의 핵심은 층별 KV 투영을 cross-layer KV pool로 교체하는 것이다. L개 층의 은닉 상태를 라우터가 k(≤L)개 채널로 혼합하고, 각 채널이 자체 KV 투영을 갖는다. 층 ℓ은 채널 j = ℓ mod k 하나만 읽으므로 KV 캐시 크기는 vanilla의 k/L이 된다. 키 브랜치와 밸류 브랜치는 별도 파라미터로 같은 구조를 쓴다. 각 소스 상태는 먼저 RMSNorm으로 정규화되어 층 간 스케일을 맞추고, 라우터는 파라미터 수를 줄이기 위해 p번째 층의 상태만 읽지만(stride p=2) 혼합 가중치는 L개 층 전체에 대해 생성한다. 가중치는 부호를 가질 수 있어 층 표현 간 차이까지 표현한다. 혼합된 채널은 다시 RMSNorm을 거쳐 K와 V로 투영되고, 키에는 채널별 정규화와 RoPE를 적용한 뒤 캐시된다.

무엇과 다른가

학습과 prefill에서는 순환 의존을 피하기 위해 엄격한 인과 어텐션을 쓴다. 토큰 i는 s<i 위치만 보고, 층 스택 평가가 끝난 뒤에 자기 채널을 만들어 캐시에 추가한다. 이때 캐시는 KV = Model(X; KV)라는 고정점으로 볼 수 있어 반복으로 푼다. Jacobi 반복은 모든 토큰을 병렬 평가하지만 갱신이 다음 패스에만 전파되어 수렴에 많은 전체 시퀀스 평가가 필요하다. 저자들은 cyclic Gauss–Seidel 반복을 도입해 토큰 위치를 g개 그룹으로 나누고 그룹을 순서대로 처리하되 그룹 내부는 병렬로 평가한다. 그룹 배정은 G_q = {i : i mod g = q}로 순환시키고 q=0…g-1 순서로 처리해 인접 위치 간 전파를 돕는다. 그룹마다 KV를 갱신하므로 같은 패스 안에서 뒤 그룹이 앞 그룹의 갱신을 읽는다. 학습에서는 마지막 n_g개 패스에만 그래디언트를 흘리는 truncated backpropagation을 쓴다.

어떻게 쓰나

실험은 Qwen3 기반 디코더 두 규모에서 진행했다. D=512에서는 16층 WhiteMatter(full k=16, half k=8)를 vanilla, LCKV, FusedKV, 그리고 24층 vanilla와 비교했고, D=1792에서는 28층 vanilla(1.351B)와 half-cache WhiteMatter(k=14, 1.326B)를 비교했다. 데이터는 Qwen3 토크나이저로 2048 길이로 패킹한 shuffled FineWeb-Edu이며, 두 규모에 각각 8B와 10B 토큰을 같은 예산으로 학습했다. WhiteMatter는 그래디언트 없는 cyclic 패스 1회 + 그래디언트 패스 2회, LCKV는 그래디언트 없는 Jacobi 패스 7회 + 그래디언트 패스 2회를 사용했다. 결과적으로 half-cache WhiteMatter는 같은 조건 vanilla 대비 held-out perplexity를 작은 규모에서 5.9%, 1.3B 규모에서 4.3% 낮췄고, 작은 규모에서는 같은 캐시 크기의 LCKV와 FusedKV도 앞섰다. full-cache WhiteMatter는 층이 50% 더 많은 vanilla와 비슷한 perplexity를 냈고, 같은 층수 vanilla보다는 7.7% 낮은 perplexity를 기록했다. zero-shot 평균 점수도 두 규모 모두 같은 조건 vanilla를 상회했고, 작은 규모에서는 50% 더 깊은 vanilla까지 포함해 모든 베이스라인을 앞섰다.

전제와 한계

수렴 속도는 정확한 자기회귀 실행으로 78.6M 토큰을 학습한 4층 참조 모델에서 측정했다. cyclic g=16은 목표 perplexity에 4패스, 7.32ms/시퀀스로 도달한 반면 Jacobi는 53패스, 91.20ms가 걸려 12.5배 빨랐다. 그룹 수를 늘리면 필요한 패스는 줄지만 패스당 비용이 커져 가장 큰 그룹 수가 최저 런타임을 주지는 않았다. 1.3B 모델, 배치 64, RTX A6000 BF16 기준으로 WhiteMatter의 prefill 처리량은 vanilla의 31%, LCKV의 1.78배, Feedback Transformer의 2.92배였고 피크 장치 메모리는 13.44GiB로 vanilla의 21.21GiB보다 낮았다. 디코딩은 네 아키텍처 모두 약 2,500 tokens/s로 비슷했지만 WhiteMatter의 피크 메모리는 10.05GiB로 vanilla 16.59GiB 대비 39.4% 줄었다.

분석 실험(16층 D=512, 20,000 스텝, 327.7M 토큰)에서는 채널 수 k를 1, 2, 4, 8, 12, 16으로 늘릴수록 품질이 좋아졌고 특히 k=1에서 2로 갈 때 개선이 컸다. 층별로 다른 혼합을 하나의 공유 혼합으로 바꾸면 k=16 독립 KV 투영을 유지해도 k=4 WhiteMatter보다 나빴다. 동적 라우터를 정적 학습 가중치로 바꾸면 perplexity가 k=1에서 3.0%, k=16에서 1.9% 올랐다. 층 ℓ이 0…ℓ 상태만 보도록 제한하면 full-cache에서도 vanilla보다는 좋았지만 full-cache WhiteMatter보다 perplexity가 4.1% 높았다. 계산량은 D=512에서 디코딩 FLOPs가 vanilla의 0.99–1.03배인 반면, 학습은 2.32–2.50배, 3패스 prefill은 3.05–3.30배였다.

실무 관점에서 이 논문이 겨냥하는 지점은 KV 캐시 메모리다. 캐시가 병목인 롱컨텍스트 서빙이나 다중 에이전트 추론에서 층 간 KV 공유로 메모리를 줄일 여지가 있다는 것이 핵심 주장이다. 다만 prefill 처리량이 vanilla의 31% 수준이고 학습 FLOPs가 2배 이상이라는 점은 프롬프트 처리 지연과 학습 비용에 민감한 서비스에서는 그대로 적용하기 어렵다는 뜻이다. 도입을 검토한다면 배치 크기와 프롬프트 길이에 따른 prefill 지연, 반복 패스 수에 따른 품질·비용 트레이드오프를 자체 워크로드에서 다시 측정해야 한다. 코드는 공개되어 있다.

저자들이 밝힌 한계는 분명하다. 자기회귀 실행이나 Jacobi 반복 대비로는 빨라졌지만 반복적 학습과 prefill은 여전히 표준 Transformer보다 비싸다. 품질 결과는 최대 1.3B 파라미터, 10B 토큰 학습까지이며 더 큰 모델과 더 많은 데이터에서의 효과는 추가 검증이 필요하다고 명시했다. full-cache 구성과 대안 KV 공유 베이스라인 비교는 작은 규모에서만 수행했다.