FOCUS cuts wasted computing in diffusion language models

About

An emerging alternative to the way today's large language models generate text could become more efficient to run at scale following research led by KAUST. A team from KAUST and The Chinese University of Hong Kong (CUHK) has developed FOCUS, a system that reduces unnecessary computing during AI text generation and allows substantially more work to be handled by the same hardware.

In experiments, FOCUS increased throughput by up to 2.32 times compared with the standard inference system used as a baseline. When the researchers tested larger blocks of text, where unnecessary computing becomes even more pronounced, the improvement reached up to 3.52 times, while the quality of generated responses was maintained or improved.

FOCUS is designed for diffusion large language models, or DLLMs, an emerging approach to generative AI. Most familiar language models build a response sequentially, predicting one piece of text after another. Diffusion language models work differently. Rather than generating a response strictly from left to right, they work across multiple token positions at once, progressively resolving the text over a series of steps.

That flexibility comes at a cost. The researchers found that these models repeatedly spend computing power on parts of a response that are not yet ready to be generated. In their experiments, only about 10% of the positions being processed were typically ready to move forward at each step, yet calculations were still being carried out across the entire group.

"We found that diffusion language models were doing a large amount of work that did not contribute to the response at that particular step," said first author Kaihua Liang, a Ph.D. student at KAUST. "That led us to ask whether the model could tell us, early enough, which parts were actually ready to move forward so that we could focus the computing effort there instead."

The team discovered that it could. Patterns in the way the model pays attention to different parts of the diffusion text provided an early indication of which positions were likely to be ready next. This relationship between the model's internal attention and what it was likely to generate became the key insight behind FOCUS.

FOCUS uses this information to concentrate computing power where it is most useful. It temporarily holds back work on positions that are unlikely to be ready and returns to them later. In doing so, the system reduces the number of positions that need to be processed through much of each generation step by approximately 65% to 80%, without retraining or changing the underlying language model.

The computing capacity saved through this process can instead be used to handle more requests. Across tests covering mathematical reasoning, coding and instruction following, FOCUS achieved its higher throughput while preserving or improving the quality of the model's responses.

"The important finding is that the model already contains information that can help us decide where computation is useful,"said Professor Marco Canini. "FOCUS turns that information into a practical way of reducing unnecessary work. As AI systems serve larger workloads, improving how efficiently we use the available computing infrastructure becomes increasingly important."

FOCUS takes a different approach from previous efforts to accelerate diffusion language models. Earlier work has concentrated largely on storing and reusing calculations the model has already made. The KAUST-led team instead asked whether some of those calculations could be avoided. Their results show that a model's own internal signals can help determine where computing resources should be directed as text is generated.

By reducing the computation needed for each response, FOCUS allows the same hardware to support larger workloads, an increasingly important consideration as generative AI systems scale. The researchers have released FOCUS as open-source software so other teams can reproduce the results and build on the approach.