Search papers, labs, and topics across Lattice.
This paper investigates the surprising phenomenon that diffusion LLMs perform better on complex reasoning tasks when generating significantly more end-of-sequence (EoS) tokens than necessary. Through controlled prompting experiments on tasks like Addition, Entity Tracking, and Sudoku, the authors demonstrate that adding EoS tokens enhances reasoning capabilities. By patching the hidden states of EoS tokens with those from counterfactual generations, they causally demonstrate that these tokens act as a hidden scratchpad for computations, indicating that diffusion LLMs "think EoS-by-EoS."
Diffusion LLMs secretly use end-of-sequence tokens as a hidden scratchpad, boosting their reasoning abilities in complex tasks.
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.
