Привет, Хабр! Меня зовут Валентин Малых, я исполняю обязанности руководителя направления фундаментальных исследований в MWS AI. Сегодня я расскажу об одном нашем исследовании по сжатию LLM. Если простыми словами, то это про то, как сделать большую модель чуть-чуть менее требовательной в плане памяти и времени выполнения. Для это придумано три базовых техники: квантизация (загрубление весов модели), дистилляция (обучение уменьшенной копии) и прунинг (удаление части сети). Этот пост как раз будет про третий способ, точнее – недавно разработанный нами в сотрудничестве с зарубежными коллегами метод структурного прунинга по глубине без дообучения, который мы назвали ReplaceMe. Например, модель LLaMA-2 после нашего сжатия на 25% сохраняет 92,5% качества. Ниже – о том, как это работает.

Современные большие языковые модели достигли масштабов в миллиарды параметров, самые продвинутые версии (например, той же LLaMA, OpenAI GPT или Gemini от Google) насчитывают сотни миллиардов обучаемых весов. Для обучения таких моделей нужны дата-центры, запустить их на обычном железе практически невозможно. Поэтому модели нередко подвергают прунингу, когда удаляются целые слои или, например, отдельные головы внимания.

Большинство существующих методов прунинга требуют последующего переобучения модели — так называемого “лечения” (healing, восстановления производительности). На это уходят часы GPU-времени и много энергии. Но наш метод – ReplaceMe – значительно более эффективный.

Ключевая идея: линейная аппроксимация блоков

Современные модели, основанные на архитектуре трансформер, — это по большей части стопка одинаковых блоков. Каждый блок делает две вещи: сначала применяет механизм внимания (attention), который “смешивает” токены, а затем пропускает результат через полносвязный слой (он же dense, он же feed-forward). Параллельно каждой операции делается “короткое замыкание” (residual connection) — к выходу прибавляется вход.

Мы заметили: если взять несколько блоков подряд и посмотреть, как они трансформируют данные, часто оказывается, что итоговое преобразование близко к линейному. То есть вместо сложной последовательности операций можно подобрать специальную матрицу, которая даст примерно тот же результат.

Вот еще аналогия, которая будет понятна людям, увлекающимся фотографией. Представьте, что у вас есть фото, и вы последовательно применяете к нему пять фильтров: каждый немного меняет яркость, контраст, насыщенность. Если изменения небольшие, вместо пяти операций можно сделать одну коррекцию цветовой кривой, которая даст похожий результат. ReplaceMe делает то же самое, но с трансформерными блоками.

Как работает метод: три шага

Шаг 1. Выбор блоков для удаления

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

Затем для каждой возможной позиции вычисляют косинусное расстояние между выходом блока i и выходом блока i+n (где n — сколько блоков планируется удалить):

где L_i — выход i-го блока, n — количество блоков, которые планируется удалить, D(⋅) — функция расстояния между активациями, i* — оптимальный индекс блока (там, где расстояние минимально).

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

Почему именно косинусное расстояние? Мы проверили разные метрики, включая обычное евклидово расстояние (L₂-норма). Оказалось, что косинусное расстояние лучше коррелирует с итоговой точностью модели после прунинга. В приложении к нашей научной статье есть график, где видно, что чем меньше косинусное расстояние, тем выше точность после удаления блоков (ссылка есть в конце этого поста).

Шаг 2. Оценка линейного преобразования

Теперь нужна матрица T, которая превратит выход блока i в то, что должно было бы получиться после блока i+n. Математически это выглядит так:

где M_i​ — выход полносвязного блока i, Y_i​ — выход внимания, L_i+n​ — итоговый выход блока, a функция h измеряет различие между двумя тензорами.

Есть два способа решить эту задачу.

Способ 1: Аналитическое решение через метод наименьших квадратов

Если использовать обычное евклидово расстояние (L₂), задача сводится к классической линейной регрессии. Решение выписывается в явном виде:

Это формула для решения системы линейных уравнений через псевдообращение матрицы. Вычисляется за один проход, без дополнительных итераций. Перемножили матрицы на GPU, получили ответ за секунды.

Способ 2: Численная оптимизация по косинусному расстоянию

Косинусное расстояние как функция потерь дает лучшие результаты, но не имеет аналитического решения. Приходится использовать градиентный спуск. Задача формулируется так:

Здесь используется упрощенная формулировка: вместо трех матриц M_i, Y_i​, L_i+n хранятся только две — M_i​ и разность L_i+n−Y_i. Это экономит примерно 66% памяти.

Матрицу T находят численно через градиентную оптимизацию алгоритмом Adam (learning rate 10⁻⁴, batch size 1024, 10 эпох по калибровочным данным). Этот способ обычно добавляет пару лишних процентов точности на итоговых метриках. Цена вопроса — время: потребуется несколько минут на оптимизацию.

Шаг 3. Интеграция в модель

Полносвязный блок в трансформерном блоке состоит из двух последовательных линейных слоев. Найденную матрицу T* перемножают с весами второго линейного слоя в блоке i.

Это работает потому, что две последовательные линейные операции математически эквивалентны одной. После перемножения матриц получается обычный линейный слой с модифицированными весами — дополнительных параметров не добавляется.

Блоки с i+1 по i+n удаляются из модели. Архитектура остается стандартной, только с меньшим количеством блоков.

Результаты

Переходим к самому интересному — что получилось на практике, когда мы испытали метод на реальных моделях. 

Основные эксперименты проводились на модели LLaMA-2-7B — ее сжали на 25%, удалив 8 блоков из 32. Результаты сравнили с пятью современными методами прунинга, каждый из которых требует дообучения (того самого “лечения”) после удаления слоев.

ReplaceMe сохранил 92,5% качества исходной модели — это лучший результат среди всех участников сравнения. Ближайший конкурент с дообучением — LLM-Streamline — показал 92,0%, а метод UIDL достиг 90,3%. Остальные подходы (LLMPruner, SliceGPT, LaCo) сохранили от 77,5% до 82,7% качества. Время сжатия составило 11 минут для ReplaceMe против нескольких часов у конкурентов.

На модели LLaMA-3-8B-Instruct провели более жесткое сравнение: все методы тестировали без healing-фазы. У метода UIDL значение перплексии (меры неуверенности модели при предсказании следующего токена) выросло с 3,11 до 2216, то есть модель практически утратила способность генерировать связный текст. У методов LLMPruner и SVD-LLM это значение оказалось в диапазоне 12-30. ReplaceMe удержал показатель на уровне 15,88 и сохранил 90,9% точности при 82-85% у конкурентов.

Разница в вычислительной эффективности значительна. Метод UIDL использует LoRA на MLP-слоях — это облегченный вариант дообучения, полное дообучение модели потребовало бы гораздо больше ресурсов. Даже такой щадящий подход занимает 32,5 часа на GPU. ReplaceMe завершает сжатие за 11,5 минут, что в 170 раз быстрее. Энергопотребление ниже в 107 раз, выбросы CO₂ меньше в 77 раз. При этом точность на бенчмарках на 2,2 процентных пункта выше.

Метод протестировали на моделях разного размера. Компактная LLaMA-3.2-1B с одним миллиардом параметров при сжатии на четверть сохраняет 87,7% качества. Средняя по размеру LLaMA-3.1-8B показывает 91,8%. Большая LLaMA-3-70B при том же коэффициенте сжатия дает 96,6%, а при более агрессивном сжатии на 37,5% сохраняет 90,6% качества. Прослеживается четкая закономерность: чем крупнее исходная модель, тем сильнее ее можно сжимать без существенных потерь. В больших моделях выше избыточность, больше блоков дублируют функции друг друга.

Метод работает и на картиночных трансформерах: CLIP-ViT при сжатии на 13% почти не теряет качество и стабильно превосходит UIDL на задачах классификации и поиска изображений.

Детали и ограничения

Для оценки матрицы преобразования достаточно 1 000 — 2 000 примеров калибровочных данных, хотя для минимизации перплексии лучше взять 8 000 — 16 000. Для моделей, дообученных на инструкциях, эффективнее использовать инструкционные датасеты, а не просто текст из интернета.

Метод работает без дообучения при сжатии до 25-37% (зависит от размера модели). При более агрессивном сжатии нужно дообучение, но можно дообучать только матрицу преобразования, а не всю модель — это в разы быстрее.

Тем, кто хочет копнуть глубже

Работа выполнена коллективом исследователей из MWS AI, Университетов ИТМО, МУИТ и Университета Крита, IACM-Forth, Archimedes Athena RC и Polynome. Спасибо коллегам за красивую идею, работающее решение реальной проблемы и открытый код.

• Препринт научной статьи: arxiv.org/abs/2505.02819

• Репозиторий: github.com/mts-ai/ReplaceMe

В материале упоминаются опенсорс-разработки компании Meta Platforms, признанной экстремистской и запрещенной в РФ.