[논문 리뷰] Diffusion LLMs can think EoS-by-EoS

Sarah Breckner, Sebastian Schuster|arXiv (Cornell University)|2026. 03. 05.

Topic Modeling인용 수 0

한 줄 요약

이 논문은 확산 LLM이 남은 시퀀스 끝 토큰(EOS)들을 숨겨진 스크래치패드로 사용하여 추론을 개선하는지 여부를 연구하고, 여러 작업에서의 행동실험과 인과적 개입 실험을 통해 증거를 제시한다.

ABSTRACT

Diffusion LLMs have been proposed as an alternative to autoregressive LLMs, excelling especially at complex reasoning tasks with interdependent sub-goals. Curiously, this is particularly true if the generation length, i.e., the number of tokens the model has to output, is set to a much higher value than is required for providing the correct answer to the task, and the model pads its answer with end-of-sequence (EoS) tokens. We hypothesize that diffusion models think EoS-by-EoS, that is, they use the representations of EoS tokens as a hidden scratchpad, which allows them to solve harder reasoning problems. We experiment with the diffusion models LLaDA1.5, LLaDA2.0-mini, and Dream-v0 on the tasks Addition, Entity Tracking, and Sudoku. In a controlled prompting experiment, we confirm that adding EoS tokens improves the LLMs' reasoning capabilities. To further verify whether they serve as space for hidden computations, we patch the hidden states of the EoS tokens with those of a counterfactual generation, which frequently changes the generated output to the counterfactual. The success of the causal intervention underscores that the EoS tokens, which one may expect to be devoid of meaning, carry information on the problem to solve. The behavioral experiments and the causal interventions indicate that diffusion LLMs can indeed think EoS-by-EoS.

연구 동기 및 목표

• 생성 길이가 추론 과제에서 확산 LLM에 어떤 영향을 미치는지 탐구한다.
• 확산 LLM에서 디코딩 단계와 뒤따르는 EoS 토큰의 역할을 구분한다.
• EoS 토큰 표현이 추론에 기여한다는 인과적 근거를 제시한다.
• 확산 모델 간 EoS-별 추론을 명시적 체인-오브-생각(CoT) 프롬프트와 비교한다.

제안 방법

• 세 가지 지시문-튜닝 확산 LLM(LLaDA1.5, LLaDA2.0-mini, Dream-v0)과 자기회귀 기준선(Llama3.1, Qwen3)을 연구한다.
• 생성 길이와 뒤따르는 EoS 토큰을 변화시키기 위해 제어된 프롬프트를 사용하여 추론 성능을 관찰한다.
• 출력에 대한 인과적 영향을 평가하기 위해 EoS 토큰의 숨겨진 상태를 패치한다(대안적 프롬프트).
• 난이도 변화에 따른 추론을 시험하기 위해 additions, entity tracking, Sudoku 데이터셋에서 평가한다.
• 다양한 토큰 예산 하에서 확산 모델의 EoS-별 추론을 체인-오브-생각 프롬프트와 비교한다.

실험 결과

연구 질문

• RQ1 생성 길이를 증가시키면 작업 전반에서 확산 LLM의 추론 성능이 향상되는가?
• RQ2 뒤따르는 EoS 토큰은 디코딩 단계와 무관하게 추론에 기여하는가?
• RQ3 EoS 토큰 표현이 모델의 답을 산출하는 데 인과적으로 관여하는가?
• RQ4 확산 모델과 자기회귀 모델에서 EoS-별 추론은 전통적 체인-오브-생각 프롬프트와 어떻게 비교되는가?

주요 결과

• 생성 길이는 여러 작업에서 확산 LLM의 성능을 향상시키고 충분한 길이에서 자기회귀 모델을 능가할 수 있다.
• 고정된 디코딩 단계 수에서 뒤따르는 EoS 토큰을 추가하면 정확도가 향상되며, 이는 EoS 토큰이 숨겨진 스크래치패드 역할을 함을 시사한다.
• EoS 토큰 표현을 바꾸는 개입이 출력을 바꾸어 EoS 토큰이 문제 해결에 사용되는 정보를 담고 있음을 나타낸다.
• CoT 프롬프팅은 자기회귀 모델에 이득을 주며, 큰 토큰 예산 하에서 확산 모델과 경쟁하거나 능가할 수 있으며, 특히 더 쉬운 작업에서 그렇다.
• LLaDA2.0은 블록-인과적 주의 설계로 인해 뒤따르는 EoS 토큰에서의 이득이 제한적이다.