Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Кэш языковой модели с повторными проходами через одни и те же слои удалось сократить без заметного провала в качестве. Команда из Virginia Tech и независимый исследователь описали HLA в препринте, который не прошёл рецензирование и содержит их собственные замеры: метод уменьшил кэш на порядок и ускорил пакетную генерацию. Для таких архитектур память GPU больше не обязательно сводит на нет экономию параметров.
Как HLA хранит историю без восстановления ключей и значений
Модель с повторными проходами применяет один стек слоёв к каждому токену несколько раз. Она получает большую вычислительную глубину без новых весов, но каждый проход создаёт свои ключи и значения для механизма внимания. Кэш ключей и значений (KV-кэш) поэтому растёт вместе с числом проходов.
Этот расход особенно мешает при генерации. Кэш занимает память, ограничивает число одновременно обрабатываемых последовательностей и целиком читается на каждом шаге. В результате небольшая по числу параметров модель может требовать много памяти именно во время эксплуатации.
HLA делит историю на две части. Для последних 128 токенов система сохраняет точные ключи и значения каждого прохода. Более старый токен записывается как компактное латентное представление, общее для нескольких проходов; первый проход получает отдельное представление меньшего размера, поскольку его данные хуже предсказываются по поздним состояниям.
Главное отличие от предыдущего подхода LLA состоит в чтении этого представления. LLA на каждом шаге восстанавливает из него ключи и значения всех проходов, поэтому экономит память, но добавляет много вычислений. HLA переводит запрос текущего токена прямо в латентное пространство, получает взвешенную сумму сохранённых представлений и переводит результат обратно. Точные и сжатые части участвуют в одной нормализации весов внимания.
Метод добавили к готовым Ouro без изменения исходных весов. Новые линейные преобразования сначала настроили методом главных компонент, а затем обучили повторять распределение внимания и выходы исходной модели. При дальнейшем обучении под задачу используются два прохода: первый записывает латентные представления, второй считает результат по правилам, близким к последовательной генерации.
Память сократилась сильнее, чем качество
На обеих моделях HLA уменьшила постоянный расход KV-кэша на токен в 10,7 раза. С учётом точного окна на GPU помещалось в 4,0–8,8 раза больше одновременных последовательностей, в зависимости от длины контекста. Пиковая скорость пакетной генерации выросла в 2,5–7,4 раза.
На MATH500, MMLU-Pro и BBH модель сохранила свыше 97% исходного результата. На задачах поиска информации RULER осталось не менее 96% результата полной версии. После обучения на конкурсной математике варианты с HLA и полным кэшем показали сопоставимую точность, а небольшая разница между ними оказалась в пределах погрешности.
Точное окно оказалось частью метода, а не необязательной оптимизацией. Ближайшие токены дают непропорционально большую долю ошибки латентного внимания, поэтому их сжатие сильнее портит результат. HLA тратит фиксированный объём памяти на недавнюю историю и сжимает ту часть контекста, где приближение работает лучше.
Когда работа меняет план развёртывания
HLA имеет смысл учитывать, если команда уже рассматривает архитектуру с повторным использованием слоёв и упирается в число параллельных генераций. Метод снимает основной инфраструктурный недостаток такой схемы: экономия весов больше не сопровождается пропорциональным ростом кэша для каждого прохода.
Для обычного Transformer это не готовая замена существующего KV-кэша. HLA использует структуру состояний между повторными проходами, добавляет новые обучаемые преобразования и требует дообучения конкретной модели. Перенести полученные коэффициенты на другую архитектуру без настройки нельзя.
Выигрыш также зависит от нагрузки. Обработка входного запроса не ускоряется: она использует исходный точный механизм внимания. При генерации одной последовательности с коротким контекстом HLA была на 26% медленнее полной версии из-за постоянных расходов на преобразования; сравнялись варианты примерно у 4 тысяч токенов. Рост скорости на коротких контекстах появляется при пакетной обработке, когда меньший кэш позволяет загрузить на GPU больше запросов.
Проверка охватывает семейство Ouro: модели на 1,4 и 2,6 млрд параметров с четырьмя проходами, GPU A100 и контексты до 16 тысяч токенов. Дальнейшее обучение проверяли только на математических задачах, причём один шаг требует двух прямых проходов вместо одного. Поэтому работа меняет расчёт инфраструктуры для Ouro-подобных моделей, но пока не доказывает тот же эффект для других семейств, большего числа проходов или более длинных контекстов.
Практический вывод зависит от профиля сервиса. Для одиночных коротких ответов HLA добавляет сложность без гарантированного выигрыша. Для длинной генерации и высокой параллельности она превращает повторное использование слоёв из экономии только на весах в экономию памяти при эксплуатации.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
