Efficient Parallel Samplers for Recurrent-Depth Models and Their Connection to Diffusion Language Models
Geiping, Yang, Su
Language models with recurrent depth, also referred to as universal or looped when considering transformers, are defined by the capacity to increase their computation through the repetition of layers. Recent efforts in pretraining have demonstrated that these architectures can scale to modern language modeling tasks while exhibiting advantages in reasoning tasks. In this work, we examine the relationship between recurrent-depth models and diffusion language models. Building on their similarities, we develop a new diffusion forcing sampler for these models to accelerate generation. The sampler advances by decoding new tokens at every forward pass of the model, while the latent states of these tokens can be further refined in parallel through recurrence. Theoretically, generation with our sampler is strictly more expressive than the baseline autoregressive generation using the same time budget on modern hardware. Moreover, this sampler, based on principles from diffusion literature, can be directly applied to existing 3.5B recurrent-depth transformers without any tuning, leading to up to a 5x speedup. Consequently, our findings not only provide an efficient mechanism for parallelizing the extra computation in recurrent-depth models at inference, but also suggest that such models can be naturally viewed as strong continuous, though causal, diffusion language models.
academic
Эффективные параллельные семплеры для моделей с рекуррентной глубиной и их связь с диффузионными языковыми моделями
В данной работе исследуется связь между языковыми моделями с рекуррентной глубиной (также называемыми универсальными трансформерами или циклическими трансформерами) и диффузионными языковыми моделями. Модели с рекуррентной глубиной увеличивают объём вычислений путём повторения слоёв и показывают преимущества в задачах логического вывода. На основе сходства двух классов моделей авторы разработали новый семплер диффузионного принуждения для ускорения процесса генерации. Семплер декодирует новые токены при каждом прямом проходе, одновременно оптимизируя потенциальные состояния этих токенов посредством рекуррентного параллелизма. Теоретически семплер обладает большей выразительностью, чем базовая авторегрессивная генерация при одинаковом временном бюджете. Что ещё более важно, семплер может быть непосредственно применён к существующему трансформеру с рекуррентной глубиной с параметрами 3,5B без какой-либо настройки, достигая ускорения до 5 раз.
Традиционные большие языковые модели используют архитектуру нейронной сети фиксированной глубины с относительно небольшим количеством слоёв (обычно двузначное число). Хотя такой дизайн хорошо работает с точки зрения эффективности обучения и производительности на большинстве задач, он имеет ограничения при решении сложных задач, требующих многошагового логического вывода, таких как математика и программирование. С точки зрения теории сложности трансформеры фиксированной глубины принадлежат классу сложности TC0, что ограничивает их выразительную способность.
Ограничения вычислительной способности: модели фиксированной глубины с трудом справляются с многошаговыми логическими цепочками, требующими концептуальных скачков
Проблемы эффективности логического вывода: хотя модели с рекуррентной глубиной обладают большей выразительностью, они генерируют медленно, так как каждая рекуррентная итерация должна выполняться последовательно
Требование параллелизации: современная архитектура GPU предоставляет возможности для параллельных вычислений, но традиционная авторегрессивная генерация не может полностью их использовать
Методы цепочки мышления: требуют экстернализации внутреннего процесса логического вывода в виде небольших шагов, что увеличивает длину последовательности
Модели с рекуррентной глубиной: хотя обладают большей выразительностью, каждый рекуррентный шаг при логическом выводе должен выполняться последовательно, что приводит к медленной генерации
Традиционные методы параллелизации: такие как спекулятивное декодирование, в основном разработаны для моделей фиксированной глубины
Теоретический вклад: уточнение связи между моделями с рекуррентной глубиной и диффузионными моделями путём установления теоретического моста через диффузионное принуждение и стратегии логического вывода на основе блоков или волн
Методологическое инновация: предложение семплера диффузионного принуждения, применимого к моделям с рекуррентной глубиной, реализующего параллелизацию процесса логического вывода
Экспериментальная верификация: проверка эффективности метода на модели Huginn-0125 с параметрами 3,5B, достижение ускорения примерно в 5 раз на эталонных тестах GSM8K, MATH500, HumanEval и MBPP при сохранении сопоставимой точности
Практическая ценность: семплер может быть непосредственно применён к существующим моделям с рекуррентной глубиной без переобучения или настройки
Учитывая модель с рекуррентной глубиной и входную подсказку x, целью является ускорение процесса генерации текста при сохранении качества генерации. Конкретно, необходимо либо сгенерировать больше токенов при одинаковом временном бюджете, либо сократить время генерации при одинаковом количестве токенов.
Рекуррентный процесс обусловлен встроенным входом e, позволяя семплеру выполнять "коррекцию пути" при изменении условий без необходимости отбрасывать частично вычисленные состояния.
Различные рекуррентные глубины могут совместно использовать кэш KV, значительно снижая использование памяти. Эксперименты показывают, что модель естественным образом поддерживает совместное использование кэша KV, требуя только хранения состояния KV последней рекуррентной итерации для каждой позиции токена.
Ограничение пакетной обработки: текущая реализация поддерживает только логический вывод одной последовательности, пакетные сценарии требуют сложного механизма логического вывода
Эффективность FLOP: хотя повышается степень параллелизма, использование FLOP фактически увеличивается
Зависимость от оборудования: оптимальные параметры зависят от конкретной конфигурации оборудования
Требования к модели: требуется, чтобы модель удовлетворяла определённым требованиям архитектуры (инъекция входных данных, надёжная рекуррентность и т. д.)
Сильная инновационность: первое установление связи между моделями с рекуррентной глубиной и диффузионными моделями, предоставляющее новую теоретическую перспективу
Высокая практическая ценность: достижение значительного ускорения логического вывода, применимого непосредственно к существующим моделям
Строгая теория: предоставление теоретического анализа масштабирования глубины vs ширины и доказательства сходимости
Полные эксперименты: верификация эффективности и надёжности метода на нескольких эталонных тестах и вариантах моделей
Статья ссылается на богатый объём связанных работ, включая:
Dehghani et al. (2019): исходная работа Universal Transformers
Chen et al. (2024a): метод Diffusion Forcing
Geiping et al. (2025): модель Huginn-0125 с рекуррентной глубиной
Rombach et al. (2022): диффузионные модели в пространстве потенциалов
Leviathan et al. (2023): метод спекулятивного декодирования
Общая оценка: Это высококачественная исследовательская работа, имеющая важные вклады как в теоретическую инновацию, так и в практическую ценность. Статья успешно устанавливает связь между двумя важными классами моделей и предлагает практический метод ускорения. Хотя существуют определённые ограничения, работа предоставляет ценные направления и основу для будущих исследований.