기술 요약: 깊이 탐색 디코딩 (Depth Exploration Decoding, DEX)

1. 문제 정의

대규모 언로 모델(LLM)의 자기회귀(Autoregressive, AR) 디코딩은 다음 토큰을 생성하기 위해 모든 토큰이 전체 레이어 스택을 통과해야 하므로 본질적으로 비용이 많이 듭니다. 최근 연구들은 깊이 차원(depth dimension)에서 상당한 중복성을 확인했습니다. 즉, 많은 토큰이 중간 레이어 단계에서 이미 예측 가능해지며, 이는 정확도를 위해 반드시 전체 모델 깊이가 필요하지는 않음을 의미합니다.

기존의 무손실 깊이 적응형 방법들은 이를 활용하기 위해 깊이 선택(depth-selection) 전략을 채택합니다. 이러한 방법들은 단일하게 사전 선택된 중간 깊이에서 토큰을 초안(draft)하고, 이를 최종 깊이 모델과 대조하여 검증합니다. 그러나 저자들은 다음과 같은 결정적인 **선택 병목 현상(selection bottleneck)**을 지적합니다:

• 늦은 종료 (Late Exit): 선택된 종료 깊이가 너무 깊으면, 이미 중복된 레이어에 계산 자원이 낭비됩니다.
• 빠른 종료 (Early Exit): 선택된 종료 지점이 너무 얕아서 예측이 최종 깊이 출력과 일치하지 않을 경우, 시스템은 폴백(fallback)을 트리거합니다. 이는 현재 토큰을 폐기할 뿐만 아니라, 의존적인 "자손(descendant)" 초안들까지 무효화하여 상당한 효율성 손실을 초래합니다.

본 논문은 현재의 선택 기반 방법들과 이상적인 "오라클(oracle)"(토큰이 안정화되는 정확한 시점을 알고 있는 존재) 사이의 간극이 매우 크다고 주장합니다. 핵심 질문은, 잘못된 종료 깊이 결정의 위험을 줄임으로써 이 중복성을 더 효과적으로 활용할 수 있는가 하는 점입니다.

2. 방법론: 깊이 탐색 디코딩 (DEX)

저자들은 단일 깊이 선택을 대신하여 여러 후보 깊이에 대해 **병렬 탐색(parallel exploration)**을 수행하는 무손실 알고리즘인 DEX를 제안합니다.

핵심 개념

• 최초 가용 깊이 (Earliest Available Depth, EAD): 저자들은 EAD를 토큰의 예측이 최종 깊이 출력과 일치하고 이후 레이어에서도 안정적으로 유지되는 가장 얕은 레이어로 정의합니다.
• 탐색 세트 (Exploration Set, $X$): DEX는 단일 출구를 선택하는 대신, $\{d_0, d_1, \dots, d_{N-1}\}$ 형태의 후보 깊이 세트를 정의하며, 여기서 $d_{N-1}$은 최종 깊이입니다.
• 해상도 (Resolution): DEX의 오버헤드는 이 세트의 "해상도"(임의의 레이어와 다음 탐색 깊이 사이의 최대 간격)에 의해 제한됩니다. 세트가 더 정밀해질수록 오버헤드는 0에 수렴합니다.

알고리즘: 확장–확정–붕괴 (Expand–Commit–Collapse)

DEX는 유한 상태 머신(FSM)에 의해 관리되는 깊이-위치 격자(depth–position lattice) 상에서 작동합니다. 프로세스는 세 단계로 구성됩니다:

1. 격자 확장 (Lattice Expansion, 병렬 탐색):

   • 모델은 $N$개의 깊이 단계로 분할되며, 각 단계는 별도의 "깊이 탐색기(depth explorer)"(잠재적으로 서로 다른 하드웨어 유닛)에 의해 처리됩니다.
   • 각 확정 위치(commit position)에서 시스템은 후보 브랜치들을 병렬로 확장합니다. 탐색 토큰은 깊이 확장(현재의 은닉 상태를 다음 레이어 탐색기로 전달) 또는 위치 확장(현재의 후보를 다음 토큰 위치 생성을 위한 접두사로 사용)을 수행할 수 있습니다.
   • 이는 서로 다른 브랜치가 서로 다른 깊이에서 만들어진 예측을 나타내는 잠재적 미래의 트리를 생성합니다.

2. 확정 (Commit, 검증):

   • "앵커 경로(anchor path)"(모든 레이어를 통과하는 브랜치)가 최종 깊이에 도달하면, 최종 깊이 참조 토큰을 생성합니다.
   • DEX는 앵커 경로로부터 나온 모든 비최종 후보 토큰들을 이 참조 토큰과 대조하여 검증합니다.
   • 최종 깊이 참조와 일치하는 가장 얕은 브랜치를 선택합니다. 만약 중간 일치 항목이 없다면, 최종 깊이 브랜치로 폴백합니다.
   • 결정적으로, 출력에 추가되는 토큰은 항상 최종 깊이 참조에 의해 생성된 것이므로, 표준 AR 디코딩과 동일한 무손실 동등성을 보장합니다.

3. 붕괴 (Collapse, 상태 프루닝):

   • 토큰을 확정(commit)한 후, 시스템은 격자를 붕괴시킵니다. 새로운 확정된 접두사와 일치하지 않는 브랜치들은 폐기됩니다.
   • 선택된 브랜치의 상태(KV 캐시, 은닉 상태)만이 다음 단계를 위해 보존됩니다.
   • 이를 통해 시스템이 여러 깊이를 탐색하면서도, 다음 AR 단계를 위해 필요한 정확한 상태를 유지하도록 합니다.

최적화

• 깊이 결합 샘플링 (Depth-Coupled Sampling, DECS): 동일한 슬롯 체인 내의 서로 다른 깊이들이 동일한 토큰을 제안하여 탐색을 낭비하는 것을 방ка하기 위해, DEX는 더 깊은 후보를 디코딩할 때 더 얕은 깊이에서 이미 제안된 토큰들을 마스킹합니다. 이는 최종 샘플링 분포를 변경하지 않으면서 중간 제안의 다양성을 장려합니다.
• 유도 어댑터 (Inducing Adapter): 중간 레이어가 자연스럽게 출력 헤드와 정렬되지 않은 표준 LLM의 경우, 저자들은 중간 깊이에 작은 잔차 어댑터(residual adapters)를 부착할 것을 제안합니다. 이 어댑터들은 자기 증류(self-distillation)를 통해 학습되어, 중간 은닉 상태를 최종 LM 헤드와 정렬함으로써 더 얕은 깊이에서의 토큰 "준비도"를 향상시킵니다.

3. 주요 기여

1. 토큰 준비도(Token Readiness)의 공식화: 본 논문은 토큰이 언제 예측 가능해지는지를 정량화하기 위해 EAD를 정의합니다. 이를 통해 기존 방법들의 선택 병목 현상을 수학적으로 입증하며, 그 오버헤드가 단순히 계산 비용 때문이 아니라 잘못된 선택의 위험성에 의해 발생함을 보여줍니다.
2. DEX 알고리즘: 확장-확정-붕괴 절차를 통해 깊이 탐색을 구현하는 무손실 디코딩 알고리즘을 제시합니다. 이는 표준 AR 디코딩과 정확히 동일한 동작을 유지하면서도 토큰 확정 비용을 줄입니다.
3. 실험적 검증: 저자들은 LayerSkip으로 학습된 CodeLlama-34B 및 Llama-2-70B와 같은 조기 종료 학습 모델과, 수학적 추론(GSM8K), 코드 생성(HumanEval), 요약(XSum) 작업을 포함한 표준 LLM(CodeLlama-34B, Llama-2-70B, Qwen3-32B)에 대해 DEX의 성능을 입증했습니다.

4. 실험 결과

저자들은 조기 종료 학습 모델(LayerSkip-trained CodeLlama-34B 및 Llama-2-70B)과 표준 LLM(CodeLlama-34B, Llama-2-70B, Qwen3-32B) 모두에서 DEX를 평가했습니다.

• 처리량 이득 (Throughput Gains): DEX는 깊이 선택 베이스라인들보다 일관되게 높은 초당 토큰 수(TPS)를 달야했습니다. LayerSkip 모델에서 DEX는 선택 방식보다 이론적인 "오라클" 헤드룸의 더 큰 부분을 실현했습니다.
• 해상도에 따른 스케일링: 탐색 해상도가 정밀해질수록(즉, 더 많은 탐색기가 추가될수록) 처리량이 향 향상되었습니다. 예를 들어, Llama-2-70B에서 탐색기를 3개에서 8개로 늘리면 오라클 속도 향상과의 간극이 크게 좁혀졌습니다.
• 경쟁력: 엔드 투 엔드 비교에서, DEX(3~4개의 탐색기 사용 시)는 매칭된 하드웨어 리소스 상에서 멀티 GPU 텐서 병렬성(TP) 및 분산 추측 디코딩(PEARL, EAGLE)과 경쟁할 만한 수준을 보였습니다.
• 절제 연구 (Ablation): 중간 브랜치를 제거한 경우(단일 종료 변형), 성능 이득이 단순히 레이어를 병렬화하는 것이 아니라 여러 깊이를 탐색하는 것에서 구체적으로 온다는 것을 보여주었습니다.

5. 의의 및 주장

본 논문은 병렬 깊이 탐색이 LLM 디코딩의 미사용된 깊이 축을 활용할 수 있는 확장 가능한 방법을 제공한다고 주장합니다.

• 패러다임 전환: 단일 종료 지점에 도박을 거는 기존 방식과 달리, DEX는 깊이 적응을 해상도로 제어 가능한 문제로 취급합니다. 여러 깊이를 탐색함으로써, "언제 종료할 것인가"에 대한 불확실성을 탐색 세트의 간격에 의해 결정되는 유한한 오버헤드로 전환합니다.
• 무손실 효율성: DEX는 (선택적 어댑터를 제외한) 표준 모델의 구조 변경 없이도 출력의 정확성(무손실)을 해치지 않으면서 이러한 이득을 달성합니다.
• 상보성: 저자들은 DEX가 "토큰 축"(얼마나 많은 미래 토큰을 예측하는가) 가속화 방법과 상보적인 "깊이 축"(얼마나 일찍 토근이 결정되는가)에서 작동한다고 언급합니다. 이들을 결합하는 것이 가능하지만, 복잡한 조정이 필요하다고 덧붙였습니다.

결론적으로, 본 논문은 단일 깊이 선택을 병렬 탐색으로 대체함으로써, 표준 자기회귀 모델의 정확한 동작을 유지하면서도 디코딩 지연 시간을 크게 줄일 수 있음을 입증합니다.