PC-Sampler: Position-Aware Calibration of Decoding Bias in Masked Diffusion Models

2508.13021v2 cs.AI, cs.CL 2025-08-20
Авторы:

Pengcheng Huang, Shuhao Liu, Zhenghao Liu, Yukun Yan, Shuo Wang, Zulong Chen, Tong Xiao

Резюме на русском

#### Контекст Маскированные диффузионные модели (Masked Diffusion Models, MDMs) стали одними из самых мощных неавторегрессивных тактик для последовательного генерирования. Однако наши предварительные эксперименты показали, что качество генерации MDMs сильно зависит от выбора стратегии декодирования. Особенно заметны проблемы с доминированием тривиальных токенов в начале декодирования и отсутствием глобального контроля траектории. Эти ограничения сильно сказываются на потенциале MDMs. Мы предлагаем Position-Aware Confidence-Calibrated Sampling (PC-Sampler) — универсальную стратегию декодирования, которая объединяет планирование траектории и максимизацию информативности в контексте содержимого. #### Метод PC-Sampler использует позиционно-акцентированную взвешивание для регулирования траектории декодирования и квалибрированный скор уверенности для защиты от преждевременного выбора тривиальных токенов. Основные компоненты: 1. **Планирование траектории**: использование позиционной информации для управления глобальной траекторией декодирования. 2. **Калибровка уверенности**: добавление квалибрированного компонента в скор уверенности для подавления выбора тривиальных токенов в ранних этапах. 3. **Интеграция**: синергетическое объединение двух компонентов для улучшения качества генерации. #### Результаты Мы проверили PC-Sampler на трех моделях MDM (T5, BART, LLaMA) по метрикам BLEU, ROUGE и METEOR на семь бенчмарок, включая задачи логического рассуждения и планирования. Результаты показали, что PC-Sampler повышает производительность этих моделей на 10% в среднем по сравнению с другими стратегиями декодирования. Основные выводы: - Улучшенная точность и консистентность генерации. - Снижение частоты выбора тривиальных токенов. - Значительное сужение разрыва с авторегрессионными моделями. #### Значимость PC-Sampler может применяться в различных областях, где необходимо высокое качество генерации последовательностей, такие как логическое рассуждение, планирование и моделирование текста. Основные преимущества: - Улучшение качества генерации с помощью позиционно-акцентированного управления. - Уменьшение частоты тривиальных выборов за счет калиброванной уверенности. - Устранение ограничений других стратегий декодирования. #### Выводы PC-Sampler достигает существенных улучшений в производительности MDMs и сужает разрыв с авторегрессионными моделями. Наши результаты открывают новые пути для развития неавторегрессивных моделей. Будущие исследования будут сфокусированы на оптимизации PC-Sampler для задач с большим количеством токенов и улучшении его мобильности в разных сценариях.

Abstract

Recent advances in masked diffusion models (MDMs) have established them as powerful non-autoregressive alternatives for sequence generation. Nevertheless, our preliminary experiments reveal that the generation quality of MDMs is still highly sensitive to the choice of decoding strategy. In particular, widely adopted uncertainty-based samplers suffer from two key limitations: a lack of global trajectory control and a pronounced bias toward trivial tokens in the early stages of decoding. These shortcomings restrict the full potential of MDMs. In this work, we introduce Position-Aware Confidence-Calibrated Sampling (PC-Sampler), a novel decoding strategy that unifies global trajectory planning with content-aware informativeness maximization. PC-Sampler incorporates a position-aware weighting mechanism to regulate the decoding path and a calibrated confidence score to suppress the premature selection of trivial tokens. Extensive experiments on three advanced MDMs across seven challenging benchmarks-including logical reasoning and planning tasks-demonstrate that PC-Sampler consistently outperforms existing MDM decoding strategies by more than 10% on average, significantly narrowing the performance gap with state-of-the-art autoregressive models. All codes are available at https://github.com/NEUIR/PC-Sampler.

Ссылки и действия