Уточнение скрытых состояний при декодировании: улучшение диффузионных языковых моделей через уточнение состояний убеждений

Latent Refinement Decoding: Enhancing Diffusion-Based Language Models by Refining Belief States

October 13, 2025

Авторы: Qinglin Zhu, Yizhen Yao, Runcong Zhao, Yanzheng Xiang, Amrutha Saseendran, Chen Jin, Philip Alexander Teare, Bin Liang, Yulan He, Lin Gui

cs.AI

Аннотация

Авторегрессионные (AR) модели остаются стандартом для генерации естественного языка, но по-прежнему страдают от высокой задержки из-за строго последовательного декодирования. Недавние подходы, вдохновленные диффузией, такие как LlaDA и Dream, смягчают эту проблему, генерируя текст параллельно, однако они имеют два ключевых ограничения: потерю информации, так как предсказательные распределения для незавершенных токенов отбрасываются на каждом шаге, и преждевременное принятие решений, когда локальные решения принимаются без достаточной глобальной координации. Мы представляем Latent Refinement Decoding (LRD) — двухэтапную структуру с Latent Refinement и Predictive Feedback Loop. На первом этапе сохраняются замаскированные позиции как распределенные смеси предсказанных токенов и маскирующего эмбеддинга, что позволяет модели формировать более глобально согласованные представления. На втором этапе постепенно фиксируются уверенные токены, в то время как неопределенные остаются для итеративной обратной связи. Динамика KL-дивергенции обеспечивает принципиальный и надежный критерий для сходимости и ранней остановки. Эксперименты в области программирования (HumanEval +6.3, MBPP +2.6) и логического мышления (GSM8K +2.9, MATH500 +3.8) показывают, что LRD повышает точность, обеспечивая ускорение до 10.6 раз, что делает его мощной и универсальной альтернативой для параллельной генерации последовательностей.

English

Autoregressive (AR) models remain the standard for natural language generation but still suffer from high latency due to strictly sequential decoding. Recent diffusion-inspired approaches, such as LlaDA and Dream, mitigate this by generating in parallel, yet they suffer from two core limitations: information loss, as predictive distributions for non-finalized tokens are discarded at each step, and premature commitment, where local decisions are made without sufficient global coordination. We introduce Latent Refinement Decoding (LRD), a two-stage framework with Latent Refinement and a Predictive Feedback Loop. The first stage maintains masked positions as distributional mixtures of predicted tokens and the mask embedding, allowing the model to establish more globally consistent beliefs. The second stage progressively finalizes confident tokens while retaining uncertain ones for iterative feedback. KL-divergence dynamics provide a principled and reliable criterion for convergence and early stopping. Experiments across coding (HumanEval +6.3, MBPP +2.6) and reasoning (GSM8K +2.9, MATH500 +3.8) show that LRD improves accuracy while delivering speedups of up to 10.6x, making it a strong and versatile alternative for parallel sequence generation.