Исследователи Google Research выпустили в открытый доступ JaxPruner — библиотеку для прореживания (pruning) и разреженного обучения (sparse training) нейросетей, построенную на фреймворке JAX. Инструмент нацелен на научные команды, которым нужен гибкий тулкит для быстрой проверки идей и сравнения методов на разных бенчмарках.
Зачем нужна разреженность
Активные исследования разреженности (sparsity) — один из способов повысить эффективность глубокого обучения, но для практического применения нужна более тесная связка железа, софта и алгоритмов. В нейросетях разреженность может относиться либо к активациям, либо к параметрам модели. JaxPruner сфокусирован именно на параметрической разреженности: более ранние работы показывали, что разреженные модели способны превосходить по качеству плотные сети с тем же числом параметров.
Авторы выделяют два подхода к получению разреженных параметров. Pruning превращает уже обученную плотную сеть в разреженную ради ускорения инференса. Sparse training, наоборот, с самого начала обучает разреженную сеть, снижая затраты уже на этапе тренировки.
Почему выбрали JAX
JAX за последние годы заметно прижился в исследовательском сообществе. Его ключевое отличие от PyTorch и TensorFlow — чёткое разделение функций и состояния: полное состояние модели хранится в одном месте, что упрощает трансформации функций (вычисление градиентов, гессианов, векторизацию) и облегчает построение общих процедур сразу для нескольких методов прореживания и разреженного обучения. Кроме того, параметрическая разреженность не зависит от данных, а значит хорошо подходит для аппаратного ускорения. Несмотря на отдельные реализации техник вроде N:M-разреженности и квантования, полноценной библиотеки для исследований разреженности в экосистеме JAX не было — это и подтолкнуло команду Google Research к созданию JaxPruner.
Три принципа разработки
Библиотека должна была отвечать на вопросы вроде «какой паттерн разреженности даёт лучший баланс точности и производительности» и «можно ли обучить разреженную сеть без предварительного обучения большой плотной модели». Разработчики опирались на три принципа.
Быстрая интеграция: ML-исследования развиваются стремительно, а кодовые базы постоянно меняются, поэтому JaxPruner построен поверх известной библиотеки оптимизации Optax и требует минимальных доработок для встраивания в существующие проекты. Переменные состояния, нужные для прореживания и разреженного обучения, хранятся вместе с состоянием оптимизатора — это упрощает параллелизацию и чекпоинтинг.
Приоритет исследований: библиотека опирается на общий API, единый для разных алгоритмов, что делает переключение между методами и базовыми решениями (baselines) простым, как и смену популярных структур разреженности.
Минимальные накладные расходы: JaxPruner следует общепринятой практике задания разреженности через бинарные маски, что требует дополнительных операций и памяти для их хранения. Главной задачей разработчиков было свести эти накладные расходы к минимуму.
Исходный код JaxPruner опубликован в открытом доступе на GitHub вместе с обучающими материалами.
Источник: MarkTechPost.