기술 요약: LEDE – 강화 학습을 통한 LLM의 경험 기반 동적 종료(Experience-Driven Dynamic Exits)

1. 문제 정의

대규모 언어 모델(LLM)은 자기회귀적 추론(autoregressive inference) 과정에서 생성된 모든 토큰이 트랜스포머 스택의 전체 레이어를 통과해야 하기 때문에 상당한 지연 시간(latency) 문제를 겪습니다. **자기 스펙큘레이티브 디코딩(Self-Speculative Decoding, SSD)**은 모델 자신의 초기 레이어를 내부 "드래프터(drafter)"로 사용하여 여러 후보 토큰을 병렬로 생성함으로써 이 과정을 가속화하는 방안을 제안했지만, 현재의 SSD 방식은 근본적으로 **정적 구성(static configurations)**에 의해 제한됩니다.

기존 방식은 고정된 드래프트 깊이(draft depth)와 고정된 스펙큘레이션 길이(speculation length)를 전체 생성 과정에 일률적으로 적용합니다. 이러한 경직성은 문맥적 희소성(contextual sparsity), 즉 예측의 난이도가 생성되는 로컬 문맥에 따라 크게 달라지는 현상을 반영하지 못합니다. 결과적으로 정적 방식은 드래프트 품질(수락률)과 스펙큘레이션 속도 사이의 트레이드오프를 최적으로 조절하지 못하여, 하위 최적의(suboptimal) 추론 성능을 초래합니다.

2. 방법론: LEDE 프레임워크

저자들은 **LEDE (Learning-based Dynamic Exit)**를 제안하며, 이는 SSD의 동적 제어를 **마르코프 결정 과정(MDP)**으로 재정의합니다. 고정된 규칙을 사용하는 대신, LEDE는 **오프라인 강화 학습(Offline RL)**을 사용하여 생성 시퀀스의 로컬 문맥에 따라 최적의 종료 레이어와 스펙큘레이션 길이를 동적으로 선택하는 정책을 훈련합니다.

2.1. 마르코프 결정 과정(MDP) 정식화

이 프레임워크는 단일 생성 단계 중의 의사결정 과정을 다음과 같은 구성 요소를 가진 MDP로 모델링합니다:

• 상태 공간 (State Space, $S$): 에이전트는 각 후보 레이어 $l$에서 상태 벡터 $s_l$을 관찰합니다. 이 벡터는 모델의 내부 토큰 분포로부터 유도되며 다음을 포함합니다:
  • 토큰 신뢰도 (Token Confidence, $C_k^l$): 상위 $k$개 토큰의 평균 음의 로그 확률(negative log-probability).
  • 토큰 엔트로피 (Token Entropy, $H^l$): 다음 토큰 예측 분포의 불확실성.
  • 평균 토큰 신뢰도 추적 (Average Trace of Token Confidence, $C_{trace}$): 슬라이딩 윈도우 내 이전 단계들의 상위-1 토큰에 대한 평균 음의 로그 확률.
• 행동 공간 (Action Space, $A$): 각 후보 레이어에서 정책(Q-네트워크)은 두 가지 행동 중 하나를 결정합니다:
  • 계속 (Continue, $a=0$): 다음 후보 레이어로 진행.
  • 종료 (Exit, $a=1$): 현재 레이어를 드래프트 단계의 최적 종료 레이어($l^*$)로 지정.
• 보상 함수 (Reward Function, $R$): 추론 효율성을 최적화하도록 설계되었습니다.
  • 양의 보상 (Positive Reward): 드래프트 토큰이 수락되었을 때($x_t^{l^*} = x_t^L$) 주어지며, 절약된 깊이($1 + \frac{L-l^*}{L}$)에 비례하여 스케일링됩니다.
  • 음의 보상 (Negative Reward): 드래프트 토큰이 거부되었을 때($-2 + \frac{L-l^*}{L}$) 주어집니다.
  • 패널티 (Penalty): 계속 진행할 경우 작은 패널티($-0.01$)가 적용되어, 적절한 시점에 조기 종료를 하도록 유도합니다.

2.2. 오프라인 RL을 통한 정책 최적화

LEDE는 과거 추론 경험의 리플레이 버퍼를 사용하여 오프라인으로 훈련된 **심층 Q-네트워크(Deep Q-Network, DQN)**를 활용합니다:

1. 데이터 수집: 초기 탐색 단계에서 에이전트가 환경과 상호작용하며 다양한 경험 튜플 $(s_t, a_t, r_t, s_{t+1})$을 수집합니다.
2. 오프라인 훈련: Q-네트워크는 n-step 시간 차(Temporal Difference, TD) 오차를 최소화함으로써 훈련됩니다. 타겟 값은 즉각적인 종료 보상과 상태의 장기적 가치 사이의 균형을 맞추는 안정적인 타겟 네트워크를 사용하여 계산됩니다.
3. 이중 제어 메커니즘 (Dual-Control Mechanism): 온라인 추론 중에 훈련된 에이전트는 두 가지 적응형 제어를 수행합니다:
   • 동적 종료 레이어 선택: 드래프팅을 시작할 최적의 깊이 $l^*$를 결정합니다.
   • 적응형 드래프팅 길이: $l^*$를 선택한 후, 에이전트는 각 드래프트 토큰마다 상태를 재평가합니다. "종료(exit)" 행동은 드래프팅을 계속함(높은 신뢰도 의미)을 나타내고, "계속(continue)" 행동은 드래프팅을 즉시 중단하고 검증 단계로 넘어감(낮은 신뢰도 의미)을 나타냅니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 설명합니다:

1. MDP 정식화: SSD의 드래프트 깊이 및 스펙큘레이션 길이의 동적 제어를 MDP 문제로 정식화하고, 제어 정책을 도출하기 위해 강화 학습을 적용한 첫 번째 연구입니다.
2. LEDE 프레임워크: 정적 구성 및 단순 휴리스틱 기반 방식의 한계를 극복하기 위해, 오프라인 RL과 경험 재생(experience replay)을 활용하여 문맥 인지형 정책을 학습하는 프레임워크를 설계 및 구현했습니다.
3. 실증적 검증: 지시 이행(instruction following), 요약, 언어 모델링, 코드 생성 등 다양한 태스크에 걸쳐 표준 자기회귀 생성 및 정적 베이스라인 대비 유의미한 속도 향상을 입증하는 포괄적인 실험을 수행했습니다.

4. 실험 결과

프레임워크는 Llama-2 (7B, 13B) 및 Llama-3 (1B), 그리고 CodeLLaMA (7B, 34B) 모델을 대상으로 지시 이행, 요약, 언어 모델링, 코드 생성 등의 다양한 태스크에서 평가되었습니다.

• 속도 향상 성능:
  • LEDE는 표준 자기회귀 디코딩 대비 평균 2.32배의 속도 향상을 달이브했습니다.
  • 피크(Peak) 속도 향상은 2.72배에 달합니다 (예: 언어 모델링 작업에서의 Llama-2-7B).
  • LEDE는 가장 우수한 정적 스펙큘레이티브 베이스라인인 LayerSkip보다 추가로 17%의 속도 향상을 제공합니다.
• 베이스라인과의 비교:
  • LayerSkip (정적): 고정된 구성에 매우 민적이며 태스크 간 일반화 능력이 떨어집니다.
  • 규칙 기반 방법 (LITE, Draft & Verify): 미리 정의된 임계값에 의해 제약되어, 깊이와 수락률 사이의 최적의 트레이드오프를 찾는 데 한계가 있습니다.
  • LEDE: 미세한 제어 정책을 학습함으로써, 생성 품질(Rouge-L로 측정)을 유사하게 유지하면서도 속도 측면에서 모든 베이스라인을 일관되게 능가합니다.
• 절제 연구 (Ablation Study): 동적 종료 레이어 선택 또는 적응형 드래프팅 길이 중 하나라도 제거하면 성능이 크게 저하되어(속도 향상이 약 2.0배로 감소), 두 메커니즘이 상호 보완적이며 필수적임을 확인했습니다.

5. 의의 및 주장

본 논문은 LEDE가 경직된 휴리스틱 중심의 추론에서 **적응형, 정책 기반 추론(adaptive, policy-driven inference)**으로의 개념적 전환을 나타낸다고 주장합니다. 실시간 문맥에 따라 계산 깊이와 스펙큘레이션 길이를 공동 최적화함으로써, 이 프레임워크는 다양한 모델과 태스크 전반에 걸쳐 더욱 견고하고 효율적인 디코딩을 가능하게 합니다.

저자들은 이 연구를 더 지능적이고 계산 효율적인 LLM 시스템을 구축하기 위한 단계로 포지셔닝합니다. 향후 연구는 이러한 이점을 더 큰 모델(70B 규모 이상)로 확장하는 데 집중할 것이라고 언급했으나, 현재의 기여만으로도 대규모 언어 모델의 추론 파이프라인을 동적으로 관리하기 위해 오프라인 강화 학습을 사용하는 것의 타당성을 입증하였습니다.