Google сократила потребление памяти ИИ-моделями в шесть раз без потери точности — с алгоритмом TurboQuant

Исследовательский отдел Google Research опубликовал работу о технологии TurboQuant — алгоритме квантизации, который сокращает разрядность KV-кеша больших языковых моделей до 3 битов без потери точности в ответах и без потребности в дополнительном обучении. В тестах на ускорителях искусственного интеллекта Nvidia H100 4-битный алгоритм TurboQuant (с четвёртым битом на коррекцию ошибок) помог восьмикратно повысить производительность при вычислении логитов внимания в сравнении с неквантованными 32-битными ключами; при этом объём KV-кеша сократился в шесть раз.

В KV-кеше хранятся ранее вычисленные данные механизма внимания, чтобы модели не приходилось пересчитывать их на каждом шаге генерации токенов. По мере роста контекстного окна этот кеш существенно разрастается; традиционные методы квантования помогают уменьшить размер кеша, но за это приходится «расплачиваться» несколькими дополнительными битами на значение из-за констант квантования, которые хранятся вместе со сжатыми данными — это своего рода аналог словаря из традиционных алгоритмов архивирования данных ZIP и RAR. При большом размере контекстного окна эти накладные расходы всё равно оказываются значительными.

TurboQuant помогает устранить эти накладные расходы в двухэтапном процессе, то есть обходится вообще без словаря. На первом этапе срабатывает метод PolarQuant — преобразование векторов данных из декартовых в полярные координаты. Каждому вектору присваивается значение радиуса (длины) и угла (направления). Угловые распределения предсказуемы и сконцентрированы, поэтому PolarQuant обходится без ресурсоёмкого этапа нормализации каждого блока, который неизбежен с традиционными квантизаторами. На выходе получается сжатие высокого качества без накладных расходов на хранение констант внимания — словарей.

На втором этапе применяется 1-битный слой коррекции ошибок с использованием квантованного алгоритма Джонсона-Линденштрауса. Остаточная ошибка квантования проецируется в пространство меньшей размерности, каждое значение сводится к одному знаковому биту, устраняя тем самым систематическую погрешность в вычислениях оценки внимания с незначительными дополнительными издержками.

Google протестировала две составляющих алгоритма TurboQuant в отдельности и их совместную работу в бенчмарках с длинным контекстом LongBench, Needle In A Haystack, ZeroSCROLLS, RULER и L-Eval на открытых моделях Gemma и Mistral. TurboQuant показал идеальные результаты в задачах класса поиска «иголки в стоге сена» добившись сжатия KV-кеша минимум в шесть раз. В группе тестов LongBench, включающей ответы на вопросы, генерацию кода и создание сводок, TurboQuant оказался не хуже, а то и лучше базового алгоритма сжатия KIVI по всем задачам.

TurboQuant показал эффективность не только в работе с большими языковыми моделями, но и в векторном поиске — его протестировали в сравнении с известными алгоритмами сжатия Product Quantization и RabbiQ на наборе данных GloVe. Даже без обучения и оптимизации разработанный в Google алгоритм показал более качественные результаты, чем его соперники, настроенные специально для работы с этим набором данных. TurboQuant вообще не требует ни обучения, ни тонкой настройки, а его выполнение сопровождается незначительными накладными расходами — он готов для развёртывания даже в условиях повышенной нагрузки.