Исследователи Sakana AI представили PC-ALM — послойную альтернативу обратному распространению, обучающую сети со 1000 слоями
Обратное распространение ошибки — это глобальный алгоритм: прямой проход, затем обратный проход, затем обновление весов, причём каждый этап зависит от предыдущего. В мозге неизвестен механизм, обеспечивающий такую синхронизацию фаз во всей сети, поэтому локальные альтернативы обучения, такие как predictive coding (PC), продолжают привлекать исследовательский интерес. Исследователи Sakana AI предлагают Augmented Lagrangian Predictive Coding (PC-ALM) — вариант PC, в котором каждое обновление остаётся локальным для слоя, но при этом восстанавливаются сигналы распространения ошибок, согласованные с обратным распространением. Исследовательская команда сообщает об обучении остаточных MLP-сетей глубиной до 1000 слоёв с результатом, отличающимся от обратного распространения примерно на 2 процентных пункта на MNIST.
Можно ли это использовать на практике? Да, в качестве исследовательского кода: эталонная реализация на JAX с лицензией MIT работает на CPU и воспроизводит матрицу экспериментов по ширине и глубине из статьи. Это метод обучения, а не модель, и он тестировался только на небольших наборах данных для классификации изображений.
Почему стандартный PC не справляется с глубокими узкими сетями
PC рассматривает каждую скрытую активацию как переменную оптимизации и штрафует квадратичное расхождение между активацией каждого слоя и предсказанием, поступающим с нижележащего слоя. Вывод выполняется посредством градиентного спуска по этой энергии; обучение — посредством шага обновления весов, подобного хеббовскому. Проблема в том, что сигнал обучения поступает на выходе и должен распространяться по цепочке локальных компромиссов. В глубоких узких сетях сигнал назначения исчезает задолго до того, как достигает входа. Innocenti и соавторы охарактеризовали разрыв между PC и BP как функцию ширины и глубины; наибольшим он оказывается, когда ширина меньше глубины.
Что изменяет PC-ALM
PC-ALM исходит из представления об обучении как об ограниченной оптимизации: минимизировать функцию потерь с учётом ограничения на каждом слое. PC представляет собой релаксацию этой задачи с квадратичным штрафом. Вместо этого PC-ALM использует расширенный лагранжиан, добавляя множитель Лагранжа к ограничению каждого слоя и сохраняя штраф PC. При λ = 0 получается в точности PC.
Вывод чередует 2 локальных шага: первичный градиентный шаг по активациям и двойственный шаг , накапливающий ошибку предсказания слоя. Выделение полного квадрата показывает, что каждый первичный шаг является стандартным шагом PC, в котором целевой показатель предсказания смещён на . После T шагов обновление весов действует на составной сигнал . Исследовательская команда интерпретирует это как PI-регулятор для каждого слоя: ошибка предсказания является пропорциональным членом, а множитель — интегральным. α = 0 соответствует PC; α = ρ при точном решении внутренней задачи даёт классический метод множителей.
Точные градиенты обратного распространения в линейном случае
В 1988 году ЛеКун отметил, что множители Лагранжа ограниченной сети в точке Каруша — Куна — Таккера равны сопряжённым переменным обратного распространения. Команда доказывает, что в линейных сетях PC при выполнении условия устойчивости спектрального радиуса PC-ALM сходится к этой точке ККТ: активации возвращаются к своим значениям прямого прохода, а каждый интегрируется до точного сопряжённого градиента BP. Граница устойчивости для каждой моды имеет вид , что при α = 0 сводится к условию PC. В отличие от монотонного градиентного потока PC, матрица итераций PC-ALM имеет комплексные собственные значения, порождающие затухающие колебания; α задаёт их частоту, но не скорость затухания.
Результаты
Исследовательская команда изучает остаточные MLP-сети с шириной и глубиной от 8 до 128 на наборах данных Fashion-MNIST и MNIST, используя параметризацию среднего поля Innocenti и соавторов и обучая модели в течение 1 эпохи. При бюджете вывода T = 2L PC-ALM соответствует обратному распространению при любых ширине, глубине и функции активации (identity, tanh, ReLU), тогда как показатели PC резко падают в глубоких узких конфигурациях. Эталонная конфигурация из репозитория (ширина 32, глубина 32, ReLU, Fashion-MNIST) сообщает о точности на тестовой выборке 78,66% для BP, 68,13% для PC и 77,75% для PC-ALM; при этом косинусное сходство градиентов с BP возрастает с 0,604 до 0,909.
Исследование расширяет эту картину: остаточные MLP-сети глубиной 1000 слоёв на MNIST (ширина 32, ReLU, 5 эпох) отстают от BP примерно не более чем на 2 пункта, а PC-ALM превосходит PC на всех проверенных наборах данных, включая ResNet-18 на CIFAR-10 и Tiny ImageNet.
Ключевые выводы
- PC-ALM добавляет множитель Лагранжа для каждого слоя в predictive coding; каждое обновление остаётся локальным для слоя.
- В линейных сетях множители сходятся к точным градиентам обратного распространения.
- Соответствует BP на матрице ширины и глубины от 8 до 128 при T = 2L; PC не справляется с глубокими узкими конфигурациями.
- Обучает остаточные MLP-сети глубиной 1000 слоёв с отставанием от BP примерно не более чем на 2 пункта на MNIST.
- Код на JAX с лицензией MIT воспроизводит результаты на CPU.
Ознакомьтесь со статьёй, блогом и репозиторием на GitHub. Все заслуги принадлежат исследователю этого проекта. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему сабреддиту о машинном обучении с более чем 150 тыс. участников и подписаться на нашу рассылку. Подождите! Вы есть в Telegram? Теперь вы также можете присоединиться к нам в Telegram.
Хотите сотрудничать с нами в продвижении вашего репозитория на GitHub, страницы на Hugging Face, релиза продукта, вебинара и т. д.? Свяжитесь с нами
Переведено автоматически с английского. Оригинал статьи — по ссылке ниже.