벡터 인덱스로 LLM 출력층과 디코딩을 가속하는 방법

배경

대규모 언어 모델은 토큰을 하나 생성할 때마다 현재 은닉 상태를 전체 어휘의 로짓으로 변환한다. 일반적인 구현에서는 출력 임베딩 행렬과 조밀한 행렬 곱을 수행한 뒤 Top-k 선택이나 샘플링을 적용한다. 어휘가 큰 다국어 모델에서는 이 행렬을 반복해서 읽는 과정이 계산량 못지않은 부담이 된다. 특히 CPU에서 소형 모델을 배치 크기 1로 실행할 때는 연산 성능보다 메모리 대역폭과 데이터 이동이 지연을 결정할 수 있다.

출력 투영을 검색 문제로 전환

이 연구의 핵심은 출력 투영과 Top-k 선택을 최대 내적 검색으로 다시 정의하는 것이다. 은닉 상태와 모든 토큰 임베딩의 내적을 계산하는 대신, 토큰 임베딩을 벡터 데이터베이스처럼 취급하고 HNSW（Hierarchical Navigable Small World）인덱스에서 내적 점수가 높은 후보를 근사 검색한다.

제안 방식은 다음과 같이 구성된다.

• 토큰 임베딩 인덱싱: 출력 임베딩을 HNSW 구조에 넣어 근사 벡터 검색을 지원한다.
• 후보 집합 축소: 매 디코딩 단계에서 점수가 높을 가능성이 큰 소수의 토큰만 가져와 전체 어휘에 대한 접근을 줄인다.
• 기존 디코딩 파이프라인과 연결: 검색된 후보의 로짓을 희소한 전체 어휘 텐서에 분산시켜 기존 샘플링 및 Top-k 인터페이스와 함께 사용할 수 있게 한다.
• 근사 검색으로 속도 확보: 전역 순위를 완전히 보장하기보다 실제 생성에 필요한 유력 후보를 유지하는 대신 메모리 접근을 줄인다.

실험 결과와 한계

저자들은 Gemma 3, Llama 3.2, Qwen 3 모델을 CPU에서 평가했다. 벡터 인덱스 기반 출력 헤드는 출력 투영 시간을 크게 줄였다. 엔드투엔드 배치 크기 1 디코딩에서는 Gemma 3 270M의 처리량이 최대 82% 향상된 것으로 보고됐다. 또한 AlpacaEval 평가에서 생성 품질이 유지됐다고 설명한다.

다만 이 결과를 모든 환경에서 조밀한 출력층이 필요 없다는 의미로 확대해서는 안 된다. 이 방식은 어휘가 크고 배치가 작으며 출력 투영이 전체 지연 시간에서 큰 비중을 차지하는 경우에 특히 적합하다. 배치가 커지면 조밀한 행렬 연산이 하드웨어 병렬성을 더 잘 활용할 수도 있다. 근사 검색에는 인덱스 구축, 메모리 배치, 후보 재현율, 희소 로짓을 되돌려 놓는 비용도 따른다. 제공된 자료만으로는 인덱스 크기와 검색 정확도, 지연 시간 사이의 전체적인 비교를 확인할 수 없으므로 실제 서비스에는 모델과 하드웨어별 검증이 필요하다.

의미와 영향

이번 연구는 LLM 출력 헤드를 고정된 조밀 행렬 연산으로만 보지 않고, 토큰 벡터에 대한 검색 모듈로 설계할 수 있음을 보여준다. 앞으로 인덱스 오버헤드와 품질·지연 시간의 균형이 더 폭넓게 검증된다면 CPU 추론, 엣지 배포, 낮은 동시성의 서비스에서 활용할 수 있는 선택지가 될 수 있다.

핵심 시사점은 메모리 접근이 디코딩의 병목일 때 전체 어휘 행렬 연산을 계속 최적화하는 것보다, 실제로 읽어야 하는 어휘 범위를 줄이는 것이 더 직접적인 해법일 수 있다는 점이다. 따라서 이 접근은 기존 출력층의 보편적 대체재라기보다 지연 시간에 민감한 소배치 추론을 위한 시스템 최적화로 이해하는 편이 적절하다.

출처: arXiv