LoRA 기반 트랜스포머의 서브 토큰 라우팅을 통한 효율적 KV 압축 기법

최근 대규모 언어 모델(LLM)의 효율적인 운영 및 배포가 중요해지면서, 트랜스포머 구조 내에서 핵심 정보를 유지하면서도 계산 비용을 줄이는 '압축' 기술이 주목받고 있습니다. 기존의 KV 압축 기법들은 주로 토큰 전체 단위나 헤드 단위 등 비교적 거친(coarse) 단위를 기준으로 정보 손실 여부를 결정하는 경향이 있었습니다.

본 논문은 이러한 한계를 극복하고, 트랜스포머 내부의 '토큰 표현 자체'에서 라우팅을 수행하여 훨씬 정밀한 제어 축을 확보합니다. 핵심 아이디어는 관련성 높은 토큰이라 할지라도 그 내부 정보가 균일하지 않으며, 제한된 유지 예산(retention budget) 하에서는 보존되는 값 그룹(value groups)이 토큰 간뿐만 아니라 토큰 내부에서도 불균등하게 분포한다는 점에 착안합니다. 따라서 KV 압축은 더 이상 '전부 아니면 전무(all-or-nothing)'의 이분법적 결정일 필요가 없습니다.

연구진은 이러한 미세 조정된 라우팅 메커니즘을 두 가지 주요 시나리오에서 탐구했습니다.

1. 압축 인식 언어 모델링 (Compression-aware Language Modeling):
이 설정에서는 쿼리(query)의 영향을 받지 않는 독립적인 설계(query-independent design)를 도입합니다. 핵심은 라우팅된 서브스페이스 LoRA와 **값 그룹 라우팅(value-group routing)**을 KV 경로에 결합하는 것입니다. 이 접근 방식은 언어 모델링 과정에서 품질과 압축률 사이의 트레이드오프를 효과적으로 개선함을 실험으로 입증했습니다.

2. 다운스트림 태스크 보존형 KV 압축 (Downstream-task-preserving KV compression):
이 설정에서는 쿼리의 영향을 고려하는 인식적 설계(query-aware design)를 도입합니다. 예측기 기반의 선택자(predictor-based selector)가 **쿼리 조건부 관련성(query-conditioned relevance)**을 사용하여 컨텍스트 토큰/값 그룹 쌍에 대한 전역 유지 예산을 할당합니다. 이 방식은 압축된 KV 정보를 사용하더라도 모델이 원래 수행하던 다운스트림 태스크의 행동을 높은 수준으로 보존하는 것을 목표로 합니다.

나아가, 연구진은 토큰 단위 라우팅과 서브 토큰 단위 라우팅 간의 관계를 분석하며 두 메커니즘이 상호보완적인 압축 축을 형성함을 보여주었습니다. 즉, 토큰 수준 방법론은 전역적으로 '어떤 토큰'이 생존할지 결정하는 역할을 수행하고, 서브 토큰 라우팅은 그 살아남은 토큰 내부의 정보를 어떻게 정교하게 압축할지를 결정합니다. 이러한 계층적 접근 방식은 트랜스포머 모델의 효율성과 성능을 동시에 극대화하는 데 기여합니다.