Neural Tangents: библиотека для обучения бесконечно широких нейросетей

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

Neural Tangents

Чтобы облегчить работу с бесконечно широкими нейросетями, исследователи из Google AI опубликовали библиотеку Neural Tangents. Neural Tangents написана на JAX, инструменте для автоматического дифференцирования функций на Python и NumPy. Библиотека позволяет строить и обучать бесконечно широкие нейронные модели. Кроме того, пользователи могут обучать нейросети с конечным размером.

Примером использования Neural Tangents является ансамблирование. В базовом случае нейросеть случайно инициализируется и затем обучается с использованием градиентного спуска. Инициализация и обучения множества таких нейросетей приводит к ансамблю моделей. Обычно исследователи усредняют предсказания разных моделей в ансамбле. Недостаток такого подхода заключается в вычислительной сложности обучения набора моделей. Для бесконечно широких нейросетей ансамбль можно описать как гауссовский процесс, где среднее и дисперсия будут считаться во время обучения.

Библиотека позволяет собирать и обучать ансамбли бесконечно широких нейросетей в 5 строчек кода.