Многоблочные диффузионные языковые модели

Аннотация

Блочные диффузионные языковые модели (BD-LM) улучшают генерацию текста на основе диффузии за счёт кэширования KV и генерации переменной длины. Естественным следующим шагом является их расширение от одноблочной диффузии (SingleBD) к многоблочной диффузии (MultiBD), где одновременно декодируется текущий набор последовательных блоков для обеспечения межблочного параллелизма. Однако существующие BD-LM в основном обучаются с использованием метода принудительного обучения (teacher forcing), при котором модель наблюдает только один зашумлённый блок при условии чистого префикса. Хотя недавняя стратегия принудительной диффузии (diffusion forcing) вводит видимость между несколькими зашумлёнными блоками, её состояния обучения всё ещё отличаются от вывода MultiBD, где декодирование выполняется над ограниченным текущим набором с гетерогенными шумовыми паттернами для отдельных слотов. Чтобы устранить этот разрыв, мы предлагаем многоблочные диффузионные языковые модели (MBD-LM), полученные путём дообучения BD-LM с использованием многоблочного принудительного обучения (MultiTF). MultiTF объединяет принудительное обучение и принудительную диффузию, выполняя обучение на ограниченных группах шума при условии чистых префиксов, с рандомизированными планировщиками шума, которые лучше соответствуют состояниям вывода MultiBD. Чтобы сделать MultiBD практически выполнимым, мы дополнительно вводим оптимизированный алгоритм декодирования на основе механизма блочного буфера (Block Buffer), который сохраняет повторное использование кэша префикса, сохраняет статическими формы входных данных и преобразует повышенный параллелизм декодирования в ускорение времени выполнения. Эмпирически MBD-LLaDA2-Mini увеличивает среднее количество токенов за прямой проход (TPF) с 3.47 до 6.19 и повышает среднюю точность с 79.95% до 81.03%; в сочетании с DMax MBD-LLaDA2-Mini-DMax достигает среднего TPF 9.34 с падением точности всего на 1.02% на тестах по математике и программированию.

English

Block Diffusion Language Models (BD-LMs) improve diffusion-based text generation with KV caching and flexible-length generation. A natural next step is to extend them from Single-Block Diffusion (SingleBD) to Multi-Block Diffusion (MultiBD), where a running-set of consecutive blocks is decoded concurrently for inter-block parallelism. However, existing BD-LMs are mostly trained under teacher forcing, where the model observes only one noisy block conditioned on a clean prefix. While the recent diffusion forcing strategy introduces visibility among multiple noisy blocks, its training states still differ from MultiBD inference, where decoding operates on a bounded running-set with heterogeneous slot-wise noise patterns. To bridge this gap, we propose Multi-Block Diffusion Language Models (MBD-LMs), obtained by post-training BD-LMs with Multi-block Teacher Forcing (MultiTF). MultiTF integrates teacher forcing and diffusion forcing by training on bounded noise-groups conditioned on clean prefixes, with randomized noise-schedulers that better match MultiBD inference states. To make MultiBD practically executable, we further introduce an optimized decoding algorithm based on the Block Buffer mechanism that preserves prefix-cache reuse, keeps input shapes static, and translates increased decoding parallelism into wall-clock acceleration. Empirically, MBD-LLaDA2-Mini increases average Tokens Per Forward pass (TPF) from 3.47 to 6.19 and improves average accuracy from 79.95% to 81.03%; when combined with DMax, MBD-LLaDA2-Mini-DMax reaches an average TPF of 9.34 with only a 1.02% accuracy drop on math and code benchmarks.