Beyond KV Reconstruction: Functional Reconstruction for MLA Draft Models in Speculative Decoding

기술 요약: KV 재구성을 넘어 – Speculative Decoding를 위한 MLA Draft 모델의 기능적 재구성 (Functional Reconstruction)

문제 정의

Multi-head latent attention (MLA)는 메모리 트래픽을 줄이기 위해 점증하는 key-value (KV) 캐시를 압축된 잠재 상태(latent states)로 대체함으로써, 롱 컨텍스트(long-context) 대규모 언어 모델(LLM) 추론의 핵심 아키텍처로 부상했습니다. TransMLA, MHA2MLA와 같은 기존 방식들은 비용이 많이 드는 재학습 없이도 사전 학습된 Multi-Head Attention (MHA) 또는 Grouped-Query Attention (GQA) 모델을 MLA로 변환할 수 있게 해주지만, 이러한 구조적 변환은 변환된 모델이 **speculative decoding의 드래프트 모델(draft model)**로 사용될 때 실패하는 경우가 많습니다.

핵심 문제는 speculative decoding이 드래프트 모델의 토큰 제안과 타겟 모델의 검증 사이의 높은 일치도(agreement)에 의존한다는 점입니다. MHA/GQA에서 MLA로의 직접적인 변환은 드래프트 모델의 일치도를 저하시키는 두 가지 주요 오류 원인을 도입합니다:

1. 저차원 인수 분해 오류 (Low-rank factorization errors): KV 프로젝션을 더 낮은 차원의 잠재 공간으로 압축하는 과정에서, 원래의 rank가 잠재 rank보다 높을 경우 원래의 스택된 K/V 맵을 정확하게 재현할 수 없습니다.
2. RoPE 처리 오류 (RoPE handling errors): 변환 과정에서 Rotary Position Embeddings (RoPE)를 분해할 때, 유지된 위치 경로(positional paths)에 잔차(residuals)가 발생합니다.

이러한 오류들은 단독 생성 시에는 용인될 수 있을지 모르나, speculative decoding 중에는 토큰 수락률(acceptance rate)을 크게 떨어뜨려 드래프트 모델을 가속기가 아닌 계산 오버헤드로 만들어 버립니다. 기존의 변환 방식들은 캐시 효율성(structural validity)에 초점을 맞추고 있으며, 제안의 일치도(functional fidelity/proposal agreement)에는 초점을 맞추고 있지 않습니다.

방법론

본 논문은 변환된 캐시 구조를 변경하거나 verifier의 감독을 요구하지 않으면서 드래프트 품질을 최적화하도록 설계된 엔드 투 엔드(E2E) 방식이자 컨버터 불가지론적(converter-agnostic)인 사후 변환 절차인 **기능적 재구성(Functional Reconstruction)**을 제안합니다.

핵심 메커니즘

이 방법은 MLA 드래프트 구축을 캐시 압축 문제가 아닌 기능적 재구성 문제로 취급합니다. 이는 세 단계로 작동합니다:

1. 초기화 (Initialization): 표준 컨버터(TransMLA 또는 MHA2MLA)가 캐시 운용이 가능한 초기 MLA 체크포인트를 제공합니다.
2. 학습 목적 함수 (Training Objective): 컨버터에 의해 도입된 query 및 KV 프로젝션만을 학습시킵니다. 출력 프로젝션($W_O$)과 원래의 MHA/GQA 블록은 동결(frozen)됩니다.
   • 주어진 calibration hidden state $h_\ell$에 대해, 시스템은 동결된 원래 GQA 블록의 출력($A^{orig}_\ell(h_\ell)$)과 학습 가능한 변환된 MLA 블록의 출력($A^{MLA}_\ell(h_\ell)$)을 계산합니다.
   • 손실 함수는 두 블록의 post-output-projection 응답 사이의 마스크된 평균 제곱 오차(MSE)를 최소화합니다:
     $$\mathcal{L}^{(\ell)}_{E2E} = \mathbb{E}_b \left[ \frac{\| M_b \odot (A^{MLA}_\ell(H^b_\ell) - \text{sg}[A^{orig}_\ell(H^b_\ell)]) \|_F^2}{\sum_t M_{b,t}} \right]$$
     여기서 $M_b$는 유효한 토큰을 위한 마스크이며, $\text{sg}$는 stop-gradient를 의미합니다.
3. 범위 (Scope): 이 목적 함수는 layer input부터 post-output-projection output까지 전체 attention function을 다루며, 저차원 KV 변환 오류와 RoPE 처리 오류를 동시에 해결합니다.

주요 특징

• 컨버터 불가지론적 (Converter-Agnostic): 동일한 목적 함수가 서로 다른 컨버터(TransMLA, MHA2MLA)로부터 나온 체크포인트를 정교화하며, 각자의 특정 캐시 토폴로지를 보존합니다.
• Verifier 감독 불필요 (No Verifier Supervision): 타겟 모델의 logits, activations, 또는 벤치마크 라벨이 필요하지 않으며, 원래 모델의 calibration hidden state만을 사용합니다.
• 추론 시 부하 없음 (Inference-Free): 추론 중에 새로운 파라미터, 브랜치, 또는 캐시 상태를 도입하지 않습니다.

주요 기여

1. 목적의 구분 (Distinction of Objectives): 본 논문은 구조적 유효성(structural conversion, 즉 유효한 MLA 파라미터화 보장)과 드래프팅(drafting, 즉 제안의 일치도 보장)이 서로 다른 목적임을 규명합니다. 모델은 구조적으로 유효할 수는 있지만, 기능적으로는 드래프트로서 성능이 낮을 수 있습니다.
2. 학습 시 최적화 (Training-Time Optimization): 기능적 목적 함수를 사용하여 원래 모델의 attention 동작을 매칭함으로써, verifier의 감독 없이 학습 단계에서만 드래프트 품질을 최적화할 수 있음을 입증했습니다.
3. 일반적 적용성 (General Applicability): 제안된 방법은 다양한 컨버터(TransMLA, MHA2MLA), 모델 제품군(Llama, Qwen), 그리고 백엔드(Hugging Face, vLLM) 전반에서 성능을 개선하거나 유지하는 컨버터 불가지론적 특성을 보여줍니다.

실험 결과

저자들은 네 가지 드래프트-타겟 쌍(Llama-1B/3B $\to$ Llama-8B; Qwen-1.5B/3B $\to$ Qwen-7B), 두 개의 컨버터, 두 개의 백엔드, 그리고 네 가지 태스크(HumanEval, Alpaca, Natural Questions, CNN/DailyMail)에 걸쳐 192개의 구성을 평가했습니다.

• 수락률 (Acceptance Rates): "Partial RoPE" 베이스라인(컨버터 로컬 재구성 방법)과 비교했을 때, 기능적 재구성은 64개의 매칭된 태스크 셀 중 37개에서 실질적인 수락률 향상을 보였습니다 (0.5 퍼센트 포인트 허용 오차 기준).
  • 26개 셀은 실질적으로 변화가 없었습니다.
  • 1개 셀은 실질적인 감소를 보였습니다.
• 처리량 (Throughput): 수락률이 향상된 37개 셀 중 12개는 실질적인 처리량 증가도 보였습니다. 이 방법은 추론 오버헤드를 도입하지 않으므로, 처리량 증가는 높은 수락률에 의해 주도됩니다.
• 상호작용 효과 (Interaction Effects): 결과는 다음과 같은 "교차(crossover)" 패턴을 드러냈습니다:
  • Llama + TransMLA: 일관되게 개선됨.
  • Qwen + MHA2MLA: 일관되게 개선됨.
  • Llama + MHA2MLA: 대부분 변화 없음 (초기 변환이 이미 기능적으로 근접했음을 시사).
  • Qwen + TransMLA: 혼합된 결과, 일부 백엔드 민감도 존재.
• 회복 한계 (Recovery Limits): 이 방법은 변환 과정에서 손실된 수락률의 상당 부분을 회복하지만, 초기 변환이나 백엔드 구현이 심각한 불일치(예: 텐서 병렬성 세맨틱스)를 유발하는 경우, 원래의 변환되지 않은 GQA 드래프트 성능을 항상 완전히 복구하는 것은 아닙니다.

의의 및 주장

본 논문은 MLA 변환과 MLA 드래프트 구축이 서로 관련되어 있지만 별개의 작업으로 취급되어야 한다고 주장합니다.

• 계층적 성공 (Hierarchical Success): 저자들은 세 가지 수준의 성공을 정의합니다: 구조적 유효성(실행 가능한 파라미터), 기능적 충실도(attention map 매칭), 드래프트 유용성(speculative 수락). 현재의 컨버터들은 첫 번째는 만족하지만, 나머지 두 가지는 실패하는 경우가 많습니다.
• 기능적 인터페이스 (Functional Interface): 기능적 재구성이 갖는 의의는 모든 컨버터가 동일한 에러 분포를 남긴다고 가정하는 대신, 공유된 기능적 인터페이스(post-$W_O$ output)에서 최적화 목적을 설정한다는 점에 있습니다. 이는 잠재 rank와 캐시 예산이라는 고정된 제약 내에서 드래프트 모델을 정교화하는 "학습 시 조절 노브(training-time control knob)" 역할을 합니다.
• 실질적 시사점 (Practical Implication): 이 연구는 향후 MLA 파이프라인을 위한 세 단계의 워크플로우를 제시합니다: (1) 캐시 효율적인 파라미터 확립, (2) attention map 최적화를 위한 기능적 재구성 적용, (3) 매칭된 수락률을 통한 검증. 이는 기능적 재구성이 일반적인 정교화 방법이지만, 잠재 rank나 백엔드 구현(예: 텐서 병렬성 세맨틱스)의 근본적인 불일치를 보완할 수는 없음을 강조합니다.

결론적으로, 기능적 재구성이 많은 시나리오에서 수락률을 실질적으로 개선하지만, 완전한 speculative speedup을 달성하기 위해서는 여전히 고품질의 초기 컨버터와 최적화된 MLA 커널이 필요함을 시사합니다.