기술 요약: Dr.LLM – LLM 의 동적 레이어 라우팅

문제 제기

대형 언어 모델 (LLM) 은 현재 정적 깊이 (static-depth) 체제 하에서 작동하며, 입력의 복잡성과 관계없이 모든 토큰을 고정된 트랜스포머 레이어 스택을 통해 처리합니다. 이 접근 방식은 두 가지 주요 비효율성을 초래합니다:

1. 낭비되는 계산: 간단한 쿼리가 모든 레이어를 통과함으로써 불필요한 계산 자원을 소모합니다.
2. 불충분한 유연성: 복잡한 추론 작업은 종종 더 깊은 처리나 반복적 정제가 필요하지만, 정적 아키텍처는 재학습이나 아키텍처 개편 없이는 이를 제공할 수 없습니다.

기존의 적응형 깊이 방법들 (예: 조기 종료, 레이어 가지치기, 혼합 깊이) 은 다음과 같은 중대한 한계에 직면해 있습니다: 종종 정확도를 희생하여 속도를 높이거나, 광범위한 재학습이나 아키텍처 수정을 요구하거나, 확장성을 저해하는 비용이 많이 드는 추론 시 검색에 의존합니다. 또한 많은 이전 접근 방식들은 효율성 향상에도 불구하고 모델 정확도를 저하시킵니다.

방법론: Dr.LLM

저자들은 Dr.LLM(LLM 을 위한 레이어의 동적 라우팅) 을 제안합니다. 이는 동결된 사전 학습 LLM에 경량화된 레이어별 라우터를 장착하는 재설계 가능한 프레임워크입니다. 이러한 라우터는 주어진 입력에 대해 특정 트랜스포머 블록을 생략 (skip), 실행 (execute), 또는 반복 (repeat) 할지 결정합니다.

1. 라우터 아키텍처

• 구조: 각 트랜스포머 블록 $B_\ell$은 선형-GELU-선형 레이어로 구성된 경량 다층 퍼셉트론 (MLP) 과 짝을 이룹니다.
• 입력: 라우터는 이전 레이어의 은닉 상태에 대한 간결한 요약에 작동합니다. 긴 컨텍스트에서의 안정성을 보장하기 위해 저자들은 윈도우 평균 풀링 (windowed mean pooling) 을 사용하여 시퀀스를 연속된 윈도우로 나누고 각 윈도우 내의 은닉 상태를 평균화합니다.
• 출력: 라우터는 $\{skip, execute, repeat\}$ 세 가지 클래스에 대한 로짓을 출력합니다.
• 오버헤드: 라우터는 입력 시퀀스당 한 번 실행됩니다. 오버헤드는 무시할 수 있으며 (토큰 수에 대해 일정), KV 캐싱과 완전히 호환됩니다.

2. MCTS 를 통한 지도 학습

강화 학습이나 추론 시 검색에 의존하는 방법들과 달리, Dr.LLM 은 오프라인 몬테카를로 트리 검색 (MCTS) 에서 도출된 명시적 지도 (explicit supervision) 를 사용하여 라우터를 학습합니다.

• 검색 과정: 각 입력에 대해 MCTS 는 계산 예산 하에서 수정된 실행 경로 (레이어 생략 또는 반복) 를 탐색합니다. 기본 전체 깊이 실행과 비교하여 작업 정확도를 유지하거나 향상시키는 경로만 유지합니다.
• 데이터셋: 검색은 약 4,000 개의 최적화된 실행 경로 (질문, 최적 레이어 동작, 답변의 튜플) 로 구성된 간결한 데이터셋을 생성합니다.
• 학습 목적: 라우터는 지배적인 "실행" 클래스와 소수인 "생략" 및 "반복" 클래스 간의 심각한 불균형을 처리하기 위해 클래스 균형 가중치를 적용한 Focal Loss를 사용하여 이 데이터셋에서 엔드 투 엔드로 학습됩니다.
• 동결: 핵심적으로 기본 LLM 가중치는 동결된 채로 유지되며, 라우터 파라미터만 업데이트됩니다.

3. 추론

추론 시, 훈련된 라우터는 검색 없이 각 레이어에 대해 탐욕적 결정 ($\hat{y}_\ell = \arg\max p_\ell$) 을 내립니다.

• 생략 (Skip): 레이어를 우회합니다 ($H^{(\ell)} = H^{(\ell-1)}$).
• 실행 (Execute): 레이어를 한 번 적용합니다.
• 반복 (Repeat): 레이어를 연속하여 두 번 적용합니다.
• 배치 처리: 배치 처리에서의 동기화를 처리하기 위해 시퀀스는 각 레이어에서의 라우팅 결정에 따라 그룹화됩니다. 저자들은 높은 라우팅 합의 (시퀀스 간 80–89% 일치) 가 동기화 오버헤드를 최소화한다고 지적합니다.

주요 기여

1. 동결된 LLM 을 위한 지도식 동적 라우팅: 레이어별 라우터를 도입하여 생략, 실행, 반복을 결정하며, 기본 가중치를 수정하지 않고 4k 개의 MCTS 기반 경로만으로 엔드 투 엔드 학습합니다.
2. MCTS 를 통한 효과적인 경로 지도: 정확도 유지 또는 향상을 위한 레이어 구성을 오프라인에서 발견하는 길이 인식 MCTS 알고리즘을 통해 추론 시 검색 없이 간결한 지도 데이터셋을 생성합니다.
3. 경량 라우터 학습: 긴 컨텍스트 안정성을 위한 윈도우 평균 풀링, 클래스 불균형을 위한 포커스 손실, 그리고 간결한 2 층 선형 라우터를 특징으로 하는 설계로, 계산 오버헤드를 무시할 수 있도록 보장합니다.
4. 정확도 및 효율성 향상: 추론 정확도 향상과 함께 실행된 레이어 수 감소를 입증하여 기존 최첨단 라우팅 방법들을 능가합니다.
5. 강건한 일반화: 특정 추론 작업에서 학습된 라우팅 정책이 최소한의 정확도 저하로 도메인 외 벤치마크에 효과적으로 전이됨을 보여줍니다.

실험 결과

도메인 내 성능 (ARC 및 DART)

LLaMA-3.2 및 Qwen-2.5 변형 등 6 개의 모델에서 평가한 결과, Dr.LLM 은 실행된 레이어의 평균 수를 줄이면서 일관되게 정확도를 향상시켰습니다:

• 정확도: DART(수학) 에서 최대 +3.4 퍼센트 포인트 (p.p.), ARC(논리) 에서 **+1.0 p.p.**만큼 향상되었습니다. 예를 들어, LLaMA-3B-Base 는 DART 에서 11.8% 에서 15.8% 로 급증했습니다.
• 효율성: 모델 간 예제당 평균 5.0 개의 레이어를 절약했으며, 특정 사례에서는 최대 11.0 개의 레이어를 절약했습니다.
• 지연 시간: 1,000 토큰 실행에서 15.3% 의 월시계 생성 시간 단축을 달성했으며, 라우터 오버헤드는 전체 지연 시간의 1% 미만을 차지했습니다.

도메인 외 일반화

재학습 없이 다양한 벤치마크 (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, AGIEval, PIQA) 에 적용했을 때:

• 라우터는 8 개 벤치마크에서 평균 0.85 p.p.의 정확도 하락만 겪으면서 효율성을 유지했습니다.
• 일부 경우 (예: LLaMA-3B 를 사용한 GPQA Diamond) 에는 정확도가 실제로 향상되었습니다 (+2.5 p.p.).
• 이는 학습된 정책이 특정 작업 패턴에 과적합되는 것이 아니라 트랜스포머 계산 내의 일반적인 구조적 중복성을 포착함을 시사합니다.

최첨단 방법과의 비교

Dr.LLM 은 추론 벤치마크에서 정확도 측면에서 기존 라우팅 방법들 (예: FlexiDepth, MindSkip, ShortGPT) 보다 최대 **+7.7 p.p.**만큼 더 나은 성능을 발휘했습니다. 수만 개의 예제에서 학습이 필요하고 종종 정확도를 희생하는 FlexiDepth 와 달리, Dr.LLM 은 훨씬 낮은 학습 오버헤드 (4k 개의 예제, 단일 GPU) 와 아키텍처 변경 없이 더 높은 정확도를 달성합니다.

라우팅 패턴 분석

라우팅 결정에 대한 분석은 구조화된 행동을 드러냈습니다:

• 초기 레이어: 일관되게 실행됨 (안정적인 사용).
• 중간 레이어: 빈번히 생략됨 (계산 절약).
• 후기 레이어: 특히 복잡한 추론 작업에서 종종 반복됨, 이는 반복적 정제 역할을 나타냄.
• 특징 학습: 이전 레이어 은닉 상태로 훈련된 라우터는 정적 입력 임베딩으로 훈련된 라우터보다 현저히 우수한 성능을 보였으며, 이는 Dr.LLM 이 내부 모델 상태를 동적으로 라우팅 결정에 매핑하는 법을 학습함을 확인시켜 줍니다.

중요성 및 주장

이 논문은 명시적으로 지도된 라우터가 동결된 LLM 을 재설계하여 예산 인식 및 정확도 주도 추론을 달성할 수 있음을 Dr.LLM 이 입증한다고 주장합니다. 주요 교훈은 다음과 같습니다:

• 재학습 불필요: 기본 모델 가중치는 변경되지 않아 모델의 원래 능력을 유지합니다.
• 정확도 긍정적 효율성: 속도를 위해 정확도를 희생하는 많은 적응형 방법과 달리, Dr.LLM 은 추론 중심 작업에서 정확도와 효율성을 동시에 향상시킵니다.
• 실용적 배포: 추론 시 검색과 아키텍처 수정의 필요성을 제거함으로써 Dr.LLM 은 적응형 추론을 위한 실용적이고 확장 가능한 솔루션을 제공합니다.
• 일반화 가능성: 학습된 라우팅 정책은 모델 패밀리 및 작업 도메인 간에 강건하게 전이되며, 이는 적응형 계산 할당이 트랜스포머 아키텍처의 근본적이고 학습 가능한 특성임을 시사합니다.

저자들은 Dr.LLM 이 재학습이나 아키텍처 재설계와 관련된 금지적인 비용 없이 효율성, 정확성, 강건성을 조화시키는 확장 가능한 적응형 추론을 향한 실용적인 단계라고 결론지었습니다.