Нейросети предложили учить через «чёрный ящик»

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

В гибридной системе часть вычислений выполняет альтернативный ускоритель, а остальную нейросеть по-прежнему обрабатывает видеокарта. Вперёд такая сборка работает без проблем: цифровые слои передают данные фотонному модулю, получают ответ и продолжают вычисления. Проблема возникает, когда нужно пройти по всей цепочке назад и определить, как скорректировать каждый блок. Фотонный модуль не раскрывает внутренние вычисления и не сообщает, как его настройки повлияли на итоговую ошибку. Для алгоритма это «чёрный ящик».

Как решить эту проблему? Учёные AIRI, Сколтеха и Центра квантовых технологий Сбера предложили метод Astralora. Настройки фотонного модуля он подбирает почти на ощупь: немного поворачивает его виртуальные ручки, снова запускает нейросеть и проверяет, стала ошибка больше или меньше. После серии таких проб алгоритм определяет, в какую сторону двигаться.

Чтобы обучать цифровые слои, расположенные перед «чёрным ящиком», Astralora создаёт его упрощённого дублёра. Он наблюдает, какие числа модуль получает и что возвращает, а затем учится приблизительно воспроизводить эту зависимость. Во время обратного прохода дублёр подсказывает предыдущим слоям, как нужно изменить их параметры.

Подход проверили на шести моделях фотонных блоков — от простого умножителя матриц до схем на микрокольцах, световых модуляторах и интерферометрах. Каждый из них по очереди ставили на место одного из цифровых блоков нейросети.

Первым тестом стал CIFAR-10 — набор из 60 тысяч маленьких изображений самолётов, машин, птиц, кошек, собак и других объектов. Гибридная нейросеть с фотонным блоком распознавала их почти так же точно, как полностью цифровая.

Затем модель учили различать 8732 короткие городские записи: лай собак, гудки машин, сирены, выстрелы, работающие двигатели, дрели, отбойные молотки и уличную музыку. Качество распознавания полностью цифровой нейросети составило 77,46%. Варианты с фотонными блоками получили от 75,6 до 78,95% — лучший слегка обогнал исходную модель.

В конце «чёрными ящиками» заменили сразу 12 крупных блоков GPT-2-подобной языковой модели. Модель продолжила обучаться, хотя обычное обратное распространение ошибки охватывало только оставшиеся цифровые слои — 190 млн параметров вместо исходных 417 млн.

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

Алгоритм позволяет подключать к обычной электронике ускорители, основанные на разных физических принципах — от фотоники до нейроморфных систем, — и совместно обучать всю конструкцию практически без потери качества. И это ли не прекрасно!

@anti_agi