Сжатие языковых моделей без дополнительного обучения

2 Дек в 16:00

Я — Валентин Малых, возглавляю направление фундаментальных исследований в MWS AI. В этой статье я представляю ReplaceMe — новый способ структурного прунинга трансформерных моделей без этапа дообучения. С его помощью мы удаляем отдельные слои глубины, сохраняя большую часть качества: например, LLaMA-2 при сжатии на 25 % демонстрирует 92,5 % от первоначальной точности.

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

Классические методы прунинга требуют «лечения» (healing) — дообучения уменьшенной копии модели, что отнимает часы GPU-времени и значительный объём энергии. ReplaceMe позволяет отказаться от этой стадии и ускоряет процесс в десятки раз.

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

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

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

Как устроен алгоритм: три этапа

Этап 1. Поиск участков «почти линейного» преобразования

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

Где L_i и L_i+n — активации блоков, а D(·) измеряет угол между векторами. Чем меньше угол, тем меньше изменения направления, и тем проще заменить блоки одной матрицей.

Этап 2. Оценка матрицы замены

Нужно найти матрицу T, такая что применение T к выходу блока i приближает выход блока i+n. Задача формулируется так:

Где M_i — активация feed-forward слоя i, Y_i — выход attention, а h измеряет расхождение тензоров. Есть два варианта решения:

1) Аналитический метод (минимизация L₂-потерь): при квадратичной ошибке задача сводится к линейной регрессии, решение вычисляется за один шаг через псевдообратную матрицу. Высокая скорость — матричные умножения на GPU занимают секунды.

2) Оптимизация по косинусному расстоянию: более точная, но требует градиентного спуска (Adam, lr=1e−4, batch 1024, ~10 эпох). Сохраняет пару процентов точности дополнительно, однако занимает минуты вместо секунд.

Оптимизация матрицы T по косинусной метрике

Этап 3. Встраивание и удаление блоков

Полученную матрицу T* умножают на веса второго линейного слоя в блоке i. Две последовательные линейные операции эквивалентны одной, поэтому после корректировки весов промежуточные слои i+1…i+n просто удаляются, не изменяя архитектуру и не добавляя параметров.

Результаты экспериментов

Сравнение ReplaceMe с другими методами прунинга

На LLaMA-2-7B мы удалили 8 из 32 блоков (−25 %). ReplaceMe сохранил 92,5 % качества после одного прохода за 11 минут, опередив методы с дообучением (LLM-Streamline 92,0 %, UIDL 90,3 %). Конкуренты потратили несколько часов на healing.

При более строгом тесте на LLaMA-3-8B-Instruct без дообучения UIDL доходил до перплексии 2216, остальные — 12–30, а ReplaceMe удерживал 15,88 и 90,9 % точности (конкуренты 82–85 %).

Эффективность впечатляет: ReplaceMe в 170 раз быстрее полного дообучения и в 107 раз экономичнее по энергии, сохраняя при этом на 2,2 п.п. выше точность.

Мы также проверили модели разного масштаба: LLaMA-3.2-1B (−25 %) дала 87,7 % качества, LLaMA-3.1-8B — 91,8 %, LLaMA-3-70B (−25 %) — 96,6 %, а при −37,5 % — 90,6 %. Чем крупнее сеть, тем большую долю блоков можно удалить при приемлемых потерях.

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

Ограничения и рекомендации

Для оценки T достаточно 1 000–2 000 примеров, но для снижения перплексии лучше использовать 8 000–16 000. Модели, дообученные на инструкциях, требуют соответствующих датасетов.

Метод без дообучения эффективен при сжатии до 25–37 %. Более агрессивный прунинг требует дообучения только T — это гораздо быстрее, чем восстановление всей модели.