최근 대규모 언어모델(LLM)의 장문 컨텍스트 처리 능력이 중요해지면서, 추론(inference) 효율성 개선에 대한 연구가 활발합니다. 특히 키-값(KV) 캐시의 크기 증가 문제를 해결하고 디코딩 메모리 트래픽을 줄이는 데 효과적인 다중 헤드 잠재 어텐션(Multi-head latent attention, MLA)이 주목받고 있습니다. 하지만 대부분의 강력한 오픈소스 LLM은 다중 헤드 어텐션(MHA)이나 그룹 쿼리 어텐션(GQA) 구조를 사용하고 있어, MLA의 효율성을 얻으려면 모델을 처음부터 다시 학습시키지 않고 변환하는 과정이 필요합니다.
문제는 기존 MHA/GQA 모델을 MLA로 직접 변환할 경우, 추측 디코딩(speculative decoding) 과정에서 드래프트 토큰(draft token)의 수용률(acceptance rate)이 크게 떨어진다는 점입니다. 추측 디코딩은 작은 드래프트 모델이 다음 토큰들을 미리 예측하고, 더 큰 타겟 모델이 이를 검증하여 추론 속도를 높이는 기술입니다. 이때 드래프트 모델의 예측이 정확해야 전체적인 속도 향상 효과를 볼 수 있는데, MLA 변환 과정에서 발생하는 어텐션 기능 오류가 드래프트 토큰 수용률을 저하시키는 원인이었습니다. 기존의 KV 재구성 방식은 캐시 압축에 초점을 맞췄지만, 이는 어텐션 기능 자체의 정확도를 떨어뜨려 추측 디코딩의 효율을 해쳤습니다.
이러한 문제를 해결하기 위해 연구진은 '기능적 재구성(functional reconstruction)'이라는 새로운 MLA 드래프트 모델 구축 방법을 제안했습니다. 이 종단간(end-to-end) 방식은 변환된 MLA 어텐션 모듈이 원본 MHA/GQA 모듈의 출력 투영(output projection) 후 응답을 최대한 재현하도록 최적화합니다. 이는 검증자 로짓(verifier logits)이나 검증자 감독(verifier supervision) 없이도 변환된 캐시와 추론 그래프를 보존하며, 캘리브레이션 은닉 상태(calibration hidden states)를 활용합니다. Llama 및 Qwen 모델을 포함한 다양한 환경에서 192가지 모델-변환기-백엔드-방법-작업 구성을 평가한 결과, 기능적 재구성은 대부분의 경우에서 드래프트 토큰 수용률을 실질적으로 개선하는 것으로 나타났습니다. 이는 MLA의 캐시 효율성과 추측 디코딩의 속도 향상을 동시에 달성하여, 장문 컨텍스트 LLM 추론의 새로운 지평을 열 것으로 기대됩니다.