[Paper] 향상된 트랜스포머를 위한 교차 레이어 퓨전을 통한 KV 캐시 재구성

Source: arXiv - 2512.03870v1

Overview

Transformer 디코더는 오늘날의 대형 언어 모델을 구동하지만, 키‑값(KV) 캐시는 시퀀스 길이에 따라 선형적으로 증가해 긴 입력에서 GPU 메모리를 빠르게 소모합니다. 논문 *“Reconstructing KV Caches with Cross‑layer Fusion for Enhanced Transformers”*는 KV 캐시를 절반으로 줄이면서 퍼플렉시티를 개선하는 영리한 방법을 제안합니다. 하위 레이어의 가장 유용한 키와 값을 학습하여 융합함으로써, 저자들은 메모리 효율적이면서도 고성능인 새로운 디코더 아키텍처를 만들었습니다.

Key Contributions

• Cross‑layer KV fusion (FusedKV): 하위 레이어 값과 중간 레이어 키를 결합해 최상위 레이어 KV 캐시를 형성하는 학습 가능한 융합 모듈을 도입합니다.
• FusedKV‑Lite: 학습 가능한 융합을 생략하고 하위 레이어 값과 중간 레이어 키를 직접 재사용하는 경량 변형으로, 추가 I/O를 없애고 메모리 오버헤드를 더욱 감소시킵니다.
• Empirical insight: 깊은 디코더에서는 값이 주로 하위 레이어에서, 키는 하위와 중간 레이어 모두에서 유용한 신호를 얻는다는 사실을 보여줍니다.
• Memory reduction: 332 M부터 4 B 파라미터 모델에 이르기까지 KV‑cache 메모리를 약 50 % 절감합니다.
• Performance boost: 캐시를 줄였음에도 불구하고 기존 Transformer 디코더보다 낮은 검증 퍼플렉시티를 제공합니다.

Methodology

1. Diagnosing KV flow – 저자들은 표준 디코더에 계측 코드를 삽입해 각 최상위 레이어 키/값이 어디서 오는지 추적했습니다. 히트맵은 명확한 분리를 보여주었는데, 값은 하위 레이어에 크게 편중되고, 키는 하위와 중간 레이어가 혼합된 형태였습니다.
2. FusedKV design –
   • Fusion module: 작은 선형 레이어(또는 MLP)가 하위 레이어 값(V_bottom)과 중간 레이어 키(K_mid)를 혼합하는 가중치를 학습합니다.
   • Post‑RoPE fusion: 혼합은 로터리 포지셔널 임베딩(RoPE)이 적용된 후에 이루어지므로, 상대 위치 정보가 이미 벡터에 포함되어 있어 다시 계산할 필요가 없습니다.
3. FusedKV‑Lite variant – 학습 가능한 융합을 건너뛰고 V_bottom과 K_mid를 그대로 최상위 레이어 캐시에 복사합니다. 이로써 추가적인 읽기/쓰기 단계를 없애고, 약간의 퍼플렉시티 상승을 감수하면서도 지연 시간을 더욱 낮출 수 있습니다.
4. Training & integration – 융합 파라미터는 일반적인 다음 토큰 예측 작업과 함께 엔드‑투‑엔드로 학습됩니다. 어텐션 연산 자체에는 변화가 없으며, 디코더는 더 작고 융합된 캐시를 읽을 뿐입니다.

Results & Findings

• Memory: 모든 규모에서 KV 캐시가 대략 절반으로 감소하여, 동일한 하드웨어에서 컨텍스트 길이를 두 배로 늘리거나 더 큰 배치를 처리할 수 있게 됩니다.
• Quality: 검증 퍼플렉시티가 일관되게 개선됩니다(낮을수록 좋음). 이는 단순한 캐시 공유 방식보다 융합된 정보가 더 풍부함을 입증합니다.
• FusedKV‑Lite: I/O 대역폭을 추가로 ~5 % 절감하지만, 퍼플렉시티는 전체 FusedKV 대비 약 0.2 %만 상승하여 여전히 베이스라인보다 우수합니다.

Practical Implications

• Long‑context inference: 개발자는 이제 이전에 4 k 토큰으로 제한되던 단일 GPU에서 8 k16 k 토큰 프롬프트를 실행할 수 있어, 문서 수준 요약이나 코드베이스 분석 같은 사용 사례가 가능해집니다.
• Cost reduction: KV 메모리를 절반으로 줄이면 동일한 컨텍스트 길이에 필요한 VRAM이 절반이 되므로, A100‑80 GB 대신 A100‑40 GB와 같은 저렴한 GPU 인스턴스로도 동일한 워크로드를 처리할 수 있습니다.
• Deploy‑time simplicity: 융합이 모델 그래프 내부에서 이루어지기 때문에 외부 캐시 관리 코드가 필요 없습니다—디코더 클래스를 교체하기만 하면 됩니다. 이는 Hugging Face Transformers, vLLM 등 기존 추론 스택에 쉽게 통합할 수 있음을 의미합니다.
• Potential for fine‑tuning: 경량 융합 파라미터는 도메인 특화 데이터에 대해 추가 미세조정이 가능해, 메모리를 늘리지 않으면서 대형 모델을 저비용으로 적응시킬 수 있는 방법을 제공합니다.

Limitations & Future Work

• Architectural scope: 본 연구는 디코더 전용 Transformer에 초점을 맞추었으며, 인코더‑디코더 혹은 순수 인코더 모델에서는 KV 동역학이 다를 수 있습니다.
• Training overhead: 융합 모듈을 도입하면 소량의 추가 파라미터와 짧은 추가 forward pass가 발생해, 매우 큰 모델에서는 학습 시간이 약간 증가할 수 있습니다.
• Generalization to extreme scales: 실험은 4 B 파라미터까지 진행했으며, 30 B 이상 모델에서 KV 패턴이 어떻게 변할지는 아직 확인되지 않았습니다.
• Future directions: 저자들은 토큰별 가중치를 갖는 적응형 융합, 멀티‑쿼리 어텐션 또는 희소성 기반 캐시로의 확장 등을 탐구할 것을 제안합니다.

Bottom line: 키와 값이 Transformer 디코더의 서로 다른 레이어에 존재한다는 점을 인식함으로써, 저자들은 KV‑cache 메모리를 절반으로 줄이면서 모델 품질을 향상시키는 간단하지만 강력한 교차‑레이어 융합 기법을 고안했습니다. 긴 컨텍스트가 필요한 LLM 기반 서비스를 구축·배포하는 모든 사람에게 FusedKV(및 Lite 변형)는 실용적인 업그레이드가 될 것입니다.

Authors

• Hongzhan Lin
• Zhiqi Bai
• Xinmiao Zhang
• Sen Yang
• Xiang Li
• Siran Yang
• Yunlong Xu
• Jiaheng Liu
• Yongchi Zhao
• Jiamang Wang
• Yuchi Xu
• Wenbo Su
• Bo Zheng

Paper Information

• arXiv ID: 2512.03870v1
• Categories: cs.CL
• Published: December 3, 2025
• PDF: Download PDF