안녕하세요. KDST 한용진입니다.
이번에 소개해 드릴 논문은 2026년도 ICLR에 발표된 "KAVA: LATENT REASONING VIA COMPRESSED KV-CACHE DISTILLATION" 입니다.
본 논문을 소개해 드리게 된 배경은 다음과 같습니다. 우리가 사용하는 KV cache가 중요한 정보를 잘 담고 있다면, KV cache를 다양한 태스크에 활용할 수 있지 않을까?라는 궁금증에서 시작되었습니다.
LLM은 Chain-of-Thought(CoT)를 활용해 추론과 같은 복잡한 문제를 잘 풀게 되었습니다. 그러나 일반적으로 사용하는 CoT는 자연어로 구성되어 있어 자연어가 가지고 있는 무의미한 단어들을 처리하는 데 불필요한 연산을 사용하는 문제점이 존재합니다. 이러한 문제를 해결하기 위해 CoT를 latent variable로 변환하여 추론하는 Latent reasoning 방법이 등장하게 됩니다. 하지만 변환된 latent variable은 사람이 해석할 수 없기 때문에 latent로 바뀌는 과정이 제대로 이루어지는지 확인할 수 없다는 문제가 있습니다.
본 논문에서는 KV cache에 내재된 정보를 활용하여 latent reasoning에서 발생하는 문제(lack of supervision)를 해결하는 방법을 제안합니다.
논문에서 제안하는 KAVA는 먼저 teacher의 KV cache를 압축하여 중요한 정보들만 캐시에 남긴 뒤, teacher의 압축된 KV cache를 student의 KV cache로 지식 증류를 함으로써 latent reasoning 모델을 학습시킨다. 이러한 과정에서 R-KV [1]와 CODI [2] 방법이 함께 사용되는데 두 방법을 간단하게 알아보도록 하자.
R-KV
추론 모델의 출력 결과는 일반적으로 장문으로 구성되기 때문에, 추론 시 캐싱된 값들이 상당히 반복적이라는 문제가 존재한다. R-KV는 Importance score(어텐션 스코어)와 Redundancy score를 통해 반복적이지 않으면서도 중요한 정보들만 저장하는 KV compression method를 제안한다.
R-KV(Redundancy-aware KV Cache Compression for Reasoning models)
추론 모델에서 장문의 출력을 생성함에 따라 반복되는 캐싱 문제를 해결하기 위한 방법을 제안한다. Importance score는 일반적인 어텐션 스코어를 사용하며 Redundancy estimation에서는 코사인 유사도를 활용한다. 두 평가지표를 활용하여 중요도와 반복성을 \lambda로 제어하여 적절한 전략을 탐색한다. 계산된 점수를 통해 Top-k의 캐시들만 선택한다.
실제 구현에서는 KV cache를 일정 예산만큼만 압축한 뒤, 나머지는 새로운 토큰값으로 채운다. 이를 통해 최신으로 들어온 토큰 정보를 유지하면서 이전에 있던 토큰들은 반복적이지 않고 중요한 토큰들만 남기도록 한다.
CODI
Explicit CoT는 자연어로 구성되어 있어 중요하지 않은 단어/토큰도 처리해야 하는 문제점이 존재한다. 이러한 문제를 해결하기 위해 Latent reasoning이 제안되었지만 latent variable은 해석할 수 없기 때문에 제대로 추론이 이루어지고 있는지 확인하기 어렵다. CODI는 teacher와 student 모델에 KV cache를 두고 teacher의 KV cache를 supervisor로 활용하는 방법을 제안한다.
CODI(Continuous Chain-of-Thought via Self-Distillation)
일반적으로 CoT 기반 추론 모델은 자연어 토큰에 의존하는 반면, 프롬프트를 continuous representation으로 표현하면 짧아진 프롬프트로 인해 더 효율적인 추론이 가능해진다. CODI는 teacher task와 student task를 두고 증류를 하는 방법을 제안한다. 또한 증류를 위한 손실 함수로 정답 직전 토큰의 hidden state를 사용하는 데, 이는 정답 직전의 토큰이 정답을 유추하는 데 가장 큰 기여를 하기 때문이라고 한다. 아래 그림을 예시로 들면, 정답 직전의 토큰은 ":"이므로 ":"에 해당되는 hidden state를 정렬하는 손실 함수를 학습한다.
KAVA가 학습하는 손실함수는 다음과 같다: 1) student/teacher loss, 2) CODI loss, 3) KV distillation
1) student/teacher loss
먼저 teacher loss는 질문이 주어졌을 때 올바른 추론 과정(CoT, C)과 올바른 답(A)을 생성하는지 확인하는 손실 함수이며 CrossEntropy loss를 사용한다. student loss는 질문과 latent variable로 표현된 추론 과정(implicit CoT, Z)이 주어졌을 때 올바른 답(A)을 생성하는지 확인하는 손실 함수이다. 두 손실 함수의 차이점은 C/Z의 위치인데, 이는 latent reasoning의 단점인 supervision이 부족한 것을 보완하기 위한 조치라고 보면 되겠다.
2) CODI loss
CODI loss는 student/teacher 모델에서 정답 직전 토큰의 hidden state를 정렬시키기 위한 손실 함수이다. Teacher 모델에서 student 모델로 지식을 증류시킬 것이므로 sg(stop gradient)를 통해 teacher에 gradient가 흐르지 않도록 한다.
3) KV distillation
KV distillation loss는 R-KV로 압축된 teacher의 KV cache에 내재된 정보들을 student의 KV cache로 증류하는 함수이다. CODI loss와 유사하게 손실 함수가 설계되었다. 이때 p는 1 또는 2의 값을 가지며, student의 KV cache 사이즈(M)에 맞추어 압축이 되었기 때문에 M으로 두 번 나눈다. (key, value 한 번씩)
Table 1은 GSM8k-AUG와 GSM8k-AUG-NL에 대한 실험 결과이다. KAVA가 대체적으로 더 좋은 성능을 보이고 있으며, 좀 더 어려운 데이터셋인 GSM8k-AUG-NL에 대해서도 정확도 감소가 덜하다는 장점이 있다고 저자들이 주장한다.
Figure 4는 explicit CoT 토큰을 일정 비율만큼 latent token으로 변경하여(20%->40%->60%->80%) 정확도와 효율성 간의 상관관계를 측정한 결과이다. x축은 CoT를 생성하는 데 필요한 forward pass의 수를 의미한다. 그림을 통해 알 수 있듯이, KAVA가 다른 방법들보다 더 좋은 Pareto frontier를 보여주고 있다.
Figure 5는 KV distillation loss에서 p를 어떤 값으로 설정해 주었는지에 따른 정확도를 보여준다. 신기하게도 p=1 (l1 norm)을 사용했을 때 더 좋은 성능을 보인다고 한다. 그러나 저자들은 이 부분에 대해서 의미 있는 해석을 제공하지는 않고 있다.
Figure 6은 R-KV를 통해 teacher의 KV cache를 압축할 때 redundancy에 초점을 맞출 것인지 importance에 초점을 맞출 것인지를 확인하는 실험이다. 그림을 통해 확인할 수 있듯이 두 평가 지표를 적절하게 모두 고려하는 것이 좋은 성능을 낸다는 것을 보여준다.
Figure 7은 latent CoT(implicit CoT)와 정답 CoT 간의 코사인 유사도를 측정한 결과이다. 왼쪽 그림은 R-KV로 KV cache를 압축하지 않은 상태에서 KV distillation loss로 학습한 결과인데, 정답 CoT와의 유사도가 거의 없는 반면, 압축된 KV cache로 학습한 결과 대각 행렬의 값이 높은 것을 통해 정답 CoT와 latent CoT의 유사도가 높은 것을 확인할 수 있다.
KV cache를 압축(compression) 또는 축출(eviction)하는 과정에서 레이어/헤드 간의 관계가 고려되지 않아 입력 토큰과의 직접적인 연계가 어려운 점이 존재한다. 그럼에도 불구하고 압축/축출된 KV cache에 있는 정보를 통해 더 좋은 추론 성능을 낼 수 있다는 점에서 유의미한 논문이라고 생각한다.
본 글에서 다루지 않은 내용 및 실험이 있으니 더 깊은 이해를 위해 논문 읽기를 권장드립니다.
https://openreview.net/pdf?id=ePrhcLbtGv
긴 글 읽어주셔서 감사합니다.
참고문헌
[1] Cai, Zefan, et al. "R-kv: Redundancy-aware kv cache compression for reasoning models." arXiv preprint arXiv:2505.24133 (2025).
[2] Shen, Zhenyi, et al. "Codi: Compressing chain-of-thought into continuous space via self-distillation." Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing. 2025.