NVIDIA, DMS 기술로 트랜스포머 모델의 추론 성능 대폭 향상

NVIDIA와 에든버러대학교 연구진이 트랜스포머 대형 언어 모델(LLM)에서 8배의 키-값(KV) 캐시 압축을 가능하게 하는 동적 메모리 희소화(DMS)를 소개했습니다. 이 기술은 메모리 사용 효율성을 높이며, 추론 시간 성능을 크게 향상시키는 동시에 모델 정확도를 유지합니다. 대형 언어 모델의 주요 과제 중 하나는 긴 시퀀스나 병렬적인 추론 체인을 생성하는 것입니다. 그러나 이는 KV 캐시의 메모리 용량에 크게 의존하며, 시퀀스 길이와 폭(병렬 스레드)에 따라 선형적으로 증가하여 GPU 메모리를 많이 차지하고, 메모리 접근이 자주 발생하면서 추론 속도가 느려지는 문제가 있습니다. 기존의 KV 캐시 최적화 기법들은 훈련되지 않은 휴리스틱 방식(예: 어텐션 가중치 기반 토큰 제거)이나, 후 훈련이 필요한 복잡한 기법(예: 동적 메모리 압축)을 사용해 왔습니다. 전자는 정확도를 저하시키는 경향이 있으며, 후자는 계산 비용이 많이 들었습니다. DMS는 이러한 한계를 극복하기 위해 혼합 접근 방식을 취합니다. DMS는 전통적인 프루닝 방법처럼 KV 캐시를 희소화하지만, 약 1,000 단계의 최소 훈련 오버헤드와 지연된 제거를 통해 중요한 맥락 정보를 유지합니다. 이를 통해 갑작스러운 정확도 하락을 피할 수 있습니다. DMS의 핵심 아이디어는 훈련 중 토큰 제거 결정을 미분 가능한 방식으로 만드는 것입니다. 이를 위해 Gumbel-sigmoid 기반 샘플링 메커니즘을 사용합니다. 미래에 제거될 것으로 예측된 토큰들은 일정 기간 동안 계속 사용 가능하며, 그 후에야 finally 제거됩니다. 이 방식은 모델이 정보의 가치를 효과적으로 흡수할 수 있도록 합니다. DMS는 기존 모델을 레트로핏하는 데 적합합니다. DMS는 각 어텐션 헤드에 추가 매개변수를 도입하지 않고, 어텐션 메커니즘의 일부(단일 뉴런)를 재사용하여 토큰 제거를 예측합니다. 이는 복잡한 그래디언트 기반 최적화 없이도 수천 단계의 훈련을 필요로 하지 않으므로, 기존 모델에 쉽게 적용할 수 있습니다. 실험 결과, DMS는 1,000 단계의 훈련만으로도 KV 캐시를 8배 압축하며, 추론 작업에서 모델 성능을 유지하거나 개선하였습니다. AIME, GPQA, LiveCodeBench 등의 추론 중심 벤치마크에서 Qwen-R1 1.5B, 7B, 32B 등 다양한 크기의 모델에서 DMS는 정확도를 9.1점, 7.6점, 9.6점 향상시키며, 같은 메모리와 계산 예산 내에서 성능을 높였습니다. Quest와 TOVA 같은 최고 성능 기준 모델들과 비교했을 때, DMS는 KV 캐시 읽기 효율성(런타임 대체 지표)과 피크 메모리 사용량 모두에서 우수한 성능을 보여주었습니다. 이는 DMS가 더 나은 파레토 경계(Pareto frontier)를 달성함을 의미합니다. 비추론 중심 작업에서도 DMS는 성능을 유지하며, MMLU, GSM8K, HellaSwag 등의 짧은 맥락 벤치마크에서 4배 압축율을 적용해도 최대 3.5점의 성능 저하만을 보였습니다. Needle-in-a-Haystack와 Variable Tracking 같은 긴 맥락 작업에서는 DMS가 기존 모델을 능가하는 성능을 보여, 긴 시퀀스에서 발생하는 정보 과다 축소 문제를 완화할 가능성을 제시했습니다. 결론적으로, 동적 메모리 희소화(DMS)는 트랜스포머 기반 대형 언어 모델의 추론 시간 효율성을 실질적으로 향상시키는 해결책입니다. 최소한의 재훈련으로 KV 캐시를 효과적으로 압축하여, 긴 시퀀스나 병렬 추론을 수행하는 데 필요한 메모리와 런타임 요구사항을 늘리지 않으면서 성능을 유지하거나 개선시킵니다. DMS는 다양한 추론 및 일반적인 작업에서 일관된 성능 향상을 보여, 실제 추론 워크로드에서 압축, 정확도, 유연한 통합 사이의 균형을 잘 맞춘 기술로 평가받고 있습니다. 업계 전문가들은 DMS가 현재까지 제시된 KV 캐시 최적화 기법들 중 가장 실용적이고 확장 가능한 해결책이라고 평가합니다. NVIDIA는 그래픽 처리 분야에서 세계적으로 인지도가 높은 기업으로, 이번 연구를 통해 다시 한 번 AI 추론 효율성 향상에 기여하고 있습니다. DMS는 특히 리소스가 제약된 환경에서 LLM의 활용성을 크게 높일 것으로 기대됩니다.