Teaching diffusion vision-language models to look at the image only when needed

Diffusion vision-language models spend most of their effort re-reading the same picture, and much of that effort may be unnecessary. In a preprint posted to arXiv, which has not yet been peer reviewed, researchers at KAIST AI describe a decoding method called LOWN, short for Look Only When Needed. It checks at each step whether the image is likely to change the answer before processing it. The authors report that it keeps 99.6% of the original model's benchmark score while using, on average, 35.3% of the computation spent on decoding.

Most familiar language models write one token, a word or fragment of a word, at a time, from left to right. Diffusion language models work differently. They start with an answer made entirely of masked placeholders and fill it in over a series of denoising steps, revealing, or unmasking, a few tokens on each pass. A diffusion vision-language model does the same while also taking an image as input, which it breaks into visual tokens of its own. According to the study, these models process their entire input at every step. On the benchmarks examined, visual tokens account for 89-97% of that input, so most of the decoding computation goes into repeatedly processing an image the model has already seen.

One existing remedy, visual token pruning, discards some of the image tokens. The authors argue that this loses the removed information and still leaves the rest of the image to be processed at every step. To measure how often the image is actually needed, they ran LLaDA-V, an existing diffusion vision-language model, with and without the image at each decoding step and compared its predictions. Taking the image away changed a token being unmasked at only 22.8% of steps on average, and in some cases at as few as 7.9%.

LOWN builds on that observation. At every step it first predicts the masked tokens without the image, skipping the visual tokens that make up most of the input, and compares the result with the most recent prediction that did use the image. If the two disagree, the model processes the image again. If they agree, the agreement acts as a certificate that the earlier image-informed prediction can be reused, and the number of tokens certified this way sets how many are unmasked in that step. The method thus decides both whether to look and how quickly to proceed. The authors also found that the need for the image differs from one benchmark to another, which they say is why an adaptive decision is required rather than a fixed schedule for when to look. LOWN needs no additional training and no hand-tuned threshold.

Beyond the headline figures, the authors report that the method works across other diffusion vision-language models based on LLaDA, can be combined with visual token pruning, and runs alongside key-value caching, a technique that stores intermediate results so they need not be recomputed, for further savings. They have released their code.

The results are those reported by the authors on ten benchmarks and on models from the LLaDA family; the abstract does not say how the approach would fare with other designs. Savings are given in FLOPs, a count of arithmetic operations, rather than in time, and because the image matters more for some tasks than others, the benefit is likely to vary with the job at hand.