Новый метод функциональной реконструкции ускоряет инференс LLM без переобучения

Группа исследователей опубликовала на arXiv препринт, посвященный повышению эффективности спекулятивного декодирования больших языковых моделей (LLM). Предложенный метод Functional Reconstruction решает проблему потери точности при конвертации моделей с multi-head attention (MHA) в multi-head latent attention (MLA), которая используется для снижения нагрузки на память при длинном контексте.

MLA становится все более востребованной для инференса длинных контекстов, так как компактные латентные состояния заменяют растущий кэш ключей и значений (KV cache), уменьшая трафик памяти. Однако большинство открытых чекпоинтов, таких как Llama и Qwen, по-прежнему используют MHA или групповое внимание (GQA). Для применения MLA требуется конвертация, но прямое преобразование может значительно снизить согласованность между черновой моделью (draft) и целевой моделью (target) в спекулятивном декодировании.

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

Авторы предлагают рассматривать построение черновой модели как функциональную реконструкцию, а не просто сжатие кэша. Их сквозной метод оптимизирует каждый конвертированный MLA-модуль, чтобы воспроизводить реакцию исходного MHA/GQA на калибровочных скрытых состояниях. Процедура не зависит от конкретного конвертера, не требует логитов верификатора и работает без дополнительного обучения с учителем.

Эксперименты проводились на четырех парах моделей Llama/Qwen, двух конвертерах (TransMLA и MHA2MLA), двух бэкендах (HF и vLLM) и четырех задачах по 200 промптов — всего 192 конфигурации. С допуском отчетности 0.5 процентного пункта Functional Reconstruction заметно улучшил принятие токенов в 37 из 64 сопоставимых ячеек, оставил 26 без изменений и ухудшил только одну.

Код и артефакты оценки доступны на GitHub по ссылке, указанной в препринте. Это позволяет другим исследователям воспроизвести результаты и интегрировать метод в свои пайплайны. Разработка может ускорить внедрение MLA в существующие открытые модели без необходимости их переобучения с нуля.