LoopMDM: зацикленные диффузионные модели учатся думать глубже, а не шире

LoopMDM: зацикленные диффузионные модели учатся думать глубже, а не шире

Все современные языковые модели — от GPT-4 до DeepSeek — строят текст последовательно, по одному токену. Но за последние два года исследователи активно развивают альтернативу: masked diffusion language models (MDM), которые генерируют текст параллельно, через множество шагов уточнения. Проблема в том, что MDM до сих пор уступали autoregressive моделям в эффективности обучения и качестве рассуждений. Новая работа LoopMDM ломает это ограничение — не добавляя параметров, а зацикливая уже существующие слои.

Что такое LoopMDM

LoopMDM — это архитектурный трюк, который берёт стандартный диффузионный трансформер и зацикливает его средние слои. Вместо того чтобы прогонять данные через 12, 18 или 24 уникальных слоя один раз, модель применяет небольшой блок из 2 средних слоёв многократно, переиспользуя одни и те же веса. Голова (head) и хвост (tail) остаются незацикленными: head готовит эмбеддинги, tail проецирует результат обратно в пространство токенов, а mid-block работает как итеративный рефайнер.

Ключевое отличие от простого увеличения глубины — количество параметров не растёт. 12-слойная LoopMDM с двумя зацикленными слоями имеет ровно столько же параметров, сколько обычная 12-слойная модель. Разница только в том, что эти два слоя прогоняются S раз за один шаг денойзинга. При S=12 эффективная глубина достигает 22 слоёв, но параметров всё ещё как у 12-слойной модели.

Авторы показывают, что этот подход даёт два независимых рычага управления. Во время обучения зацикливание создаёт эффект масштабирования глубины без добавления параметров — модель учится глубже обрабатывать каждый шаг диффузии. Во время инференса можно менять количество циклов S на лету: больше циклов — выше качество, меньше циклов — быстрее генерация. Это делает LoopMDM первой диффузионной языковой моделью, которая масштабирует вычисления при инференсе так же легко, как это делают chain-of-thought промпты у autoregressive моделей.

Почему зацикливание работает лучше увеличения глубины

Интуиция простая: если у вас есть 12 слоёв, и вы хотите большей выразительности, есть два пути. Первый — добавить ещё 6-10 слоёв с новыми весами. Это работает, но каждый новый слой требует градиентов, памяти и времени на обучение. Второй — взять существующие слои и применить их несколько раз, позволяя модели итеративно уточнять представления. LoopMDM показывает, что второй путь эффективнее.

В экспериментах на трёх корпусах — OpenWebText, LM1B и FineWeb-Edu — LoopMDM достигает той же перплексии, что и базовая MDM, но с до 3.3× меньшими затратами на обучение. При этом на downstream-задачах зацикленная модель не просто не уступает, а систематически превосходит базовый вариант. На девяти бенчмарках commonsense-рассуждений LoopMDM с S=12 улучшает результаты на всех датасетах, с самыми заметными приростами на BoolQ (+10.8 points) и RACE (+3.2 points).

Ещё более убедительное сравнение — с более глубокими моделями. Авторы обучили 21-слойную MDM с сопоставимыми вычислениями на шаг (per-step FLOPs). LoopMDM с 14 слоями и зацикливанием победила её на GSM8K — бенчмарке математических рассуждений — с отрывом +2.7 points при S=8 и +3.4 points при S=16. Это означает, что зацикливание даёт не просто экономию параметров, а качественно другую динамику обучения: модель учится лучше рассуждать, потому что итеративное применение слоёв заставляет её пересматривать промежуточные гипотезы.

Математические рассуждения: +8.5 points на GSM8K

GSM8K — золотой стандарт для оценки математических способностей языковых моделей. Это набор из 8.5 тысяч школьных задач по арифметике, где правильный ответ требует не просто вычислений, а пошагового логического вывода. Autoregressive модели традиционно доминировали здесь, потому что их последовательная генерация естественным образом подходит для цепочек рассуждений. Диффузионные модели, генерирующие текст параллельно, historically отставали.

LoopMDM меняет эту картину. На FineWeb-Edu чекпоинтах с сопоставимыми обучающими вычислениями зацикленная модель показывает монотонный рост точности с увеличением числа циклов. При S=16 — значении, превышающем максимальное количество циклов во время обучения (S_max=8) — прирост достигает +8.5 points на GSM8K по сравнению с 14-слойной базовой моделью. Это говорит о двух вещах: во-первых, зацикливание действительно улучшает способность к многошаговым рассуждениям, а во-вторых, модель обобщается за пределы обучающего режима — она умеет использовать больше циклов, чем видела при тренировке.

Интересно, что при S=1 LoopMDM немного уступает базовой модели. Это ожидаемо: один прогон зацикленных слоёв эквивалентен поверхностной обработке, и преимущество зацикливания раскрывается только при повторном применении. Но как только S достигает 6 и выше, зацикленная модель уверенно выходит вперёд и продолжает расти даже при S=24 — вдвое превышая обучающий максимум.

Sudoku как лаборатория для глобального планирования

Чтобы понять, почему зацикливание помогает, авторы придумали хитрый эксперимент: заставили модель решать Sudoku с фиксированным порядком заполнения клеток слева направо. Обычно MDM решают Sudoku, используя адаптивный порядок демаскирования — сначала заполняют самые лёгкие клетки. Это даёт хорошие результаты, но не показывает, может ли модель действительно планировать глобально.

Авторы убрали это преимущество: клетки заполняются строго по порядку, независимо от сложности. При таком ограничении любое улучшение может приходить только от внутришаговых вычислений — то есть от зацикливания. Результат оказался драматичным: при S=1 точность решения составила всего 10.6%, почти как у однослойной модели (10.9%). Но при S=2 количество ошибок упало с 24 до 7, а при S=3 модель решала головоломки почти идеально. Это прямое доказательство того, что зацикливание позволяет маскированным позициям взаимодействовать и корректировать глобально несогласованные предсказания до фиксации токена.

Адаптивное зацикливание: умное распределение вычислений

Не все шаги диффузии одинаково сложны. В начале процесса модель видит почти полностью замаскированную последовательность и должна принять грубые стратегические решения. В конце — остаётся только несколько неопределённых позиций, и работы мало. Авторы обнаружили, что польза от дополнительных циклов неравномерна: она пиковая на промежуточных шагах, где у модели достаточно контекста, но ещё остались нерешённые позиции, и минимальна на краях процесса.

Это наблюдение привело к простой, но эффективной стратегии: адаптивное зацикливание. Вместо фиксированного S на всех шагах модель динамически выбирает количество циклов в зависимости от текущего timestep. На практике это снижает среднее число циклов с 12 до примерно 5, сохраняя при этом большую часть качества. Для продакшена это означает, что можно получить 80% прироста за 40% вычислений — без каких-либо изменений в архитектуре или обучении.

Механизм: маскированные позиции начинают разговаривать друг с другом

Чтобы понять, что происходит внутри зацикленных слоёв, авторы проанализировали паттерны внимания. Они измерили долю внимания, которую маскированные позиции уделяют другим маскированным позициям (mask-to-mask attention), при разных значениях S. Результат оказался однозначным: с ростом числа циклов маскированные позиции всё активнее взаимодействуют между собой.

Это важно, потому что в стандартной MDM маскированные позиции в основном «смотрят» на уже раскрытые токены и пытаются угадать себя по контексту. Зацикливание даёт им время для диалога: каждая замаскированная позиция может учитывать предположения соседей, корректировать собственные гипотезы и сходиться к глобально согласованному решению. Это похоже на то, как люди решают сложные задачи — не мгновенно, а через итеративное уточнение, где каждый шаг учитывает промежуточные результаты других мыслительных линий.

Теоретически это соответствует концепции «вычислительного рабочего пространства» (computational workspace), которая недавно появилась в литературе о padded looped transformers. Маскированные позиции в MDM играют роль параллельного хранилища промежуточной информации, а зацикливание увеличивает эффективную глубину обработки без роста ширины. В логарифмической шкале ширины MDM и зацикленные трансформеры оказываются тесно связанными вычислительными моделями — они оба выполняют итеративное параллельное уточнение, просто в разных представлениях.

Что это значит для будущего языковых моделей

LoopMDM вписывается в более широкий тренд: исследователи ищут способы масштабировать вычисления при инференсе независимо от масштабирования параметров. У autoregressive моделей это давно работает через chain-of-thought: длиннее рассуждение — лучше результат. У диффузионных моделей такого рычага не было — до LoopMDM.

Теперь у MDM появляется аналогичный механизм. Нужен быстрый черновик — ставим S=4. Нужен тщательно выверенный ответ на сложный вопрос — ставим S=16. Это открывает путь к адаптивным системам, которые динамически распределяют вычисления в зависимости от сложности запроса, точно так же, как человек тратит больше времени на сложную задачу и меньше на простую.

Ещё один важный аспект — эффективность обучения. 3.3× меньше FLOP при сопоставимом качестве означает, что диффузионные модели могут стать конкурентоспособными при значительно меньших бюджетах. Для академических лабораторий и стартапов это может снизить барьер входа в разработку собственных foundation models. А для крупных компаний — открыть возможность обучать больше моделей, перебирать архитектуры и экспериментировать с меньшими затратами.

Важно понимать, что LoopMDM — не замена autoregressive моделям, а дополнение. Она показывает, что в парадигме диффузионного языкового моделирования есть ещё нераскрытый потенциал. Если зацикливание двух слоёв даёт +8.5 points на GSM8K при 170M параметрах, что будет при масштабировании до 7B, 70B или 400B параметров? И что произойдёт, если комбинировать зацикливание с другими недавними прорывами — например, с непрерывными эмбеддингами из ELF или с адаптивным демаскированием? Поле для экспериментов огромно.

Часто задаваемые вопросы

Чем LoopMDM отличается от обычного увеличения глубины трансформера?

Обычное увеличение глубины добавляет новые слои с новыми весами, что растёт параметры и затраты на обучение. LoopMDM переиспользует существующие веса, создавая эффективную глубину без роста параметров. Эксперименты показывают, что зацикливание средних слоёв эффективнее простого добавления слоёв: LoopMDM побеждает 21-слойную модель при сопоставимых per-step FLOP.

Можно ли применить зацикливание к autoregressive моделям вроде GPT?

Идея зацикливания слоёв не нова — она известна как Universal Transformers и применялась к autoregressive моделям. Но в контексте masked diffusion зацикливание работает иначе: маскированные позиции получают возможность итеративно уточнять предсказания до фиксации токена, чего нет в autoregressive генерации, где каждый токен фиксируется навсегда.

Как адаптивное зацикливание влияет на скорость генерации?

Адаптивное зацикливание снижает среднее число циклов с 12 до ~5, сохраняя большую часть качества. Это означает, что практическая скорость генерации может быть выше, чем у неадаптивной модели с фиксированным S=12, при сопоставимом качестве. Точный trade-off зависит от задачи и аппаратной реализации.

Итог

LoopMDM доказывает, что в диффузионных языковых моделях важнее не ширина, а глубина мысли — буквально. Зацикливая средние слои трансформера, авторы получили модель, которая тренируется в 3 раза дешевле, решает математические задачи на 8.5 points лучше и масштабирует вычисления при инференсе так же гибко, как chain-of-thought у autoregressive моделей. Это не просто архитектурный трюк — это новый способ думать о том, как языковые модели обрабатывают информацию: не за один проход, а через итеративное уточнение, где каждый цикл делает представление чуть более согласованным, чуть более осмысленным, чуть более человечным.

← Все записи