Оптимизация, вдохновлённая физикой, улучшает глубокий прунинг LLM

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

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

Предположение о независимости становится всё более ограничивающим по мере того, как языковые модели становятся глубже, а их слои — более специализированными. Например, удаление блока 20 может иметь разный эффект в зависимости от того, отсутствует ли также блок 19 или блок 24. Эти зависимости аналогичны взаимодействиям между спинами. Их игнорирование может позволить сохранить приемлемое качество при умеренных уровнях сжатия, но способно привести к неудачному выбору, когда необходимо одновременно удалить множество блоков. Перебор всех комбинаций позволил бы учесть эти взаимосвязи, однако число возможных подмножеств экспоненциально растёт с глубиной модели, поэтому непосредственная оценка каждой сокращённой модели становится непрактичной.

Чтобы смоделировать как индивидуальную важность, так и взаимодействия, метод вводит переменную связи в остаточный путь каждого трансформерного блока. Нулевое значение означает сохранение блока, а единица — его удаление. Затем исследователи приближённо оценивают влияние этих переменных на функцию потерь модели с помощью разложения Тейлора второго порядка. В результате получается приближённая матрица Гессе. Её диагональные элементы описывают изолированный эффект изменения отдельных блоков, а внедиагональные элементы отражают попарные взаимодействия между различными решениями об удалении. Эти внедиагональные элементы содержат информацию, которую отбрасывают методы независимого ранжирования блоков.

Итоговая задача состоит в том, чтобы выбрать ровно M блоков из N, минимизируя квадратичную энергию xᵀH⁰x. На языке оптимизации это задача бинарной оптимизации с ограничениями. На языке физики — спиновое стекло Изинга со связями между всеми парами и сохраняющейся намагниченностью, поскольку число выбранных спинов «вверх» — а следовательно, и число удалённых блоков — фиксировано. Каждая низкоэнергетическая конфигурация спинов соответствует определённому подмножеству трансформерных блоков, которые следует удалить. В статье сообщается, что эта недорогая в вычислении величина энергии служит надёжным приближённым показателем реального качества модели в последующих задачах: конфигурации с меньшей энергией, как правило, дают сокращённые модели с лучшими результатами на бенчмарках.

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

В случаях, когда пространство конфигураций остаётся управляемым, исследователи выполняют полный перебор на одном GPU. Оценка нескольких миллионов конфигураций занимает считаные секунды, а метод можно масштабировать до десятков миллиардов вариантов, поскольку вычисление энергии для каждой конфигурации обходится недорого. Крупнейший описанный пример точного поиска предусматривает выбор восьми удаляемых блоков из 80-блочной модели Llama-3.3-70B. Такой поиск охватывает около 29 миллиардов конфигураций и занял примерно два дня.

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

Метод особенно эффективен при агрессивном сжатии, когда взаимодействия между решениями об удалении имеют наибольшее значение. При сжатии Llama-3.3-70B-Instruct на 50% предложенная стратегия выбора улучшает результат MMLU почти на 23 процентных пункта по сравнению с наиболее сильным из конкурирующих методов удаления блоков, оценённых авторами. Этот результат показывает, что рассмотрение отсечения как глобальной задачи взаимодействий, а не как списка независимых оценок слоёв, позволяет сохранить значительно больше возможностей модели, когда необходимо удалить большую часть сети.