Diffusion LLMs can think EoS-by-EoS

TL;DR

Diffusion LLMs utilize end-of-sequence tokens as a hidden scratchpad, enabling complex reasoning by leveraging EoS representations, which enhances their problem-solving capabilities beyond traditional expectations.

Contribution

This paper introduces the idea that diffusion LLMs think EoS-by-EoS, demonstrating how EoS tokens serve as a hidden computational space for reasoning tasks.

Findings

Adding EoS tokens improves reasoning performance.

Intervening on EoS hidden states alters outputs, confirming their informational role.

Diffusion models can effectively use EoS tokens for complex reasoning.

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.