BUDDY: 적응형 대규모 언어 모델 추론을 위한 예산 기반 동적 깊이 라우팅

대규모 언어 모델 (LLMs)은 그 깊이와 파라미터 규모로 인해 높은 추론 비용이 발생합니다. 깊이 가지치기 (Depth pruning)는 중복되는 Transformer 블록을 건너뜀으로써 지연 시간 (latency)을 줄일 수 있지만, 기존 방법들은 (i) 사용자별 계산 예산 (compute budgets) 하에서 제어 능력이 제한적이며, (ii) 일반적으로 라우팅 경로를 고정하여 디코딩 (decoding) 과정 중 컨텍스트 (context)가 성장함에 따라 적응하지 못한다는 단점이 있습니다. 우리는 예산 기반의 동적 깊이 라우팅 프레임워크인 Buddy를 제안합니다. Buddy는 경량화된 결정 모듈 (Decision Module)을 사용하여 입력 조건에 따라 중간 레이어 (intermediate layers)의 점수를 매기고, 주어진 예산을 충족하기 위해 상위 k개의 레이어를 결정론적으로 실행합니다. 디코딩 시점의 적응을 지원하기 위해, Buddy는 오버헤드가 적은 전역 컨텍스트 (global context) 소스로서 첫 번째 레이어의 KV 캐시 (KV cache)를 재사용하며, 각 라우팅 결정 전에 이를 최신 토큰 표현 (token representation)과 함께 풀링 (pooling)합니다. 명시적인 예산이 제공되지 않을 경우, 선택 사항인 예산 예측기 (Budget Predictor)가 품질과 효율성의 균형을 맞추기 위해 입력에 따른 계산 수준을 추정합니다. Llama 제품군 및 Qwen 모델에 대한 실험 결과, Buddy는 강력한 정적 가지치기 (static pruning) 베이스라인들과 경쟁할 만한 성능을 보였으며, 엄격한 예산 제어, 디코딩 시점의 재라우팅 (rerouting), 그리고 단일 학습 모델 내에서의 다중 예산 지원을 독보적으로 수행하면서 정확도-계산 트레이드오프 (accuracy-compute trade-off)를 종종 개선함을 보여주었습니다.