SparseD: Разреженное внимание для диффузионных языковых моделей

Аннотация

Хотя диффузионные языковые модели (DLMs) представляют собой перспективную альтернативу авторегрессивным моделям (ARs), существующие открытые DLMs страдают от высокой задержки при выводе. Этот узкий момент в основном обусловлен квадратичной сложностью механизма внимания относительно длины контекста при вычислении всех пар запрос-ключ. Интуитивно, для снижения этой сложности естественной стратегией является ограничение внимания разреженными паттернами, сохраняющими только наиболее релевантные связи. Такие подходы хорошо зарекомендовали себя в ARs, где внимание следует фиксированным и четко определенным разреженным паттернам. Однако в DLMs мы наблюдаем иное поведение разреженности: (1) паттерны внимания различаются между головами, (2) паттерны внимания в каждой голове остаются высоко схожими на всех шагах денизинга, и (3) ранние шаги денизинга критически важны для генерации. Эти наблюдения делают методы разреженного внимания, разработанные для ARs, в значительной степени неприменимыми к DLMs, так как они не учитывают специфичные для голов структуры и рискуют ухудшить генерацию при применении на ранних шагах денизинга. Для решения этих проблем мы предлагаем SparseD — новый метод разреженного внимания для DLMs. Используя наши наблюдения, SparseD требует предварительного вычисления специфичных для голов разреженных паттернов только один раз, повторно используя их на всех шагах. Это предотвращает повторное вычисление разреженных паттернов на каждом шаге денизинга. В то же время SparseD использует полное внимание на ранних шагах, а затем переключается на разреженное внимание на более поздних шагах для сохранения качества генерации. В совокупности это делает SparseD практичным и эффективным решением для применения DLMs в задачах с длинным контекстом. Экспериментальные результаты показывают, что SparseD обеспечивает ускорение без потерь, достигая до 1,50-кратного ускорения по сравнению с FlashAttention при длине контекста 64k и 1024 шагах денизинга.

English

While diffusion language models (DLMs) offer a promising alternative to autoregressive models (ARs), existing open-source DLMs suffer from high inference latency. This bottleneck is mainly due to the attention's quadratic complexity with respect to context length in computing all query-key pairs. Intuitively, to reduce this complexity, a natural strategy is to restrict attention to sparse patterns that retain only the most relevant connections. Such approaches are well-established in ARs, where attention follows fixed and clearly defined sparse patterns. However, in DLMs, we observe distinct sparsity behaviors: (1) attention patterns vary across heads, (2) attention patterns in each head remain highly similar across denoising steps, and (3) early denoising steps are critical for generation. These findings render sparse attention methods designed for ARs largely incompatible with DLMs, as they fail to capture head-specific structures and risk degrading generation when applied in early denoising steps. To address these challenges, we propose SparseD, a novel sparse attention method for DLMs. Leveraging the observations, SparseD only requires pre-computing head-specific sparse patterns one time, and reuses them across all steps. This prevents recomputing sparse patterns at each denoising step. Meanwhile, SparseD uses full attention in the early steps, then switches to sparse attention later to maintain generation quality. Together, these establish SparseD as a practical and efficient solution for deploying DLMs in long-context applications. Experimental results demonstrate that SparseD achieves lossless acceleration, delivering up to 1.50times speedup over FlashAttention at a 64k context length with 1,024 denoising steps.

SparseD: Разреженное внимание для диффузионных языковых моделей

SparseD: Sparse Attention for Diffusion Language Models

Аннотация

Support