Compact Language Models via Pruning and Knowledge Distillation
[Статья][Код есть, хоть это поначалу и не очевидно]

Введение

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

Но если вы не гаражный стартап, а компания с триллионной капитализацией - вам он под силу.

Подданые кожаной куртки выпустили запруненные версии Nemotron-15B под названием Minitron-4B и Minitron-8B, которые по качеству якобы конкуретноспособны текущей SOTA в легком весе.

Метод

У трансформеров можно прунить:
☠️ Целые блоки (про что был обзор некоторое времчя назад)
☠️ Размерность признаков (был у вас d_model - 768, стал 512)
☠️ Головы в Attention
☠️ Внутреннюю размерность MLP

Первый метод про прунинг в глубину, а остальные про ширину.

Для оценки важности используются критерии, основанные на активациях:
⭐️ Средний attention по голове (при прунинге головы)
⭐️ Средняя активация по каналу (для MLP)
⭐️ Средняя активация LayerNorm (для оценки важности признаков)

Для прунинга в глубину смотрят на среднее косинусное расстояние между входом и выходом блока и изменение перплексии при выкидывании блока.

Рассматривают оценку важности one-shot, когда выкидываем все за раз, или итеративно, по чуть-чуть.

Далее выбирают некоторое (около 20) число конфигураций сжатых моделей с целевой степенью сжатия, сжимают модели и дообучают на 1.8B токенов. Берется самый перспективный кандидат и дообучается уже на ~100B токенов. Немало, но все же на 2 порядка меньше, чем у современных LLM.

На дообучении рассматривают сумму следующих лоссов:
1️⃣ Кросс-энтропию (исходный лосс)
2️⃣ KL-дивергенцию на логитах с несжатой моделью
3️⃣ L2 лосс между промежуточными признаками (так как сжатая модель может быть уже, есть матрица переводящая в размерность учителя)