Всем привет! Меня зовут Денис Кузнеделев, я работаю в команде Yandex Research. Моё основное направление деятельности на данный момент — задача сжатия и ускорения больших языковых и картиночных моделей. Затраты на обучение, инференс и деплой LLM стали одной из ключевых инфраструктурных проблем индустрии: дефицит вычислительных ресурсов, нехватка видеопамяти и высокие требования языковых моделей к вычислительным ресурсам препятствуют масштабированию решений. 

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

Зачем сжимать большие языковые модели

За последние несколько лет возможности языковых моделей значительно возросли. Современные ведущие LLM — универсальные эксперты, способные решать олимпиадные задачи по математике, бизнес‑задачи и редактировать большие репозитории с кодом.

В то же время успех LLM во многом обусловлен их размером и значительным объёмом вычислений при каждом прямом проходе. Развёртывание больших нейросетей требует дорогостоящей инфраструктуры и значительного энергопотребления. Поэтому снижение затрат на их использование представляет существенный интерес. Кроме того, многие хотят запускать языковые модели локально на ПК с ограниченным объёмом видеопамяти и количеством вычислительных ядер. В таких случаях сжатие модели становится необходимым, чтобы обеспечить адекватную скорость работы.

Методы сжатия языковых моделей

Существует множество методов сжатия больших языковых моделей. В рамках сегодняшнего рассказа мы рассмотрим три.

Неструктурированный прунинг. Не все параметры нейронной сети одинаково важны. Определённую долю из них (в зависимости от модели и задачи) можно отбросить (занулить), почти не потеряв в качестве. 

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

Однако современные ускорители вычислений (GPU, TPU, ASIC) обычно обрабатывают параметры большими группами, и разреженные операции общего вида не имеют аппаратной поддержки. Поэтому неструктурированный прунинг интересен в основном для инференса на CPU и специальных архитектурах (например, чипах Cerebras).

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

Квантизация. Исходное представление весов модели (чаще всего fp16/bf16 для современных моделей) избыточно. Параметры модели можно представить в типе меньшей точности (INT8, INT4, FP8, FP4), который принимает меньшее количество возможных значений, почти не изменяя выход модели. Этот подход популярен для сжатия LLM благодаря простоте использования, доступности и удобным для пользователей реализациям (GGUF, Ollama и другие фреймворки). Подробный гайд по квантизации уже публиковал мой коллега.

Равномерное и неравномерное сжатие

Наиболее популярная стратегия сжатия — однородное равномерное сжатие, при котором все слои сжимаются одинаково. Для прунинга это означает одинаковую степень прореживания всех слоёв, а для квантизации — одинаковую битность и формат квантизации во всех слоях.

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

Но как определить, какие веса важны, а какие нет? Это сложный вопрос, и на него пока нет полностью удовлетворительного ответа.

В предыдущих работах на эту тему (Weight subcloning: direct initialization of transformers using larger pretrained ones, Shortened LLaMA: Depth Pruning for Large Language Models with Comparison of Retraining Methods, ShortGPT: Layers in Large Language Models are More Redundant Than You Expect, The Unreasonable Ineffectiveness of the Deeper Layers) были предложены различные эвристики для определения важности трансформерных блоков:

• косинусная близость между входами и выходами (если блок важный, он сильно «поворачивает» активации; если нет — выполняет операцию, близкую к тождественному преобразованию, и от него можно избавиться);

• приращение функции потерь при отбрасывании блока.

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

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

Хороший метод неравномерного сжатия должен учитывать такие нюансы. Наша цель — разработать метод, свободный от предположений о линейности и монотонности ошибок сжатия. И у нас это получилось.

Метод EvoPress

Сначала мы решили проверить, справедливо ли предположение о линейности и монотонности. Мы обнаружили интересный пример: у модели Llama-3-8B есть конфигурация прореживания блоков, при которой отбрасывание дополнительного количества слоёв приводит к улучшению качества модели. Это опровергает гипотезу о линейности и монотонности.

Как быть?

Мы предложили эволюционный алгоритм нахождения оптимальных конфигураций сжатия. Его суть в следующем:

1. Определяется мера качества конфигурации (fitness function), которую нужно оптимизировать. В нашей работе мы использовали KL‑дивергенцию между логитами исходной несжатой модели и сжатой на подвыборке калибровочных данных как меру различия между распределениями вероятностей двух моделей.

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

3. Прогоняются итерации эволюционного поиска. На каждой итерации создаётся несколько дочерних конфигураций из родительской с помощью случайных (специфичных для метода сжатия) преобразований. Лучшая дочерняя конфигурация становится родительской на следующей итерации.

Итерация поиска выглядит так:

• создаются дочерние конфигурации из родительской;

• на первом шаге оценивается качество всех конфигураций на небольшой подвыборке, сохраняются несколько лучших и начинается следующий раунд;

• на следующем раунде оставшиеся конфигурации прогоняются на большей выборке данных, отбираются лучшие;

• процесс повторяется 2–3 раза: количество кандидатов уменьшается, а выборка увеличивается, пока не останется одна конфигурация;

• если она лучше родительской, то становится новым родителем.

Случайные преобразования для рассмотренных методов сжатия:

• Неструктурированный прунинг: для каждого слоя создаются конфигурации с разной степенью прореживания с помощью алгоритма (например, SparseGPT); на каждой итерации степень прореживания слоя можно увеличить или уменьшить.

• Структурированный прунинг: отбрасывается один из «живых» трансформерных блоков и «воскрешается» один из «мёртвых», при этом общая степень сжатия сохраняется.

• Квантизация: битность квантизации понижается в некоторых слоях и повышается в других.

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

Для первой проверки эффективности мы решили выяснить, сможет ли эволюционный алгоритм найти лучшую конфигурацию на маленькой модели (Tiny‑Llama), где её можно найти полным перебором. Оказалось, что алгоритм справляется с этой задачей.

Эксперименты

Для проверки эффективности подхода мы применили его к трём вышеупомянутым техникам сжатия больших языковых моделей на семействах Llama-2, Llama-3, Mistral и Phi-3. Для оценки качества моделей использовались стандартные академические бенчмарки из Language Model Evaluation Harness.

В случае структурированного прунинга EvoPress заметно опережает предыдущие подходы, основанные на эвристиках, особенно при более высоких степенях сжатия. 

Мы также проверили метод на неструктурированном прунинге при сжатии на 50, 60 и 70%. EvoPress позволяет получить несколько дополнительных процентов качества по сравнению с однородным сжатием и методом OWL, основанным на распределении выбросов.

Проанализировав полученные конфигурации, мы заметили, что EvoPress обычно стремится сильнее всего проредить Q‑проекцию в Self‑Attention и как можно меньше трогать V‑проекцию.

И в конце концов, мы попробовали метод в контексте квантизации. При 3-битном сжатии с использованием GPTQ в качестве базового алгоритма EvoPress позволяет выжать примерно процент качества в среднем по задачам.

Паттерн сжатия напоминает неструктурированный прунинг. Q‑ и K‑проекции сжимаются сильнее, а V — слабее всех остальных слоёв.

Сколько ресурсов требует эволюционный поиск

Стоимость вполне приемлема. Быстрая версия поиска достигает почти полной сходимости примерно за час на одной видеокарте (consumer‑grade GPU), что доступно многим энтузиастам. Также потребуется около 100–200 ГБ места на диске для хранения конфигураций при неструктурированном прунинге и квантизации.

Заключение

Итак, мы рассмотрели EvoPress — гибкий и универсальный подход к поиску оптимальной сжатой конфигурации для различных техник сжатия моделей. В отличие от традиционных однородных подходов и простых эвристик, EvoPress использует эволюционный алгоритм для автоматического поиска оптимальной, неравномерной конфигурации сжатия по слоям. За счёт этого модели удаётся сжимать сильнее там, где это возможно без существенной потери качества, и «щадить» наиболее чувствительные параметры.

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