Новый метод обучения нейросетей: что предлагает Thinking Machines?

Исследователи стартапа Thinking Machines, возглавляемого Мирой Мурати, представили новую методику обучения нейросетей, которая обещает повысить стабильность и эффективность работы с моделями. В отличие от традиционных методов регулирования весов и активаций, авторы предлагают подход, который меняет саму структуру нейросетей.

Проблема, с которой сталкиваются нейросети
Во время обучения нейросетей одна из главных трудностей — это контроль за масштабами тензоров: весов, активаций и градиентов. Когда эти значения становятся слишком большими или маленькими, возникает ряд численных проблем, таких как взрывы и исчезающие градиенты, которые мешают корректному обучению. Для борьбы с этим часто используются такие методы, как gradient clipping, weight decay и layer norm, но они не решают проблему в корне.

Новая идея: обучение на многообразии
Предлагаемый подход идет намного глубже: вместо того, чтобы просто корректировать размеры тензоров, авторы предлагают ограничить сами структуры тензоров, заставив их "жить" в определённом многообразии (или manifold).

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

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

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

— Равномерное движение слоёв: Для равномерной стабилизации предложено использовать "бюджет шагов", чтобы все слои модели двигались синхронно, что повышает общую стабильность.

Результаты и ограничения
На тесте с CIFAR-10 новый метод показал лучшие результаты по сравнению с популярным оптимизатором AdamW, в том числе по стабильности обучения. Однако, несмотря на перспективы, у метода есть и важные вопросы:

• Как правильно выбирать пространства для разных типов слоёв?
• Будет ли метод эффективно работать на больших моделях и на float16?
• Как его масштабировать для более сложных задач?
• И, конечно, серьёзные вычислительные затраты.

Подход Thinking Machines представляет собой интересную и фундаментальную инновацию в обучении нейросетей. Хотя этот метод ещё далёк от практического применения, его потенциал в решении проблем численной стабильности и масштабируемости нейросетей нельзя недооценивать.

🔥 — Подход имеет потенциал для масштабирования и решения текущих проблем
🤔 — Сложности с вычислениями и масштабированием остаются слишком большими

Data Science