Изучение политик демаскировки для диффузионных языковых моделей

Аннотация

Диффузионные (крупные) языковые модели (dLLM) в настоящее время достигают аналогичной производительности на последующих задачах по сравнению с их авторегрессионными аналогами, сохраняя при этом перспективу большей эффективности на этапе вывода. Особенно успешным вариантом является дискретная диффузия с маскированием, при которой буфер, заполненный специальными маскирующими токенами, постепенно заменяется токенами, сэмплированными из словаря модели. Повышение эффективности может быть достигнуто за счет параллельного раскрытия нескольких токенов, однако одновременное раскрытие слишком большого их количества чревато снижением качества генерации. Таким образом, ключевым аспектом проектирования dLLM является процедура выборки, которая на каждом шаге диффузионного процесса определяет, какие токены заменить. Действительно, недавние исследования показали, что эвристические стратегии, такие как пороговая фильтрация по уверенности, приводят к более высокому качеству и пропускной способности по токенам по сравнению со случайным раскрытием. Однако у таких эвристик есть недостатки: они требуют ручной настройки, и мы наблюдаем, что их производительность ухудшается с увеличением размера буфера. В данной работе мы вместо этого предлагаем обучать процедуры выборки с помощью обучения с подкреплением. В частности, мы формализуем выборку в маскированной диффузии как марковский процесс принятия решений, в котором dLLM выступает в качестве среды, и предлагаем облегченную архитектуру политики на основе однослойного трансформера, которая отображает уверенности токенов dLLM в решения о раскрытии. Наши эксперименты показывают, что эти обученные политики соответствуют производительности передовых эвристик в сочетании с полуавторегрессионной генерацией, превосходя их в условиях полной диффузии. Мы также исследуем переносимость этих политик, обнаруживая, что они способны обобщаться на новые базовые dLLM и большие длины последовательностей. Однако мы также наблюдаем, что их производительность снижается при применении к данным из других предметных областей, и что тонкая настройка компромисса между точностью и эффективностью может быть сложной задачей в нашем подходе.

English

Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary. Efficiency can be gained by unmasking several tokens in parallel, but doing too many at once risks degrading the generation quality. Thus, one critical design aspect of dLLMs is the sampling procedure that selects, at each step of the diffusion process, which tokens to replace. Indeed, recent work has found that heuristic strategies such as confidence thresholding lead to both higher quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger buffer sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy architecture based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive generation, while outperforming them in the full diffusion setting. We also examine the transferability of these policies, finding that they can generalize to new underlying dLLMs and longer sequence lengths. However, we also observe that their performance degrades when applied to out-of-domain data, and that fine-grained tuning of the accuracy-efficiency trade-off can be challenging with our approach.

Изучение политик демаскировки для диффузионных языковых моделей

Learning Unmasking Policies for Diffusion Language Models

Аннотация

Support