LLM 디코딩을 위한 깊이 탐색 (Depth Exploration)

자기회귀(Autoregressive) LLM 디코딩은 많은 토큰이 중간 깊이(intermediate depths)에서 예측 가능해짐에도 불구하고, 모든 생성된 토큰을 전체 레이어 스택(layer stack)을 통해 평가합니다. 기존의 손실 없는 깊이 적응형(lossless depth-adaptive) 방법들은 단일한 비최종 종료 깊이(non-final exit depth)를 선택하고 이를 최종 깊이 모델로 검증함으로써 이러한 중복성을 활용합니다. 그러나 우리의 측정 결과에 따르면, 이러한 선택 기반 전략은 상당한 여유 공간(headroom)을 남겨둡니다. 종료 시점을 너무 늦게 선택하면 계산이 낭비되고, 너무 일찍 선택하면 폴백(fallback)이 발생하여 의존적인 초안(dependent drafts)을 버리게 됩니다. 우리는 단일 깊이 선택을 여러 후보 깊이에 대한 병렬 탐색(parallel exploration)으로 대체하는 손실 없는 디코딩 알고리즘인 Depth Exploration Decoding (DEX)를 제안합니다. 각 커밋 위치(commit position)에서 DEX는 최종 깊이 참조 모델(final-depth reference)을 통해 후보들을 검증하고, 정확히 최종 깊이의 토큰을 커밋하며, 재사용 가능한 브랜치 상태(branch states)만을 유지하기 위해 탐색 격자(exploration lattice)를 붕괴(collapse)시킵니다. 이 확장-커밋-붕괴(expand--commit--collapse) 절차는 표준 자기회귀 디코딩과의 동등성을 유지하면서 각 토큰을 커밋하는 비용을 줄여줍니다. 조기 종료 학습(early-exit-trained)된 LLM과 표준 LLM 모두에서, DEX는 대표적인 깊이 선택 베이스라인(depth-selection baselines)보다 우수한 성능을 보이며, 투기적 디코딩(speculative decoding) 및 분산 디코딩(distributed decoding) 방법들과 비교해도 경쟁력 있는 엔드 투 엔드(end-to-end) 처리량을 달성합니다. 또한, DEX는 탐색되는 깊이가 더 세밀해질수록 성능이 향상되며, 이는 병렬 깊이 탐색이 LLM 디코딩의 미사용된 깊이 축(depth axis)을 활용하는 확장 가능한 방법임을 보여줍니다.