Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Из вывода Transformer с повторно используемыми блоками научились убирать вычисления, которые почти не меняют результат. В препринте Wanqi Yang и Shiwei Liu, который не рецензирован и где все числа получили сами авторы, FlashLoop ускорил полный цикл вывода до 1,64 раза и сократил KV-кеш до 6 раз. Метод работает поверх готовой модели без дополнительного обучения, поэтому его можно проверить, не меняя веса.
Поздние циклы меняют лишь часть состояния
Looped Transformers многократно применяют один и тот же набор блоков. Так модель увеличивает вычислительную глубину, не сохраняя отдельные параметры для каждого слоя, но каждый новый цикл снова обрабатывает контекст и создаёт собственный KV-кеш — сохранённые ключи и значения механизма внимания.
Меньше параметров поэтому не означает более дешёвый вывод. На контексте длиной 32K исходная Ouro обрабатывала вход за 27 секунд, тогда как сопоставленная обычная LLaMA — за 3 секунды. Причина не только в арифметике: новые наборы ключей и значений также приходится хранить и читать из памяти GPU.
Работа показывает, что изменения между соседними циклами распределены неравномерно. По мере повторений большинство токенов меняется всё слабее, а заметные обновления остаются у небольшой группы. Эта группа постепенно сужается: токены, активные в позднем цикле, обычно уже были активны в предыдущем.
Похожая картина возникает внутри внимания. Разницу между результатами соседних циклов определяет небольшая и достаточно стабильная группа ключей. Для Ouro обновление 10% столбцов внимания в позднем цикле восстанавливало более 90% точного результата внимания.
Наконец, соседние наборы ключей и значений отличаются слабее, чем полные состояния. Поэтому выгоднее сжимать не каждый KV-кеш отдельно, а базовое состояние и небольшие поправки к нему.
FlashLoop объединяет три вида экономии
Сначала FlashLoop выполняет ранние циклы полностью. Затем он ранжирует токены по тому, насколько изменились их скрытые состояния, и пересчитывает только активные. Остальные токены повторно используют состояния и записи KV-кеша из предыдущего цикла.
Во время внимания метод выбирает ключи по статистике предыдущего цикла. Простое удаление остальных ключей исказило бы нормировку: оставшиеся вероятности сложились бы в единицу и получили слишком большой вес. FlashLoop сохраняет общую массу вероятности выбранной группы и использует её при пересчёте, чтобы приблизить результат полного внимания.
KV-кеш хранится как сжатая база и последовательность сжатых поправок. Для оставшихся активными токенов FlashLoop записывает низкоразрядную разницу с предыдущим состоянием, а для неактивных не создаёт новую запись. В основных опытах поправки сжимали до 4 бит; более агрессивный вариант уже заметно ухудшал качество.
Разреженность сама по себе не гарантирует ускорения на GPU: выбор отдельных токенов создаёт нерегулярные чтения памяти и много мелких операций. Реализация собирает выбранные данные в плотные буферы, а обновление и сжатие KV-кеша объединяет внутри вычислительного ядра. Это превращает сокращение операций в выигрыш по времени, а не только в меньшую теоретическую оценку вычислений.
Метод меняет планы только для моделей с повторяющимися блоками
FlashLoop проверяли на пяти моделях семейств Ouro и Huginn и на пяти наборах задач, включая MATH-500, GSM8K и ARC-Challenge. Средние результаты оставались близки к исходным моделям, хотя варианты с усиленным пошаговым рассуждением оказались чувствительнее к сжатию. Замеры полного времени выполняли на одном GPU NVIDIA A100, поэтому перенос ускорения на другие ускорители и серверные конфигурации требует отдельной проверки.
Для команды, которая уже использует Looped Transformer или рассматривает большую глубину циклов и длинный контекст, работа меняет оценку стоимости вывода. Линейный рост вычислений и KV-кеша больше не обязательно принимать как свойство архитектуры: часть этого роста можно убрать после обучения модели. Практический прототип при этом потребует специальных GPU-ядер и настройки доли обновляемых токенов и ключей.
Для обычного Transformer оснований менять архитектуру только ради FlashLoop нет. Метод использует устойчивость состояний именно между повторениями общего блока и не доказывает, что Looped Transformers в целом выгоднее обычных моделей. Он скорее снимает один из главных инженерных барьеров у команд, которые уже выбрали повторные циклы ради меньшего числа параметров или управляемой вычислительной глубины.
Источники
Иллюстрация: рисунок из статьи «FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates», Wanqi Yang, Shiwei Liu, CC BY 4.0
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
