Model Pruning

Прунинг модели удаляет параметры, которые почти не влияют на результат обученной нейронной сети, — от отдельных весов до целых фильтров, голов внимания или слоёв, чтобы сделать модель компактнее и дешевле в запуске. Большинство сетей избыточно параметризованы, поэтому можно удалить значительную часть их весов с незначительным снижением точности, особенно если впоследствии модель дообучить. Метод прунинга с последующим переобучением, предложенный в 2015 году, сократил число параметров AlexNet с 61 миллиона до 6,7 миллиона, а VGG-16 — со 138 миллионов до 10,3 миллиона без потери точности на ImageNet. Прунинг — стандартный этап подготовки моделей для телефонов, камер и другого оборудования для периферийного ИИ; теперь его также применяют к большим языковым моделям (LLMs).

Как работает прореживание моделей#

При обрезке каждому потенциальному весу или структурному элементу присваивается оценка важности. Элементы с самыми низкими оценками удаляются, а оставшиеся веса модели проходят тонкую настройку, чтобы восстановить точность. Самый простой и по-прежнему распространённый критерий — модуль веса: веса, близкие к нулю, или фильтры с небольшой нормой L1 или L2 меньше всего влияют на выход при удалении. В других критериях градиенты или активации используют для оценки того, насколько вырастет функция потерь без каждой части модели. Одновременное удаление большой доли параметров обычно снижает точность, поэтому обрезку, как правило, выполняют в несколько этапов, каждый из которых завершается тонкой настройкой; процесс продолжают, пока модель не достигнет целевого размера или точность не опустится ниже допустимого уровня:

Обрезка также изменила представления исследователей о размере сетей. Согласно гипотезе лотерейного билета, плотная сеть содержит небольшие подсети, которые при обучении отдельно, начиная с исходных начальных весов, достигают точности всей сети. Авторы гипотезы стабильно находили на MNIST и CIFAR-10 такие подсети размером менее 10–20% от исходной модели.

Неструктурированная и структурированная обрезка#

При неструктурированной обрезке отдельные веса обнуляются в произвольных местах тензора. Она позволяет добиться высокой разреженности при небольшом снижении точности, но форма тензоров не меняется, поэтому стандартные плотные ядра выполняют те же вычисления; для ускорения нужны разреженные ядра или поддержка на уровне оборудования. GPU NVIDIA начиная с поколения Ampere поддерживают один из таких шаблонов — разреженность 2:4, при которой обнуляются два из каждых четырёх последовательных весов, — и обеспечивают скорость вычислений до двух раз выше, чем для плотных данных. При структурированной обрезке удаляются целые фильтры, каналы, головы или слои. В результате получается более компактная плотная модель, которая быстрее работает на любом процессоре, однако удаление того же числа весов сильнее сказывается на точности. Удаление фильтра также требует удалить соответствующий входной канал следующего слоя, а библиотеки, такие как Torch-Pruning, автоматически отслеживают такие зависимости. Утилиты обрезки в PyTorch показывают различия между этими подходами:

import torch
from torch.nn.utils import prune

# Two identical convolution layers, one for each pruning style
torch.manual_seed(0)
a = torch.nn.Conv2d(16, 32, kernel_size=3)
b = torch.nn.Conv2d(16, 32, kernel_size=3)

# Unstructured: zero the 50% of individual weights with the smallest magnitude
prune.l1_unstructured(a, name="weight", amount=0.5)

# Structured: zero the 50% of output filters with the smallest L2 norm
prune.ln_structured(b, name="weight", amount=0.5, n=2, dim=0)

for name, layer in [("unstructured", a), ("structured", b)]:
    zeros = (layer.weight == 0).float().mean().item()
    empty = (layer.weight.flatten(1) == 0).all(dim=1).sum().item()
    print(f"{name}: {zeros:.0%} of weights zero, {empty} of 32 filters empty")

В обоих слоях обнулено 50% весов. Пустые фильтры есть только в структурированно обрезанном слое — 16 из 32; их удаление оставляет плотный слой с вдвое меньшим числом фильтров, что сокращает вычисления вдвое на обычном оборудовании. Оборудование, рассчитанное на работу с разреженными данными, также выигрывает от неструктурированной обрезки: интеграция CVflow компании Ambarella для Ultralytics YOLO, которую Ambarella поддерживает в собственной версии пакета, постепенно обрезает веса во время обучения YOLO и сохраняет разреженность при экспорте в ONNX для компиляции под периферийные ИИ-чипы Ambarella.

Прореживание и родственные методы#

• Квантование модели: Снижает точность представления всех весов, например с 32-битных чисел с плавающей запятой до 8-битных целых чисел, не удаляя сами веса. Обрезку и квантование часто сочетают: обрезка уменьшает число операций, а квантование удешевляет каждую из них.
• Дистилляция знаний: Обучает небольшую модель-ученика имитировать большую модель-учителя; часто её применяют после структурированной обрезки, чтобы восстановить точность, потерянную при удалении целых фильтров или слоёв.
• Dropout: Временно обнуляет случайные активации во время обучения, выступая в роли регуляризации; при инференсе используется полная сеть.

Обрезка больших языковых моделей#

Повторное обучение LLM после каждого этапа обрезки обходится слишком дорого, поэтому методы однократной обрезки удаляют параметры за один проход, используя небольшой калибровочный набор. Метод SparseGPT сократил разреженность моделей OPT с 175 миллиардами и BLOOM со 176 миллиардами параметров до 60% менее чем за 4,5 часа, лишь незначительно увеличив перплексию; он также поддерживает шаблон 2:4, который ускоряют GPU. Структурированная обрезка в сочетании с дистилляцией знаний позволяет создавать компактные плотные модели: на основе предварительно обученной модели с 15 миллиардами параметров компания NVIDIA создала модели Minitron с 8 и 4 миллиардами параметров, обрезав каналы эмбеддингов, нейроны MLP и головы внимания, а затем повторно обучив их, используя до 40 раз меньше токенов, чем при обучении с нуля.