[Paper] 향상된 Transformers를 위한 Cross-layer Fusion을 이용한 KV Caches 재구성

Source: arXiv - 2512.03870v1

개요

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

주요 기여

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

방법론

1. KV 흐름 진단 – 저자들은 표준 디코더에 계측 코드를 삽입해 각 최상위 층 키/값이 어디서 오는지 추적했습니다. 히트맵은 명확한 분리를 보여줍니다: 값은 하위 층에 크게 편중되고, 키는 하위와 중간 층이 혼합됩니다.
2. FusedKV 설계 –
   • 융합 모듈: 작은 선형 레이어(또는 MLP)가 하위 층 값(V_bottom)과 중간 층 키(K_mid)를 혼합하는 가중치를 학습합니다.
   • RoPE 이후 융합: 혼합은 회전 위치 임베딩(RoPE)이 적용된 뒤에 이루어지므로, 상대 위치 정보가 이미 벡터에 포함돼 재계산이 필요 없습니다.
3. FusedKV‑Lite 변형 – 학습 가능한 융합을 건너뛰고 V_bottom과 K_mid를 그대로 최상위 층 캐시에 복사합니다. 이로써 추가 읽기/쓰기 단계가 사라지고, 약간의 퍼플렉시티 상승을 감수하면서 지연 시간을 더욱 낮출 수 있습니다.
4. 학습 및 통합 – 융합 파라미터는 표준 다음 토큰 예측 작업과 함께 엔드‑투‑엔드로 학습됩니다. 어텐션 연산 자체에는 변경이 없으며, 디코더는 더 작고 융합된 캐시를 읽기만 하면 됩니다.

결과 및 발견

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

실용적 함의

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

제한 사항 및 향후 연구

• 아키텍처 범위: 본 연구는 디코더 전용 Transformer에 초점을 맞추었으며, 인코더‑디코더 혹은 순수 인코더 모델은 KV 동역학이 다를 수 있습니다.
• 학습 오버헤드: 융합 모듈을 도입하면 소량의 추가 파라미터와 짧은 전방 패스가 발생해, 매우 큰 모델에서는 학습 시간이 약간 증가할 수 있습니다.
• 극한 규모에 대한 일반화: 실험은 4 B 파라미터까지 진행했으며, 30 B 이상 모델에서 KV 패턴이 어떻게 변할지는 아직 확인되지 않았습니다.
• 향후 방향: 저자들은 토큰별 가중치를 갖는 적응형 융합, 멀티‑쿼리 어텐션 또는 희소성 기반 캐시로의 확장 등을 탐색할 것을 제안합니다.

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

저자

• 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

논문 정보

• arXiv ID: 2512.03870v1
• 분류: cs.CL
• 발표일: 2025년 12월 3일
• PDF: Download PDF