Кронекером по Фишеру: как при сжатии LLM учесть кривизну второго порядка и не умереть

На связи Вика Чекалина из команды «Мультимодальный ИИ» AIRI. Сегодня мы поговорим о том, как выкручиваться, когда вычислительных ресурсов мало и для деплоя необходимо сжать LLM, но так, чтобы потеря в качестве была минимальной. Логично ^[1], что нужно использовать дополнительную информацию — например, матрицу Фишера (FIM), которая содержит в себе информацию о важности и взаимодействии параметров модели. 

Но FIM реальных LLM слоев просто огромные, и работать с ними напрямую вычислительно тяжело или просто невозможно. Люди придумывают различные упрощения: например, можно отталкиваться от предположения, что матрица Фишера диагональна — это сильно сокращает объём памяти ^[2], однако корреляция между различными параметрами слоя теряется навсегда. Мы решили не идти на такие компромиссы и придумали метод ^[3], который быстро и легко позволяет параметризировать полную недиагональную матрицу Фишера без потери информации о корреляции.

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

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

Обычный SVD минимизирует норму реконструкции весов и считает все параметры слоя равнозначными. Но для качественного сжатия — когда модель показывает минимальное падение качества — нужно знать, какие веса критичны, а какими можно пожертвовать, а также как ошибки ^[4] в разных местах могут компенсировать или усиливать друг друга. Это описывается важностью параметров и информацией об их кросс‑корреляциях.

Веса слоя как матрично-нормальное распределние

Допустим, что модель предобучена и находится в локальном минимуме, а веса слоя по матрично-нормальному закону (matrix normal, MN) с плотностью

По определению, ковариация выражается через произведение Кронекера, а обратная к ней матрица точности есть информация Фишера этой гауссианы по параметру среднего:

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

Оптимальное разложение весов слоя GFWSVD

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

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

Теорема 3.1. Пусть  — матрица весов одного слоя обученной нейронной сети. Предположим, что выполнены следующие условия:

1. Модель обучена на supervised‑задаче с MLE‑функционалом (например, кросс‑энтропийной функцией потерь).

2. Эмпирическая матрица Фишера одного слоя , допускает кронекерову факторизацию .

3. Веса взяты из MVN‑распределения по формуле (2) в окрестности оптимума.

При этих условиях оптимальное ранг‑ приближение , минимизирующее приращение потерь , задаётся формулой

где и получены из соответствующих разложений Холецкого, и ; получается ранг‑ усечённым SVD вспомогательной матрицы .

Здесь , — кронекеровы факторы информации Фишера.

Доказательство

В предположении, что функция потерь порождена методом максимального правдоподобия (MLE), гессиан в точке оптимума совпадает с информацией Фишера. Следовательно, уравнение (2) можно заменить суррогатной задачей

при и . Подставляя и применяя разложение Холецкого к множителям A и B, получаем:

Было установлено, что в предположении оптимальное решение этой задачи может быть получено с помощью стандартного SVD вспомогательной матрицы (SVD есть решение MLE для распределений вида (2)). Итоговое решение находится в два шага:
1) поиск оптимального решения ранга для вспомогательной задачи: ,
2) восстановление оптимального решения исходной задачи посредством обратного преобразования: , что даёт наилучшее решение для уравнения (7).

Следовательно, разложение представляет собой оптимальное сжатие для уравнения (6), которое, в свою очередь, обеспечивает минимальный прирост ошибки для заданной задачи, определённой уравнением (4).

Конец доказательства.

Другими словами, наша теорема гарантирует нам оптимальность низкорангового сжатия, если, конечно, мы сможем получить и  — инвертированные матрицы ковариаций и . Рисунок 1 иллюстрирует, что такой вид сжатия слоя весов — с перевзвешиванием на корень из матриц ковариаций — является обобщением обычного SVD, где не учитывается ни важность, ни корреляции параметров, а также взвешенно-диагональной версии SVD (FWSVD, о нём чуть ниже), где учитывается только важность параметров.

FWSVD как частный случай

FWSVD был предложен ^[5] в 2022 году. Этот подход взвешивает единичную матрицу SVD по диагональной матрице Фишера. Покажем коротко, что диагональный FWSVD — частный случай нашего разложения GFWSVD. FWSVD минимизирует

то есть перевзвешивает веса одной диагональной матрицей . В наших терминах это соответствует вырожденной кронекеровой аппроксимации Фишера, где одна сторона — единичная матрица, а другая — диагональная: . Решение задачи даёт

то есть с точностью до константы . Подставляя в Теорему 3.1, получаем факторы

где  — ранг‑ SVD матрицы . Это ровно решение задачи FWSVD. Таким образом, в рамках GFWSVD диагональный FWSVD всегда с нами, но теперь мы учитываем недиагональные части фактор‑матриц с обеих сторон (и, соответственно, недиагональные элементы матрицы Фишера).

От гессиана к факторам: разложение большой матрицы без матрицы

Казалось бы, дело за малым — написать разложение Кронекера матрицы Фишера и спокойно делать оптимальное сжатие! Однако это матрица с квадратичной размерностью относительно размера линейного слоя, то есть для слоя LLaMA 2 размерами несколько тысяч на несколько тысяч параметров матрица Фишера будет матрицей размера миллион на миллион. Её проблематично даже загрузить на гпу, не то, что делать её разложение.

Помня, что мы работаем в точке минимума, мы считаем гессиан и матрицу Фишера тождественными. Матрица Фишера в нашей задаче определяется внешнее прозведение вытянутых градиентов:

где  — матрица градиентов весов слоя на ‑м батче, а  — соответствующий ей вектор.

Кронекерово разложение сводится к SVD матрицы Фишера после перестановки (), а SVD — к операции matvec/rmatvec, то есть умножению некоего вектора на матрицу . Матрица после перестановки будет иметь вид:

Как уже было отмечено, для объекта размером несколько миллионов на несколько миллионов (для LLM‑слоя размера имеет размер) провести разложение затруднительно. Перевод же в спарс‑вид сильно замедляет matvec/rmatvec, который предполагает проход по всем элементам матрицы.

Мы обошли эти ограничения, используя свойство произведения Кронекера

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

Новый Алгоритм 1 разложения Кронекера для гессиана линейного слоя выглядит так:

Требования: список градиентов {g_i}, число батчей |D|
1: I_F  ← (1/|D|) Σ g_i g_iᵀ            # не материализуется
2: Ĩ_F  ← (1/|D|) Σ G_i ⊗ G_i           # не материализуется
3: (u, σ, vᵀ) ← старший сингулярный триплет Ĩ_F (результат SVD)
4: b ← u · σ
5: a ← v
6: B ← reshape(b, (m, m))
7: A ← reshape(a, (n, n))
8: return (B, A)

Здесь старший сингулярный триплет определяется SVD, прочитанным с новыми ускоренными matvec/rmatvec.

Алгоритм 1 имеет теоретическую сложность относительно обычного , но практически расчёт факторов для одного слоя популярной модели LLaMA 2 7B проходит примерно за минуту (ранее это было более 9 часов).

Замеры скорости получения фактор‑представлений для разных LLM — на рисунке ниже:

Валидация

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

Алгоритм разложения примерно одинаков для любых DL‑моделей (в том числе LLM):

1. на калибровочном датасете собираем градиенты на нескольких шагах (для LLM — примерно 100);

2. по формуле (8) они дадут приближение матрицы Фишера в оптимуме(так называемый эмпирический Фишер);

3. раскладываем его с помощью Алгоритма 1 и получаем , ;

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

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

Учёт корреляций на синтетических данных

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

Целевые слои раскладываются нашим методом GFWSVD, и валидационная точность сравнивается с диагональным приближением Фишера и с обычным, невзвешенным SVD (yf hbceyrt отмечено как baseline). Результат подтверждает интуицию ^[6]: в изотропном режиме диагональ работает не хуже, но для скоррелированных данных GFWSVD выходит вперёд, и разрыв максимален в области сильного сжатия или низких рангов (левая часть картинки).

Llama3.1 8B

Второй тест — реальная LLM: Llama 3.1 8B, обученная на очень большом объёме данных. Её веса сильно скоррелированы и слабо избыточны, поэтому она особенно сложна для сжатия. Мы сравниваем GFWSVD с несколькими бейзлайнами:

• FWSVD — уже упомянутый диагональный Фишер: учитывает важность параметров, но не их корреляции;

• ASVD, SVD‑LLM — активационные методы: важность оценивается по нормам активаций (прокси, а не прямая чувствительность);

• Basis Sharing — сильный интересный бейзлайн. В то время как мы ловим корреляции между параметрами, Basis Sharingловит корреляции между слоями модели.

Кроме стандартного замера перплексии, мы валидируем сжатую модель на датасетах здравого смысла — ARC Easy, ARC Challenge, HellaSwag, PIQA, WinoGrande, OpenBookQA. Эти бенчмарки покрывают взаимодополняющие аспекты рассуждения: ARC Easy содержит школьные вопросы по естественным наукам с вариантами ответа, где ARC Challenge собран из вопросов, на которых ошибаются поисковые и словарные бейзлайны; HellaSwag проверяет выбор правдоподобного продолжения ситуации, PIQA — понимание физических свойств объектов и повседневных действий, WinoGrande — разрешение местоименных кореференций, требующее контекстных знаний, а OpenBookQA — применение известного научного факта к новой ситуации.

Таблица 3. Llama 3.1 8B Instruct, сжатие от 20% до 50% на WikiText-2 (WT2), PTB и среднего по шести датасетам commonsense‑рассуждений. Для перплексии меньше — лучше (↓), для точности больше — лучше (↑).

Значения точности приведены со стандартным отклонением примерно ±0.01, опущенным для читаемости.

Анализ таблицы показывает, что при сжатии от 20% до 50% GFWSVD заметно обходит методы, сжимающие слои независимо (FWSVD и ASVD), и превосходит Basis Sharing на всех уровнях сжатия вплоть до 50%. То есть для модели с сильно скоррелированными весами учёт внутрислойных корреляций даёт больше, чем межслойных. На 50% (это довольно сильный коэффициент сжатие для низкоранговых методов) почти все методы, кроме GFWSVD и basis sharing, показывают скор рандомного гадания. Единственные методы, при которых модель не рассыпается — GFWSVD и Basis Sharing, причем наш метод обгоняет его на 3%.

BERT

Мы также посмотрели, как метод ведёт себя на модели BERT — с неплотными, несильно скоррелированными параметрами. Мы применили GFWSVD к её слоям и провалидировали на бенчмарке GLUE (General Language Understanding Evaluation) — стандартном наборе из девяти задач понимания естественного языка. 

На рисунке ниже приведен результат: красная линия соответствует нашему методу GFWSVD, зеленая — FWSVD (чистодиагональный), ASVD — перевзвешиванию весов на информации из активаций. Видно, что учет корреляций начинает тем больше, чем больший уровень сжатия мы делаем. 

Заключение

В этой работе мы рассматриваем задачу низкорангового сжатия предобученных слоев LLM. Мы показали, что оптимальное сжатие получается, когда веса слоя перевзвешены на корень из обратных фактор‑матриц информации Фишера. В недиагональных элементах этих факторов хранится информация о корреляциях весов модели, и мы предложили метод, который делает расчет полных факторов возможным и быстрым.

Применение к реальным LLM показало, что наибольшей производительности метод достигает моделях с плотными, коррелированными весами (LLaMA3 8B). На иных моделях важность информации и корреляциях становится важна на высоких степенях сжатия. Больше подробностей в нашей статье ^[3], которую мы опубликовали на ICML 2026.

А ещё этот подход также работает и для квантизации ^[7], но об этом как-нибудь в другой раз…

Буду рада ответить на ваши вопросы!