EntRGi: Энтропийно-ориентированное управление вознаграждением для диффузионных языковых моделей

EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models

February 4, 2026

Авторы: Atula Tejaswi, Litu Rout, Constantine Caramanis, Sanjay Shakkottai, Sujay Sanghavi

cs.AI

Аннотация

Метод управления с помощью вознаграждения успешно применяется для адаптации непрерывных диффузионных моделей во время тестирования; он обновляет каждый шаг денизинга с использованием градиентов от последующей модели вознаграждения. Мы исследуем управление с помощью вознаграждения для дискретных диффузионных языковых моделей, где невозможно провести дифференцирование через естественные выходы модели, поскольку они представляют собой дискретные токены. Существующие подходы либо заменяют эти дискретные токены непрерывными аппроксимациями, либо используют такие методы, как straight-through estimator. В данной работе мы демонстрируем недостатки обоих методов. Первый подход ухудшает обратную связь по градиенту, поскольку модель вознаграждения никогда не обучалась на непрерывных входах. Второй подход приводит к некорректной оптимизации, поскольку градиент, вычисленный на дискретных токенах, используется для обновления непрерывных логитов. Нашим ключевым нововведением является преодоление этого компромисса за счет введения нового механизма под названием EntRGi: Entropy aware Reward Guidance, который динамически регулирует градиенты от модели вознаграждения. Путем модуляции непрерывной аппроксимации с использованием уверенности модели наш подход существенно улучшает управление с помощью вознаграждения, одновременно обеспечивая надежные входные данные для модели вознаграждения. Мы эмпирически подтверждаем наш подход на диффузионной языковой модели с 7 миллиардами параметров, используя 3 различные модели вознаграждения и 3 многозадачных тестовых набора, демонстрируя устойчивое улучшение по сравнению с современными методами.

English

Reward guidance has been applied to great success in the test-time adaptation of continuous diffusion models; it updates each denoising step using the gradients from a downstream reward model. We study reward guidance for discrete diffusion language models, where one cannot differentiate through the natural outputs of the model because they are discrete tokens. Existing approaches either replace these discrete tokens with continuous relaxations, or employ techniques like the straight-through estimator. In this work, we show the downsides of both these methods. The former degrades gradient feedback because the reward model has never been trained with continuous inputs. The latter involves incorrect optimization because the gradient evaluated at discrete tokens is used to update continuous logits. Our key innovation is to go beyond this tradeoff by introducing a novel mechanism called EntRGi: Entropy aware Reward Guidance that dynamically regulates the gradients from the reward model. By modulating the continuous relaxation using the model's confidence, our approach substantially improves reward guidance while providing reliable inputs to the reward model. We empirically validate our approach on a 7B-parameter diffusion language model across 3 diverse reward models and 3 multi-skill benchmarks, showing consistent improvements over state-of-the-art methods.