Learning to Read the Contextual Tokens in Diffusion Transformers

Learning to Read the Contextual Tokens in Diffusion Transformers. By Omer Dahary and four others, from Tel Aviv University and Cornell University. Published on the fifth of October. No code repository is linked to it, though it does publish its weights.

Multimodal Diffusion Transformers, or MM-DiTs, process visual and textual representations jointly throughout generation. As they do, they keep updating the text tokens through multimodal attention, so those tokens stop being a fixed copy of the prompt and become dynamic contextual tokens. What they are for has not been well understood. The claim of this work is that they carry a rich, global representation of the scene being generated, and that this can be read out in plain language.

The reading instrument is a lightweight bottleneck network. It maps intermediate contextual tokens into the input space of a frozen Large Language Model, so that the language model can answer questions about the emerging image directly from those hidden representations. Interrogated this way, the tokens give up a lot. Generation-specific semantics become accessible surprisingly early in denoising, including attributes the prompt left underspecified. Finer details become readable progressively, as denoising goes on.

One result stands out. The information remains decodable even when the MM-DiT is given an empty prompt. With no text to start from, whatever the contextual tokens hold must have accumulated from the evolving visual representation itself. The authors also find that generations whose contextual representations are more readable tend to receive higher human-preference scores.

That last observation turns the reader into a lever. The authors introduce Contextual Alignment, a training technique that explicitly reinforces the visual-semantic information encoded in the contextual tokens. It improves generation quality and distributional coverage.

The claims concern MM-DiTs, the architecture in which text tokens are updated alongside the image. The view into the tokens is through one instrument, a trained bottleneck feeding a frozen language model, so what counts as readable is readable to that reader. The link between readability and human preference is a tendency, not a guarantee for any single image. And Contextual Alignment is reported to improve quality and coverage, without the abstract quantifying by how much.

Two roles end up assigned to the same tokens. They are an interpretable window onto what an MM-DiT is building as it denoises, and they are a target that training can push on to build it better. A part of the model that looked like plumbing for the prompt turns out to hold a running description of the image.