Trie 기반 빔 서치 - LLM 디코딩의 메모리와 속도를 동시에 잡다
목차
개요
이 논문은 대규모 언어 모델(LLM)에서 빔 서치(Beam Search) 디코딩의 메모리 비효율 문제를 해결하기 위한 새로운 접근법을 제안한다. 기존 배치 기반 빔 서치는 각 빔마다 별도의 KV 캐시를 유지해야 하므로 메모리 사용량이 빔 수에 비례하여 증가하는 문제가 있었다. 저자들은 공통 접두사를 가진 빔들 간에 단일 KV 캐시를 공유하는 Trie 기반 병렬 디코딩 방식을 제시하였다. 이 연구는 EMNLP 2025 메인 컨퍼런스에 채택되었다.
방법론
Trie 구조를 활용한 디코딩
Trie(접두사 트리)는 공통 접두사를 공유하는 문자열을 효율적으로 저장하는 자료 구조이다. 이 논문에서는 빔 서치 과정에서 생성되는 여러 후보 시퀀스들이 공통 접두사를 가진다는 점에 착안하여 Trie 구조를 디코딩에 적용하였다. 각 빔이 독립적으로 디코딩되는 기존 방식과 달리, Trie 구조를 통해 공통 접두사 부분의 연산과 메모리를 공유할 수 있다.
KV 캐시 공유 전략
기존 빔 서치에서는 각 빔마다 개별 KV 캐시를 유지해야 했다. 이 논문의 핵심 기여는 공통 접두사를 가진 빔들이 하나의 KV 캐시를 공유하도록 설계한 것이다. Trie 구조에서 같은 경로를 따르는 빔들은 동일한 KV 캐시를 참조하며, 분기가 발생하는 시점에서만 새로운 캐시가 생성된다. 이를 통해 중복 메모리 사용을 대폭 줄일 수 있다.
주의 메커니즘별 평가
이 논문에서는 세 가지 서로 다른 주의 메커니즘을 가진 모델에 대해 제안 방법을 평가하였다.
| 주의 메커니즘 | 평가 모델 |
|---|---|
| Multi-Head Attention | Phi-3.5-mini-instruct |
| Grouped Query Attention | Llama-3.1-8B-Instruct |
| Sliding Window Attention | Mistral-Small-24B-Instruct-2501 |
각 주의 메커니즘은 KV 캐시의 구조와 크기가 다르므로, Trie 기반 캐시 공유의 효과도 메커니즘에 따라 달라진다. 다양한 주의 메커니즘에서의 범용성을 검증하기 위해 세 가지 모델을 선정하여 실험을 수행하였다.
주요 결과
메모리 절감
Trie 기반 디코딩은 기존 배치 기반 빔 서치 대비 메모리 사용량을 4~8배 감소시켰다. 공통 접두사에 대한 KV 캐시 공유가 메모리 절감의 핵심 요인이다. 이러한 메모리 절감은 메모리 제약이 있는 환경에서 LLM 배포를 용이하게 만든다.
디코딩 속도 향상
메모리 절감과 함께 디코딩 속도도 최대 2.4배 향상되었다. KV 캐시의 중복 연산이 줄어들면서 전체 디코딩 과정이 가속화된 결과이다. CNN/DailyMail 요약 태스크와 HumanEval 코드 생성 벤치마크에서 평가하였으며, 생성 품질은 기존 방식과 동일하게 유지되었다.
| 평가 항목 | 결과 |
|---|---|
| 메모리 절감 | 4~8배 감소 |
| 디코딩 속도 | 최대 2.4배 향상 |
| 생성 품질 | 타협 없이 유지 |
| 평가 벤치마크 | CNN/DailyMail, HumanEval |
한계와 주의사항
이 논문은 특정 세 가지 모델과 두 가지 벤치마크에서만 실험을 수행하였다. 다른 모델이나 태스크에서의 일반화 가능성은 추가 검증이 필요하다. 또한 빔 간 공통 접두사의 비율이 낮은 경우 Trie 기반 캐시 공유의 이점이 줄어들 수 있다.
결론
이 논문은 Trie 기반 병렬 디코딩을 통해 LLM 빔 서치의 메모리와 속도 문제를 동시에 해결하는 방법을 제안하였다. 공통 접두사를 가진 빔들 간의 KV 캐시 공유라는 핵심 아이디어를 통해 메모리를 4~8배 절감하고, 디코딩 속도를 최대 2.4배 향상시켰다. 세 가지 서로 다른 주의 메커니즘에서 효과를 검증하였으며, 생성 품질의 타협 없이 효율성을 개선하였다. 메모리 제약 환경과 대규모 배포 상황에서 유용한 디코딩 방법론이다.