Stochastic KV Routing: 적응형 깊이별 캐시 공유 활성화

Transformer 모델의 추론 시 발생하는 KV 캐시 메모리 병목 현상을 해결하기 위해 층 간 캐시 공유를 제안한다. 기존의 시간축 압축 방식과 달리 깊이 차원의 중복성을 활용하여 메모리 사용량을 획기적으로 줄이면서도 추론 속도를 개선한다.

Transformer의 Self-Attention은 매 레이어마다 입력 토큰을 고차원 벡터인 Key와 Value로 변환하여 저장하는데, 시퀀스가 길어질수록 이 KV 캐시가 모델 파라미터보다 커지는 메모리 병목이 발생한다. 기존 연구들은 중요도가 낮은 토큰을 삭제하는 시간축 최적화에 집중했으나, 본 논문은 인접한 레이어들이 서로 유사한 정보를 인코딩한다는 깊이 차원의 중복성에 주목한다.

R-CLA는 학습 시 레이어 l이 확률 p로 자신의 KV를 생성하거나, 확률 1-p로 이전 레이어 l'의 KV를 강제로 참조하게 한다. 이는 구조적 수준에서 캐시 누락(Cache Fault)을 시뮬레이션하는 것과 같으며, 모델이 특정 레이어의 고유한 특징 공간에만 의존하지 않고 이전 레이어의 일반적인 의미 표현에서도 필요한 정보를 추출하도록 유도한다.

결과적으로 모델은 추론 시점에 일부 레이어의 캐시 생성을 건너뛰고 이전 레이어의 것을 재사용해도 성능 하락이 거의 없게 된다. 이는 메모리 대역폭 요구량을 줄여 초당 토큰 처리량(Throughput)을 높이고 더 긴 문맥을 처리할 수 있게 한다.

전체적인 접근 방식은 학습 시 레이어 간의 엄격한 KV 의존성을 깨는 것이다. 표준 Self-Attention이 Q_l, K_l, V_l을 사용하여 연산하는 것과 달리, R-CLA는 Bernoulli 분포에서 샘플링된 결정 변수 d를 도입한다.

d ~ Bernoulli(p)를 계산하여 d=1이면 표준 Self-Attention을 수행하고, d=0이면 이전 레이어 집합 {1, ..., l-1} 중 하나를 균등 분포로 선택하여 해당 레이어의 K_l', V_l'을 참조한다. [레이어 인덱스 입력 → 확률적 샘플링 연산 → 참조할 KV 소스 결정 → 레이어 간 캐시 공유 학습].

추론 시에는 고정된 캐시 공유 전략 S를 정의한다. S에 포함되지 않은 레이어는 가장 가까운 이전 캐시 레이어의 데이터를 재사용한다. [레이어 인덱스 l 입력 → S 포함 여부 확인 → 포함 시 신규 KV 저장, 미포함 시 이전 KV 로드 → 메모리 로드 횟수 감소].

Llama-3.1-8B 모델 실험 결과, HotpotQA 데이터셋에서 캐시를 25%만 유지(75% 삭제)했음에도 불구하고 R-CLA 모델은 F1 점수 0.237을 기록하여 기본 모델(0.080) 대비 약 196%의 성능 향상을 보였다.

Qwen3-8B 모델의 경우 100% 캐시 유지 설정에서 R-CLA 적용 시 F1 점수가 0.357로 기본 모델(0.233)보다 53% 이상 향상되는 정규화 효과가 관찰됐다.

추론 효율성 측면에서 Qwen3-8B 기준 32K 문맥 길이에서 4개 레이어마다 캐시를 공유할 경우, KV 캐시 메모리는 4.6GB에서 1.1GB로 4배 감소했으며 전체 피크 메모리는 약 3.5GB 절감됐다. 또한 초당 토큰 처리량은 22.8 tok/s에서 26.1 tok/s로 개선됐다.

R-CLA 아키텍처는 레이어 l의 쿼리 프로젝션 Q_l이 이전 레이어들의 다양한 KV 분포와 상호작용하도록 강제하여 견고성을 확보한다. 이는 수학적으로 레이어 간 피처 정렬(Feature Alignment)에 대한 불변성을 학습하는 과정이다.

구현 시 비리더(non-leader) 레이어는 K, V 프로젝션 연산을 완전히 생략하고 HBM(High Bandwidth Memory)으로부터의 로드 및 업데이트 과정을 건너뛴다. 이는 Algorithm 2에 명시된 것처럼 Load/Update 연산을 선택적으로 수행함으로써 메모리 대역폭 오버헤드를 직접적으로 줄인다.

실험 데이터셋으로는 HotpotQA, SQuAD v2, MSMarco, TriviaQA, RepLiQA 등 문맥 유지 능력이 중요한 QA 태스크들을 활용하여 캐시 공유 시의 정보 손실 여부를 엄격히 검증했다.

R-CLA를 활성화하기 위해서는 모델의 전체 파라미터 업데이트를 포함하는 추가적인 학습 자원이 필요하다. 또한 본 연구는 Mixture-of-Experts(MoE) 아키텍처에 대한 평가는 수행하지 않았으며, 주로 QA 태스크에 집중된 평가가 이루어졌다.