엔비디아가 AI 에이전트가 여러 모델을 번갈아 활용하는 장시간 작업 중 모델 전환 시 발생하는 막대한 계산량과 처리 지연 문제를 완화할 수 있는 새로운 기술을 발표했다.
엔비디아는 최근 한 AI 모델의 KV(Key-Value) 캐시를 다른 모델이 바로 활용하도록 변환해 주는 '크로스 모델 KV 캐시 전송(Cross-Model KV Cache Transfer)' 기술을 온라인 아카이브를 통해 공개했다.
기존 방식에서는 AI가 대화 중 더 강력한 모델이나 비용 효율적인 모델로 전환할 때마다 새 모델이 축적된 대화 내용을 처음부터 다시 계산해야 했지만, 이 기술은 기존 모델의 KV 캐시를 새 모델에서 곧바로 사용할 수 있는 형태로 변환한다.
KV 캐시와 모델 전환의 과제
대형언어모델(LLM)은 사용자의 입력을 처리할 때 먼저 전체 맥락을 한 번 처리하는 '프리필(prefill)' 단계를 거쳐 키와 값 정보를 KV 캐시에 저장한다. 이후 토큰을 생성하는 과정에서는 이 저장된 정보를 재활용하므로, 긴 대화를 매번 처음부터 다시 처리할 필요가 없다.
문제는 AI 에이전트가 작업 특성에 따라 여러 모델을 오갈 때 나타난다. 서로 다른 모델은 KV 캐시를 저장하고 해석하는 방식이 다르기 때문에 기존 캐시를 그대로 쓸 수 없고, 새 모델이 지금까지의 모든 맥락을 다시 프리필해야 한다. 특히 맥락이 길어질수록 모델 크기와 입력 길이에 따라 계산량과 지연시간이 급격히 증가한다.
선형 변환을 통한 해결책
엔비디아 연구진은 이 문제를 해결하기 위해 한 AI 모델의 KV 캐시를 다른 모델이 사용할 수 있게 변환하는 방법을 개발했다. 핵심 아이디어는 같은 계열의 모델들이 크기가 달라도 KV 캐시가 유사한 패턴(선형적 관계)을 따른다는 발견에 기반한다. 이를 통해 복잡한 재학습 없이 간단한 수학적 연산(선형 연산)만으로 캐시를 다른 모델에 맞게 변환할 수 있다는 가능성을 제시했다.
'큐원3 14B'와 '32B' 모델을 대상으로 한 실험에서 14B 모델의 한 레이어 정보만으로도 32B 모델의 키(Key) 정보 변화량의 56%, 값(Value) 정보 변화량의 32%를 설명할 수 있었다. 여러 레이어의 정보를 활용하자 이 수치가 키는 79%, 값은 65%까지 상승했다. 즉, 작은 모델이 이미 처리한 대화 맥락을 더 큰 모델이 처음부터 다시 읽지 않아도 활용할 수 있다는 가능성을 확인한 것이다.
선형 매퍼를 이용한 구현 방식
연구진은 헤드별 릿지 회귀(ridge regression)라는 상대적으로 단순한 선형 계산 기법(선형 매퍼)을 활용했다.
먼저 대상 모델의 각 레이어마다 소스 모델의 어떤 레이어가 가장 도움이 될지 선별한다. 그 다음 위치 정보를 담은 RoPE를 제거하고, 적은 양의 데이터를 이용해 소스 모델의 KV 캐시를 대상 모델에 맞게 변환하는 수식을 만든다.
이 과정에 사용된 조정 데이터는 1024토큰 길이의 '파인웹-Edu(FineWeb-Edu)' 데이터 500개에 불과했다.
실험 결과와 성능 개선
큐원3, 라마 3.1, 미스트랄 3 등 3개 모델 계열의 6개 조합에 대해 시험한 결과, 선형 매퍼로 KV 캐시를 변환했을 때 일부 조합에서는 대상 모델이 맥락을 처음부터 직접 처리했을 때의 정확도를 73~98% 수준으로 유지할 수 있었다.
속도 측면에서도 눈에 띄는 개선이 있었다. 긴 대화를 처음부터 다시 처리하는 기존 방식과 비교하면, KV 캐시를 변환해 전달하는 방식이 2.7배에서 최대 25배까지 빨랐다. 모델 전환 시 발생하던 재계산 부담을 크게 줄이면서도 기존 모델의 성능을 상당 부분 유지할 수 있다는 뜻이다.
현재의 한계와 향후 개선 방향
다만 현재로서는 모든 AI 모델 간에 자유롭게 KV 캐시를 주고받기는 아직 어렵다. 이번 연구는 KV 헤드 수와 헤드당 차원이 동일한, 즉 같은 모델 계열 내에서 호환되는 모델들을 중심으로 진행됐다. 따라서 서로 다른 모델 계열이나 구조가 크게 다른 모델들 사이에서도 동일한 방식이 작동하는지는 추가 연구가 필요한 상황이다.
일부 미스트랄 3 모델 조합에서는 단순한 선형 방식만으로는 정확도가 상당히 떨어지는 제약도 드러났다. 이에 연구진은 선형 매퍼 대신 비선형 MLP(다층 퍼셉트론)를 도입했고, 성능이 저하된 사례에서 정확도를 최대 37%포인트까지 향상시킬 수 있었다.
이는 선형 방식으로는 풀기 어려운 모델 조합에 대해서는 더욱 정교한 변환 기법이 필요할 수 있음을 시사한다.