Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Декодирование Looped Transformers удалось ускорить, начав вычислять будущие токены до окончательной проверки текущего. В препринте Seoul National University и KAIST, который не прошёл рецензирование и содержит замеры самих авторов, пиковое ускорение достигло 6,83 раза при сохранении распределения токенов целевой модели. Метод не требует переобучения или отдельной модели для черновых предсказаний.
Повторные проходы дают черновик почти бесплатно
Looped Transformers хранят небольшой набор блоков Transformer и применяют его к скрытому состоянию несколько раз. Так модель увеличивает вычислительную глубину без пропорционального роста числа параметров, но при генерации каждого токена снова загружает одни и те же веса.
Из-за этого декодирование упирается в пропускную способность памяти. Обычная схема завершает все повторные проходы для текущего токена и только затем начинает следующий, поэтому GPU ждёт последовательную цепочку вычислений.
Промежуточное состояние уже содержит распределение следующего токена, хотя модель ещё не закончила проходы. LoopSpec берёт из раннего состояния черновой токен, добавляет его к текущему префиксу и запускает вычисление следующей позиции. Основная ветка тем временем продолжает проходы и получает окончательное распределение.
Ветки находятся на разной рекуррентной глубине, то есть прошли разное число повторов одного блока. LoopSpec обрабатывает их одним пакетным вызовом. Когда ранний черновик принят, часть работы для следующего токена уже выполнена, и последовательная задержка сокращается.
Вторая ветка не даёт конвейеру остановиться
Один ранний черновик создаёт уязвимое место: если итоговое распределение его отвергает, все вычисления зависимой ветки приходится удалить. Декодирование возвращается к подтверждённому префиксу и снова проходит полную рекуррентную глубину.
LoopSpec готовит запасной вариант из более глубокого состояния. Он строится не из всего нового распределения, а из его остатка после первого предложения. Поэтому вторая ветка ищет токен среди вариантов, которые должны заменить отвергнутый первый черновик, а не дублирует его.
Запасная ветка появляется только тогда, когда более глубокое состояние снизило уверенность в первом токене. Этот затвор сдерживает рост числа параллельных веток и оставляет дополнительное вычисление для случаев, где оно вероятнее пригодится.
При окончательной проверке система сначала пробует принять первый токен, затем запасной и лишь после двух отказов перезапускает конвейер. В экспериментах с базовыми моделями полный перезапуск потребовался не более чем для 2,95% подтверждённых токенов.
Для жадного декодирования проверка сохраняет тот же итоговый токен, который выбрала бы исходная модель. При случайной выборке отдельная последовательность может отличаться, но её вероятность остаётся такой же, как при обычном декодировании. Ускорение поэтому не означает, что модель стала точнее: LoopSpec меняет расписание вычислений, а не её ответы как распределение.
Глубину обеих черновых веток не обязательно подбирать полным перебором. Расстояние между промежуточным и окончательным распределениями в замерах убывало по степенному закону; на этой основе работа выводит подходящие глубины аналитически, и расчёт совпал с измеренным оптимумом.
Метод меняет планы только для рекуррентных моделей
Проверка охватила семь версий моделей из двух семейств: Ouro и Raven. В неё вошли задачи на общие рассуждения, математику и код из GSM8K, MATH-500, BBH, HumanEval+, MBPP+ и AIME. Скорость измеряли относительно обычного авторегрессионного декодирования в реализации на SGLang.
У Ouro выигрыш достиг 3,36 раза, а более глубокие Raven получили больший эффект, указанный в лиде. На моделях с рассуждением ускорение сохранилось и при случайной выборке токенов, где оно доходило до 2,83 раза. LoopSpec также обошёл требующий обучения метод DFlash по времени декодирования до 1,7 раза.
Сокращение числа последовательных повторов не превращается в такое же сокращение реального времени автоматически. У Raven заметную долю затрат создают слои до и после повторяющегося блока, а ранние черновики чаще запускают эти слои и новые ветки. Поэтому конфигурация с наименьшим числом последовательных шагов может проиграть другой конфигурации по задержке.
Для команды, которая уже разворачивает Looped Transformer, работа добавляет вариант оптимизации без нового обучения и отдельного чернового чекпойнта. Цена внедрения переносится в сервер вывода: ему нужно вести кэш ключей и значений отдельно для каждой ветки, удалять отвергнутые продолжения и собирать вычисления через CUDA Graph.
Планы вокруг обычного Transformer метод напрямую не меняет: LoopSpec зависит от повторного применения общих блоков и пригодных для чтения промежуточных состояний. При выборе рекуррентной архитектуры потенциальный выигрыш в задержке теперь стоит проверять вместе с расходом памяти на ветки и стоимостью слоёв вне цикла, а не оценивать модель только по числу параметров.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
