AQUA-KV: адаптивная квантизация KV-кэша

На ICML 2025 команда Yandex Research представила шесть статей (каких именно — читайте в одном из предыдущих постов) — среди них есть работа, посвящённая методу адаптивной квантизации KV-кэша. Один из авторов, исследователь Yandex Research Алина Шутова, рассказала нашему каналу, в чём суть предложенного в публикации способа.

Одна из ключевых проблем эксплуатации LLM — экспоненциальный рост потребления памяти графических процессоров при обработке длинных контекстов. Это связано с необходимостью хранения KV-кэша. Для современных моделей, таких как Llama 3.2 70B, и контекстов в 131 тысячу токенов, объём KV-кэша может достигать 42,9 ГБ на последовательность, что существенно ограничивает практическое применение и увеличивает стоимость вычислений. Традиционные методы сжатия, такие как примитивное квантование или прунинг, демонстрируют значительную деградацию качества генерации при агрессивных режимах сжатия, особенно в области 2–3 бит на значение.

Предложенный авторами статьи метод AQUA-KV (Adaptive QUAntization for Key-Value) представляет принципиально новый подход, основанный на фундаментальном наблюдении: векторы ключей и значений в соседних слоях трансформера обладают высокой степенью корреляции. Эта структурная избыточность позволяет прогнозировать значительную часть информации слоя k+1 на основе данных слоя k.

Вместо независимого квантования каждого слоя AQUA-KV использует обученные линейные предикторы. Один предиктор предсказывает ключи слоя k+1 на основе ключей слоя k, другой предсказывает значения слоя k+1 по комбинации предсказанных ключей этого слоя и значений слоя k. Обучение этих компактных адаптеров проводится в ходе одноразовой калибровки на целевой модели.

Критический шаг метода — переход от квантования векторов целиком к квантованию только остаточной информации, то есть разности между фактическими векторами слоя и их предсказаниями. Поскольку остаток содержит лишь ту информацию, которую нельзя получить из предыдущего слоя, его информационная энтропия существенно ниже. Эта остаточная компонента подвергается экстремальному квантованию (до 2–2,5 бит на элемент) с применением векторного квантования без данных (data-free VQ), адаптивно оптимизирующего распределение битов под дисперсию остатков. Для восстановления KV-векторов во время инференса используются те же предикторы и деквантованный остаток.

Эксперименты демонстрируют эффективность AQUA-KV. На моделях семейств Llama 3.2 и Qwen 2.5 применение метода с квантованием до 2 бит на значение привело к снижению объёма памяти KV-кэша в 16 раз (с ~43 ГБ до ~2,7 ГБ для контекста в 131K токенов) при сохранении практически исходного качества генерации. Относительное увеличение перплексии составило менее 1%, а деградация точности на задачах длинного контекста из бенчмарка LongBench не превысила 1%. AQUA-KV совместим с любыми методами квантизации, и, как продемонстрировано в работе, заметно улучшает качество всех рассмотренных методов. Метод демонстрирует совместимость с техниками прунинга, такими как H2O, обеспечивая дополнительную экономию памяти. Код AQUA-KV можно найти на GitHub.

ML Underhood

#YaICML25

Как в Яндекс Погоде использовали фичу для таргета, чтобы улучшить точность предсказания осадков

Яндекс запустил новые модели краткосрочного прогноза осадков — подробнее о них можно почитать на Хабре. А специально для нашего канала ML-разработчик в Яндекс Погоде Дмитрий Стефеев рассказал о важной идее, которая позволила обойти возникшие в начале работы трудности.

Цель — создать модель регрессии для прогноза выпавших за час осадков (в мм). Основными фичами являются прогнозы глобальных численных моделей (NWP).

Главная проблема заключалась в том, что радарные данные, которые мы используем в качестве таргета для обучения, значительно отличаются от прогноза численных моделей на входе. Всё из-за сложности прогноза осадков. Очень легко промазать как по времени, так и пространству. Все функции потерь и их комбинации, которые мы пробовали, приводили к тому, что модель прогнозировала значительно меньше осадков, чем нужно, боясь промазать. Можно было бы попробовать перейти от регрессии к классификации и подбирать пороги, но их, скорее всего, пришлось бы подбирать для каждого шага прогноза (как делали, к примеру, в Metnet от Google) и, вероятно, они бы отличались для разных мест на Земле и разных сезонов. Это значительно усложнило бы поддержку модели и увеличило вероятность непредсказуемого поведения.

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

В ходе экспериментов мы пришли к простой формуле для весов. Сначала считаем медиану для прогнозов осадков всех численных моделей в каждой точке прогноза. Далее считаем среднюю абсолютную разность прогнозов численных моделей и медианы, а затем делим полученное значение на медиану, клипая результат по подобранным границам. Полученное значение W мы применяем для лосса к радарам, а для лосса к прогнозу численной модели используем вес 1 - W. И такой трюк позволяет полностью устранить проблему с уменьшением объёма осадков — по сути это можно считать регуляризацией выхода модели к численному прогнозу, так как мы привязываем прогноз к физической модели, не давая ему отойти слишком далеко.

Далее мы добавили домножение полученных весов на дополнительные веса, которые линейно зависят от шага прогноза, чтобы по мере увеличения шага снижать вес лосса к радарам и увеличивать к численному прогнозу до подобранного граничного значения.

На изображении выше в первой строке показан пример радарного таргета по первым 12 часам прогноза. Ниже — прогноз нашей модели и прогноз численной модели в дополнительном таргете. Для удобства отображения значения мм приведены к трём классам осадков. Видно, что на первых шагах, прогноз модели больше похож на радары, а далее становится ближе к прогнозу численной модели, что позволяет улучшать качество прогноза на первых шагах за счёт исторических радарных данных, если они доступны на входе модели.

ML Underhood