BeNice!는 루프 모델의 KV 캐시 공유로 추론 메모리 증가를 막는다

무엇인가

이 논문은 루프형(looped) 언어모델의 추론 메모리 문제를 다룬다. 루프 구조는 같은 블록을 여러 번 통과해 파라미터 수를 고정한 채 품질을 높이지만, 추론 시 각 반복이 독립 KV 캐시를 만들어 KV 메모리가 반복 횟수 T에 선형으로 증가한다. 기존 연구는 마지막 반복의 KV 캐시를 이전 반복에 재사용하는 KV-sharing이 일부 모델에서 품질 손실 없이 가능하다고 관찰했지만, 왜 어떤 모델은 되고 어떤 모델은 안 되는지, 동일 학습 설정에서도 왜 달라지는지는 설명되지 않았다.

어떻게 동작하나

저자들은 KV-shareability가 반복 블록의 contractive residual dynamics, 즉 잔차 상태가 고정점으로 수렴하는 성질에서 나온다고 본다. 은닉 상태 h_t가 고정점 h_hat으로 수렴하면 key/value 투영도 수렴하고, 중간 캐시를 마지막 캐시로 바꿨을 때 attention 출력 변화가 입력 잔차 차이 또는 고정점 잔차로 bound된다는 분석을 제시한다. 그래서 마지막 반복 캐시가 고정점에 가장 가까워 broadcast 대상으로 최적이라고 주장한다.

무엇과 다른가

이를 유도하려고 1단계에서 stochastic depth 사전학습을 쓴다. 배치의 각 시퀀스마다 깊이 T를 Poisson(mu_rec)에서 샘플링하고, backprop depth는 floor(mu_rec/2)로 제한한다. 여러 깊이에서 같은 readout으로 낮은 손실을 내야 하므로 잔차가 드리프트하지 않고 수렴하도록 압력을 준다. 하지만 stochastic depth만으로는 충분하지 않다. 770M·1.3B Parcae처럼 고정점 수렴은 보여도 KV-sharing 시 품질이 떨어지는 경우가 있었고, 저자들은 KV-sharing이 첫 반복 첫 레이어에서 토큰 독립적인 basis shift를 만들어 공유/비공유 고정점을 다르게 만든다고 분석한다. 공유/비공유 최종 잔차의 평균 상대 차이는 공유 가능 모델이 최대 1.69%, 실패 모델이 최대 177.5%였다.

어떻게 쓰나

2단계 BeNice!는 이 첫 반복 오차를 교정하는 후학습이다. 손실은 CE(y, p_theta^(T)) + gamma CE(y, p_theta^(1<-T))로, 전자는 기존 T회 반복 예측의 품질을 유지하고 후자는 첫 반복의 query/residual을 쓰되 마지막 반복의 KV 캐시를 쓰는 hybrid predictor를 학습시킨다. 즉 KV-sharing rollout을 토큰 단위로 직접 돌리지 않고도 첫 반복이 미래 반복의 KV와 호환되는 잔차를 내도록 만든다. T는 사전학습과 같은 깊이 분포에서 샘플링한다. MLA나 KDA에는 KV 캐시 대신 마지막 반복의 latent 또는 state를 재사용하도록 확장한다.

전제와 한계

실험은 Parcae-style 140M, 370M, 770M, 1.3B 모델을 100B 이상 ClimbMix 토큰으로 학습해 진행했다. 테스트 시 반복은 T=mu_rec. KV-sharing 실패는 모든 스케일에서 나타나 Core와 Core-Extended를 최대 5.58점, 3.31점 떨어뜨렸다. BeNice!는 공유 시 Core와 Core-Extended를 최대 8.35점, 4.56점 개선했고, 비공유 조건에서는 표준 finetuning과 비슷한 품질을 유지했다. 공유할 때 validation perplexity는 최대 45.2%(9.75 PPL 포인트) 낮췄다.

KV-sharing은 다른 KV 압축과 직교한다고 주장한다. MLA와 KDA sequence mixer에서도 BeNice!가 KV-shareability를 가능하게 했고, 모든 스케일·mixer에서 공유로 인한 Core/Core-Extended 변화는 평균 -0.09점, -0.04점이었다. 저정밀 KV 양자화와 조합했을 때도 Core/Core-Extended 변화는 평균 -0.02점, +0.06점, ClimbMix validation perplexity 증가는 평균 0.67%였고, 양자화와의 상호작용은 최대 0.53% perplexity로 보고됐다. 추론에서는 sglang에 KV-sharing을 구현했고, 1.3B Parcae를 Core-Extended, 배치 64, 문맥 2048로 돌릴 때 global HBM 사용량이 63.82GB에서 21.60GB로 66.2% 줄었다. mu_rec=8에서 토큰 용량은 3.33배, mu_rec=16에서 6배가 됐다. Mooncake 요청을 AIPerf로 재생한 시뮬레이션에서 2 req/s일 때 캐시 적중률 8.1포인트(최대 31.4%), 처리량 23.6% 증가, p90 TTFT는 139.2초에서 23.2초로 83.6% 감소했다. 동시성 frontier에서는 고정 median decode rate에서 처리량 6.7% 증가, 고정 처리량에서 median decode rate 34.8% 증가를 보였다.

개발자에게 이 논문은 루프형·반복 깊이 모델을 긴 문맥으로 서빙할 때 KV 캐시가 병목이면, 모델을 stochastic depth로 사전학습하고 BeNice! 후학습을 적용해 반복별 KV 캐시를 하나로 접는 선택지를 보여준다. 다만 임의의 공개 모델에 바로 적용하는 기법은 아니고, 중간 루프 구조와 stochastic-depth 사전학습, 그리고 별도 후학습이 전제다. 원문은 별도 한계 절을 두지 않았지만, 품질 실험은 최대 1.3B Parcae-style 모델과 ClimbMix/Core 벤치마크에 한정되고, 추론 이득은 sglang 구현과 Mooncake 기반 시뮬레이션, GLM 4.7-FP8 블록을 루프시킨 proxy 모델에서 측정됐다. 실제 도입 전에는 자신의 모델에서 공유/비공유 품질 차이, sequence mixer·양자화 조합, 메모리 제약과 부하 조건을 확인해야 한다.