Telegram-канал dl_stories - DLStories: Неотсортированное

DLStories

01 августа 2024 18:24

Я тут недавно перед одним интервью на стажировку решила-таки немножко заботать тему inference-time ускорения и сжатия больших моделей. В конце концов, модельки становятся все больше, требуют все больше памяти и времени для работы. Именно поэтому тема ускорения и сжатия сейчас очень на слуху — хочется понимать, что там в целом происходит.

Так вот, разных подходов к ускорению/сжатию моделей, конечно, очень много. Некоторые из них используют особенности железа (строения CPU/GPU), модифицируя некоторые операции внутри моделей так, чтобы они выполнялись как можно эффективнее на конкретном железе. Примеры таких подходов — FlashAttention и Mamba. Есть и другие подходы, которые работают с самой моделькой, стараясь добиться эффективного сжатия/ускорения вне зависимости от железа. И вот тут становится интересно, тут часто возникает красивая математика. И вот на прошлой неделе на ICML 2024 представили одну такую работу от исследователей Yandex Research и IST Austria по сжатию LLM с красивой математикой и довольно внушительными результатами.

Статья называется Extreme Compression of Large Language Models via Additive Quantization. Это метод экстремального сжатия LLM (до 2-3 бита на параметр!), который при этом сохраняет в среднем 95% качества. Метрики в статье репортят на LLama 2 и Mixtral, плюс на HF уже выложили много разных вариантов моделек, квантизованных с помощью метода.

Общая идея такая: авторы адаптируют для LLM технику квантизации Additive Quantization (AQ), которая изначально была предложена для экстремального сжатия многомерных векторов с сохранением метрики расстояния между ними. Transformer-based LLM обычно состоит из блоков вида Multi-Head Attention+MLP. Параметры Attention — обучаемые матрицы W_Q, W_K, W_V, параметры MLP слоя — также обучаемые матрицы W и b. Вот эти все матрицы мы и будем квантизовать с помощью Additive Quantization. При этом нам важно, чтобы квантизованные матрицы сохранили dot product similarity, и в этом как раз преимущество метода AQ — он позволяет сохранять метрики.

AQ работает так: делит матрицы на части по p последовательных строк, и каждая часть затем представляется в виде суммы M векторов, взятых из предварительно обученных M кодовых книг (codebooks). Вот в том, как обучаются codebooks и как из них затем выбираются векторы для представления частей матрицы, и кроется красивая математика) Я пока в ней разобралась не до конца, для меня это прям отдельная новая тема сжатия информации, и надо в нее чуть закопаться.
Скажу только, что сама идея обучать codebooks для представления весов модельки не нова. Для сжатия трансформеров похожее предлагали в этой статье, например. Но разница в том, как именно обучать эти codebooks, — в этом и состоит главное преимущество Additive Quantization.

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

Как-то так. Для меня статья получилась хорошим стартом в область сжатия моделек. Буду рада советам о том, что еще надо обязательно почитать, чтобы примерно понимать, что происходит в этой теме (я просто вообще-вообще раньше этим не интересовалась)

Ссылки:
✔️Статья
✔️Код на GitHub
✔️Квантизованные модельки на HF
✔️Статья на Хабре