💬 NLP

Just on Time: Token-Level Early Stopping for Diffusion Language Models

이 논문은 경량화된 예측 신호를 기반으로 안정적인 토큰을 동적으로 식별하고 확정함으로써, 다양한 작업에 걸쳐 생성 품질을 유지하면서도 계산 비용을 크게 줄이는 디퓨전 언어 모델을 위한 학습이 필요 없는 토큰 수준의 조기 종료 방법을 소개한다.

원저자: Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi, Serhii Dmytryshyn, Dmytro Khamula, Michal Zakrzewski, Damian Rynczak, Jacek Małecki, Taras Rumezhak, Volodymyr Karpiv

게시일 2026-08-04

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zakhar Kohut, Severyn Shykula, Mykola Vysotskyi, Serhii Dmytryshyn, Dmytro Khamula, Michal Zakrzewski, Damian Rynczak, Jacek Małecki, Taras Rumezhak, Volodymyr Karpiv

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대한 직소 퍼즐을 맞추고 있다고 상상해 보세요. 하지만 상자에 그려진 그림을 보는 대신, 대신 빈 회색 사각형들로 가득 찬 상자에서 시작합니다. 이 퍼즐을 풀려면 각 사각형이 어떤 색이어야 하는지 추측한 다음, 전체 그림을 보고 어떤 추측이 틀렸는지 깨달아 다시 색칠해야 합니다. 여러분은 그림이 마침내 제대로 보일 때까지 이 "추측하고 수정하기" 과정을 수백 번 반복합니다. 이것이 바로 **디퓨전 언어 모델(Diffusion Language Model)**이라 불리는 새로운 종류의 AI가 글을 쓰는 방식입니다. 한 단어씩 타이피스트처럼 써 내려가는 기존의 AI와 달리, 이 모델들은 빈 페이지에서 시작하여 캔버스에 세부 사항을 천천히 더해가는 화가처럼 문장 전체를 한꺼번에 정교하게 다듬어 나갑니다.

문제는 이 채색 과정이 믿기 힘들 정도로 느리다는 점입니다. AI가 문장의 쉬운 부분(예: "The"나 "cat")을 단 몇 번의 붓질만으로도 파악하더라도, 만약을 위해 그 부분을 계속 덧칠하곤 합니다. 이는 마치 요리사가 수프 맛을 보고 완벽하다는 것을 깨달았음에도 불구하고, 손님에게 내놓기 전까지 한 시간 동안 계속 저으면서 맛을 보는 것과 같습니다. 이는 엄청난 양의 컴퓨터 연산 능력과 시간을 낭비합니다. 과학자들이 던지는 질문은 이것입니다. "어떻게 하면 AI에게 '이봐, 이 부분은 제대로 됐으니 그만 작업하고 어려운 부분에 집중해'라고 말할 수 있을까?"

여기서 **JoT(Just on Time)**라는 새로운 방법이 등장합니다. JoT를 AI 화가 옆에 서 있는 똑똑한 감독관이라고 생각해 보세요. 전체 팀에게 동시에 멈추라고 명령하는 대신, JoT는 캔버스 위의 개별 단어를 하나하나 관찰합니다. AI가 특정 단어(예를 들어 "cat"이라는 단어)에 대해 확신이 매우 높아지는 순간(예: 99% 확신할 때), JoT는 "그 단어는 고정해! 더 이상 건드리지 마!"라고 외치고, AI의 주의를 다음의 흐릿하고 불확실한 단어로 옮깁니다.

JoT를 개발한 연구진은 이러한 "토큰 수준의 조기 종료(token-level early stopping)"가 놀라운 효과를 낸다는 것을 발견했습니다. 그들은 Dream-7B와 LLaDA-8B라는 두 가지 강력한 AI 모델을 대상으로 수학 문제 풀이, 상식 퀴즈, 문장 완성, 그리고 코드 작성이라는 네 가지 과제를 테스트했습니다. 결과는 놀라웠습니다. GSM8K라는 수학 추론 테스트에서 JoT는 AI를 단 2.3점 정도의 미미한 점수 하락만으로 5.5배 더 빠르게 만들었습니다. 코딩 챌린지인 HumanEval에서의 속도 향상은 훨씬 더 극적이었는데, AI가 거의 모든 정답을 맞히면서도 속도가 무려 20배 가까이(19.6배) 빨라졌습니다.

이 논문은 모든 단어를 동시에 멈추게 하는 "일률적인(one-size-fits-all)" 방식에 반대합니다. 대신, JoT는 영리한 기술을 사용합니다. AI가 이미 완성한 단어 옆에 있는 단어들에 대해서는 확신의 기준치를 낮추는 것입니다. 만약 AI가 문장의 시작 부분을 이미 완벽하게 해냈다면, 그다음 단어는 조금 더 일찍 신뢰할 수 있습니다. 이를 통해 AI는 정말로 혼란스러운 부분에만 에너지를 집중하고, 불필요한 반복 작업을 건너뛸 수 있습니다.

이 방법은 매우 빠르지만, 저자들은 이것이 마법은 아니라는 점을 주의 깊게 언급합니다. 우리는 AI가 까다로운 수학 단계에서 막혀서 숫자를 너무 빨리 확정 지을 때 실수가 발생한다는 것을 측정했습니다. 하지만 이러한 오류는 드물며, 속도 측면에서의 이점이 압도적입니다. 논문은 각 단어가 자신만의 완벽한 순간에 "종료"되도록 허용함으로써, 이 강력하지만 느린 AI 모델들을 재학습하거나 핵심 설계를 변경하지 않고도 일상적인 용도로 훨씬 더 실용적으로 만들 수 있다고 제안합니다. 이는 AI가 더 열심히 일하는 것이 아니라, 더 똑똑하게 일하도록 만드는 단순하고도 훈련이 필요 없는 해결책입니다.

기술 요약: 확산 언어 모델(Diffusion Language Models)을 위한 Just on Time (JoT)

문제 정의

확산 언어 모델(Diffusion Language Models, DLMs)은 완전히 마스킹된 시퀀스를 깨끗한 출력으로 반복적으로 정제하며 텍스트를 생성한다. 이러한 방식은 병렬 토큰 예측과 양방향 문맥 통합을 가능하게 하지만, 디코딩 효율성 측면에서는 비효전적이다. 표준 DLM은 전체 시퀀스에 대해 고정된 수의 정제 단계( $T$ )를 요구하는데, 이는 많은 토큰이 최종 단계에 도달하기 훨씬 전부터 이미 안정적이고 높은 신뢰도의 예측에 도달함에도 불구하고 동일한 단계를 적용하기 때문이다. 이는 이미 수렴한 토큰들을 계속해서 정제하게 만듦으로써 불필요한 계산 오버헤드를 초래한다. 기존의 조기 종료(early-stopping) 방법들은 주로 전역적인 기준을 적용하거나 특정 작업에 맞춘 적응에 의존하며, 단일 생성 과정 내에서 개별 토큰의 다양한 수렴 속도를 동적으로 처리하지 못한다.

방법론: JoT (Just on Time)

저자들은 각 위치에서 독립적으로 수렴을 식별하도록 설계된 학습이 필요 없는(training-free) 토큰 단위 조기 종료 메커니즘인 JoT를 제안한다. JoT는 서로 다른 토큰들이 시퀀스 내에서 각기 다른 단계에서 "종료(exit)"할 수 있도록 허용하여, 필요한 곳에만 계산을 집중시킨다.

핵심 구성 요소

신뢰도 지표 (Confidence Metric):
JoT는 각 마스킹된 위치 $i$ 에서의 예측 확실성을 상위 두 예측 확률( $p_1$ 및 $p_2$ )의 비율을 사용하여 정량화한다:
$r_i = \frac{p_1}{p_2 + \epsilon}$
이 지표는 스케일링되지 않은 로짓(logits)에서 유도되므로 샘플링 온도(temperature) 변화에 불변한다. 높은 비율은 최상위 후보에 대한 결정적인 선호를 나타내며, 비율이 1에 가까우면 불확실성을 의미한다.
공간적 변조 (Spatial Modulation):
이미 해결된(unmasked) 문맥에 인접한 토큰은 더 풍부한 국소 정보와 더 낮은 유효 노이즈의 혜택을 받는다는 점을 인식하여, JoT는 공간적으로 적응형 임계값을 도입한다.
- 기하학적 커널(geometric kernel)은 윈도우 $D$ 내에서 미해결 토큰과 위치 $i$ 사이의 근접성을 기반으로 가중치 $w_i$ 를 계산한다.
- 이 가중치는 소프트닝 인자(softening factor) $\phi_i \in [0, 1]$ 로 정규화된다.
- 위치 $i$ 에 대한 수용 임계값 $\tau_i$ 는 최대 임계값 $\tau_{max}$ 와 최소 임계값 $\tau_{min}$ 사이에서 보간된다:
  $\tau_i = \tau_{max} - (\tau_{max} - \tau_{min}) \cdot \phi_i$
- 해결된 문맥 근처의 토큰은 낮은 임계값을 받아 (더 쉽게 확정되고), 마스킹된 영역의 중앙에 있는 토큰은 더 엄격한 임계값을 유지한다.
조기 종료 결정 (Early-Exit Decision):
각 단계에서 마스킹된 토큰의 신뢰도 비율 $r_i$ 가 적응형 임계값 $\tau_i$ 를 초과하면, 해당 토큰은 확정(argmax 또는 샘플링을 통해 선택)되고 마스킹된 위치 집합에서 제거된다. 이 과정은 모든 토큰이 확정되거나 최대 단계 수에 도달할 때까지 계속된다.

이론적 분석

본 논문은 (Proposition 1)을 통해 "선두 쌍 역전(leading-pair reversals)"—즉, 확정된 토큰이 더 많은 문맥이 드러남에 따라 나중에 차순위 후보에게 추월당하는 현상—의 확률에 관한 공식적인 분석을 제공한다. 모델의 예측이 실제 조건부 분포와 일치한다(calibration)는 가정하에, 이러한 역전의 확률은 $2/(1+\tau)$ 로 제한된다. 저자들은 높은 임계값에서의 잔류 오차는 조기 종료 메커즘 자체보다는 캘리브레이션 실패에서 기인할 수 있다고 언급한다.

주요 기여

JoT 프레임워크: 해결된 문맥과의 공간적 근접성에 따라 수용 임계값을 조정하는 학습이 필요 없는 토큰 단위 조기 종료 방법.
포괄적 평가: 두 가지 최첨단 모델(Dream-7B 및 LLaDA-8B)을 대상으로 네 가지 다양한 벤치마크(GSM8K(수학), MMLU(지식), HellaSwag(상식), HumanEval(코드))에서 실험 수행.
절제 연구 (Ablation Studies): 임계값 선택과 공간적 변조의 효과를 분리하여 분석함으로써 하이퍼파라미터 설정에 대한 실질적인 가이드를 제공.

실험 결과

저자들은 JoT를 전체 디코딩 및 다른 조기 종료 베이스라인(Prophet, KLASS, Fast-dLLM)과 비교 평가하였다.

속도-품질 트레이드오프 (Speed-Quality Trade-off):
- Dream-7B: JoT는 GSM8K에서 최대 5.54배의 속도 향상을, HumanEval에서 19.60배의 속도 향상을 달이트했다. 정확도 하락은 미미했다 (예: GSM8K에서 -2.3 포인트, HumanEval에서 -0.6 포인트). 대부분의 벤치마크에서 성능은 전체 디코딩 대비 약 3% 이내의 차이를 유지했다.
- LLaDA-8B: JoT는 모든 벤치마크에서 전체 디코딩 대비 정확도 차이를 3.1 포인트 이내로 유지하면서 경쟁력 있는 속도 향상을 제공했다 (예: GSM8K에서 3.75배).
비교: JoT는 일반적으로 더 큰 정확도 하락을 겪은 Prophet보다 우수한 성능을 보였으며, 여러 작업에서 KLASS 및 Fast-dLLM보다 더 나은 속도-정확도 파레토 프런티어(Pareto frontier)를 제공했다.
결합 가능성 (Composability): JoT는 단계당 레이턴시를 줄이는 Fast-dLLM의 DualCache와 성공적으로 통합되었으며, 이는 단계 축소와 단계당 최적화가 서로 직교하며 결합될 수 있음을 입증한다.
실패 모드: 500개의 GSM8K 샘플 분석 결과, 오류는 주로 모델이 중요한 추론 분기(예: "/ 2" 대신 "× 2"를 선택하는 경우)에서 토큰을 성급하게 확정할 때 발생했다. 이러한 오류는 새로운 형태의 실패 모드가 아니라, 전체 디코딩과 비교했을 때 다른 빈도로 발생하는 기존의 오류였다.

의의 및 주장

본 논문은 JoT가 재학습이나 작업별 미세 조정 없이도 확산 언어 모델의 디코딩 비효율성에 대한 설득력 있는 해결책을 제공한다고 주장한다. 충분한 신뢰도를 보이는 즉시 토큰을 확정함으로써, JoT는 필요한 총 확산 단계를 실질적으로 줄인다.

저자들은 JoT를 Dream-7B에서 테스트된 조기 종료 방법 중 최상의 속도-품질 트레이드오프를 달성하고 LLaDA-8B에서도 경쟁력을 유지하는 방법으로 규정한다. 그들은 이 방법이 특히 효과적인 이유가 "명백한" 문맥 토큰에 대한 중복적인 정제를 건너뛰면서도 복잡한 추론 단계에 필요한 계산 노력은 보존할 수 있기 때문이라고 강조한다. 논문은 최적의 임계값이 작업마다 다를 수 있지만, 하이퍼파라미터에 대한 방법론의 강건성(robustness) 덕분에 다양한 애플리케이션에 실용적이라고 결론짓는다.