Recurrent Independent Mechanisms (RIMs)
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, Bernhard Schölkopf
Статья: https://arxiv.org/abs/1909.10893

Ещё пройдёмся по малоизвестным, но интересным статьям.

Рекуррентные независимые механизмы, (Recurrent Independent Mechanisms или RIMs), от весьма достойного коллектива авторов целятся в создание модульной рекуррентной архитектуры, способной поймать динамику среды.

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

В этой архитектуре для моделирования динамической системы вся модель разделяется на k небольших подсистем, каждая из которых рекуррентная (чтобы можно было выучить динамику). Эти подсистемы и называются RIM’ами, и каждый RIM выучивает что-то своё.

Верхнеуровнево группы рекуррентных ячеек (RIM’ов) работают сравнительно независимо и изредка через механизм внимания взаимодействуют между собой, а также со входом. Каждый RIM в работе представляется какой-то рекуррентной сетью (LSTM, GRU) со своим набором параметров и вектором состояния.

Если более детально, то работа RIM’ов состоит из нескольких шагов:

1. Каждый RIM обращается через механизм внимания ко входу системы. RIM генерирует query (Q), а ключи и значения (K,V) получаются из элементов входа. Можно считать, что на входе есть множество каких-то переменных и каждый RIM должен выбрать что-то релевантное для себя.

2. Активироваться в итоге должно какое-то небольшое количество RIM’ов (это гиперпараметр, k_A). Реализация этого проста: ко входу добавляем вектор нулей, смотрим какое “внимание” обращает на этот вектор каждый из RIM, ранжируем по степени активации и выбираем k_A тех, кто обратил на нулевой вектор наименьшее внимание. Это множество активировавшихся модулей, S_t.

3. Активировавшиеся модули из множества S_t обсчитывают свою динамику по своим выученным функциям (то есть делают обычный шаг рекуррентной архитектуры, LSTM, например, h_{t+1} = LSTM(h_t, attention_to_input). Неактивировавшиеся модули остаются неизменёнными.

4. Активировавшиеся модули также могут считать состояние всех других модулей, включая неактивные (там тоже может быть что-то ценное). Это тоже делается через механизм внимания, каждый модуль генерит Q, а состояния остальных модулей дают K,V. Итоговое новое состояние модуля считается как сумма результатов механизма внимания и своего состояния из пункта 3 (в статье, похоже, опечатка, там состояние взято от предыдущего момента времени, но это неверно).

У архитектуры по сути два гиперпараметра: общее число модулей (K) и число активных модулей на каждом шаге (k_A). В принципе также есть тип ячейки (LSTM, GRU, whatever) с её параметрами, а также число голов внимания.

Сетки в принципе гоняли маленькие. Скрытое состояние суммарно 600 юнитов, K=6, k_A обычно 4, головы тоже 4.

Эксперименты проводили на задачах типа копирования (на вход подаётся последовательность символов, затем много пустых входов, а затем сигнал на воспроизведение оригинальной последовательности и система должна её выдать); sequential MNIST (классификация символов из MNIST, вытянутых в последовательность пикселей); задачка с прыгающими мячами (множество шариков разных масс и размеров движется внутри окошка, отскакивая друг от друга в соответствии с базовой физикой); а также RL-задачка на object picking из BabyAI, где много дистракторов.

RIM показали себя достойно, лучше LSTM, NTM (Neural Turing Machine), RMC (Relational Memory Core), трансформера, DNC (Differentiable Neural Computer). На Sequential MNIST аналогично. Модель неплохо обобщает, работает на последовательностях длин сильно больших, чем те, что она видела при обучении. На дистракторы тоже реагирует лучше, чем бейзлайны (меньше “отвлекается”).

В играх Атари при замене RIM'ом обычной LSTM и обучении через PPO результат ощутимо улучшается.

На языковом моделировании тоже получше, чем близкая по параметрам LSTM, так что можно и не только для RL применять.

Всякими абляциями показывают, что внимание на входе важно, коммуникация между RIM’ами важна, разреженность в активности модулей тоже. Если каждый механизм сложен, то легко можно получить доминирование одного модуля над всеми и тогда интересного обучения не получается. Если же модули не слишком сложные, то они вынуждены кооперироваться. Что-то в этом есть.

В общем это интересный inductive bias про то, что всё не обязано взаимодействовать со всем.

На ICLR 2020 не взяли :)

У этой движухи есть продолжения. Stay tuned!

Learning to Combine Top-Down and Bottom-Up Signals in Recurrent Neural Networks with Attention over Modules
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, Yoshua Bengio
Статья: https://arxiv.org/abs/2006.16981

Статья про дальнейшее развитие RIM — BRIM, Bidirectional RIM. Эту взяли на ICML 2020.

Серьёзное ограничение RIM (как и многих других архитектур, работающих с последовательностями наблюдений) в том, что в них отсутствуют top-down сигналы. По идее, предыдущий опыт даёт различную ценную информацию в виде убеждений, ожиданий, контекста и другой верхнеуровневой информации, могущей серьёзно повлиять на восприятие входных данных. BRIM собственно и предлагает совместить оба пути, bottom-up (обычный) и top-down.

Вот, кстати, недавно описанный Feedback Transformer тоже предлагал аналогичный путь для информации в сети, и он был неплох по результатам. В типовых же архитектурах передача информации сверху вниз отсутствует напрочь (рекуррентные сети — это не то, там связи лишь внутри одного слоя, сверху вниз их нет). А в мозгу таких связей навалом.

Ещё одно ограничение оригинальной работы про RIM — это отсутствие иерархии. В модели не было множества слоёв, были лишь конкурирующие модули, они все существовали на едином уровне иерархии.

Текущая работа добавляет к RIM, во-первых, иерархию слоёв, а, во-вторых, связи от более высоких в иерархии слоёв к нижним. Каждый слой по-прежнему состоит из некоторого числа отдельных модулей (как в RIM).

Общение между модулями теперь устроено богаче. Модуль из слоя L в момент времени t генерит ключи Q и обращается через механизм внимания к 1) нижележащим модулям со слоя L-1 в тот же момент времени t; 2) вышележащим модулям со слоя L+1 за предыдущий момент времени t-1; 3) ну и ко всё тому же вектору нулей, чтобы по активациям к этому вектору отобрать активные модули.

Внутри слоя общение модулей устроено как раньше — активные модули обновляют своё состояние по всем модулям текущего слоя, неактивные сохраняют старое состояние.

Никаких дополнительных лоссов модель не требует, это что называется drop-in replacement для LSTM/GRU.

Получается красивая модульная архитектура с иерархией, где потоки информации регулируются механизмами внимания. Это в общем уже близко к теории сознания под названием Global Workspace Theory (теория глобального рабочего пространства).

Проводят эксперименты с различными бейзлайнами: варианты LSTM (иерархические, с фидбеком, со вниманием, или со всем сразу), трансформеры, RMC, RIM, иерархический RIM. BRIM берут двуслойный.

На sequential MNIST/CIFAR BRIM генерализует на более высокие разрешения лучше остальных. Из прикольного, чем больше в данных на входе шума, тем больше модель смотрит на верхние модули.

На задаче сложения (идёт поток чисел на вход, затем сигнал какие два из потока нужно сложить результат лучше аналогичной LSTM.

На stochastic moving MNIST (видео 64x64 в котором одно или два символа летают и отскакивают от стенок, и надо предсказать будущие кадры) результат лучше, чем у другой модели.

На отскакивающих мячиках (в том числе с закрывающей часть картинки ширмой) BRIM даёт результат лучше, чем LSTM или RIM.

Ну и языковое моделирование тоже лучше.

В RL на десяти первых по алфавиту Атари играх работает лучше LSTM/RIM.

В общем с точки зрения перформанса норм.

Важность внимания и модульности показали ещё в RIM. Здесь дополнительно из абляций показали, что двунаправленность и иерархия также очень важны.

Coordination Among Neural Modules Through a Shared Global Workspace
Anirudh Goyal, Aniket Didolkar, Alex Lamb, Kartikeya Badola, Nan Rosemary Ke, Nasim Rahaman, Jonathan Binas, Charles Blundell, Michael Mozer, Yoshua Bengio
Статья: https://arxiv.org/abs/2103.01197

И ещё одна релевантная статья в значительной степени от тех же авторов, что и про RIM/BRIM. На этот раз суперсвежая, от 1 марта.

В современном DL есть большое движение в сторону более богатых и модульных репрезентаций. С RIM это вообще очевидно, но RIM пока не массовая архитектура. Есть активное движение к архитектуре Mixture of Experts (MoE), применявшейся в свежих Switch Transformer или переводчике M2M-100. Но даже обычные трансформеры по сути сегментируют вход по позициям, где у каждой позиции есть своя репрезентация.

Обычно в таких архитектурах взаимодействия между отдельными репрезентациями (или модулями, где они более выражены) реализованы попарно. Проблема таких попарных взаимодействий в том, что они могут не дать какого-то когерентного интегрированного представления или глобальной координации, а это может повредить задачам, ради которых оно делается.

И тут мы снова возвращаемся к упомянутой в прошлый раз теории глобального рабочего пространства (Global Workspace Theory, GWT), где вводится ограниченный канал (то самое глобальное рабочее пространство), через который взаимодействуют различные модули-специалисты, при этом соревнуясь за доступ к нему.

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

В целом, система составленная из модулей-специалистов с разреженными взаимодействиями обладают хорошими свойствами типа 1) возможности масштабирования (легко добавить новых специалистов); 2) лучшей устойчивости (выбивание части специалистов может не сильно повредить системе); 3) вычислительной эффективности (вычисления производятся в основном локально, что хорошо). Но синхронизация всё равно нужна.

GWT предлагает простой и понятный способ это устроить. Есть глобальная общая репрезентация, называющаяся рабочим пространством или доской. Она может быть модифицирована теоретически любым специалистом и далее рассылается всем специалистам в системе. При этом, чтобы в неё что-то записать, специалист соревнуется с остальными и по факту возникает некая приоритизация сигналов специалистов.

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

Для работы с общим пространством (SW, shared workspace) используется механизм внимания (это обычное multihead dot-product attention как в трансформерах), где элементы этого пространства, как и специалисты, рассматриваются как неупорядоченное множество. SW — это в каком-то смысле память. Скорее рабочая память.

Данная архитектура накручивается как на RIM, так и на трансформеры. В RIM всё и так уже модуляризовано (каждый модуль — это специалист), а в трансформерах специалистом считается каждая отдельная позиция.

По сути предложенная архитектура заменяет имеющиеся в RIM или трансформерах механизмы внимания на новый, работающий через глобальное рабочее пространство очень ограниченного размера (по аналогии с человеческой рабочей памятью авторы считают, что это число меньше 10).

Процесс работы устроен следующим образом:

1. Вход обрабатывается, чтобы получить его репрезентацию каждым специалистом. По сути здесь мы получаем какие-то начальные представления, с которыми дальше работать.

2. Специалисты соревнуются за право записи в глобальное пространство. Этот механизм по сути обеспечивает реагирование только на критически важные сигналы, так чтобы всякая ерунда не захламляла GWT. Соревнования сделано через внимание, Q поступает из памяти (то самое GWT), а K и V от специалистов. Память по результату обновляется либо через soft-attention от всех специалистов (но по факту мы знает, что обычно там доминирует их небольшое число, ибо софтмакс), либо вообще через top-k softmax. Первые модели дальше будут называться SSW (shared workspace with soft competition), вторые HSW (hard attention). Содержимое памяти обновляется через хитрый механизм с гейтами как в RMC (https://arxiv.org/abs/1806.01822).

3. Содержимое памяти бродкастится всем специалистам. Специалисты обновляют своё состояние снова через механизм внимания (генерят свои Q, получают K,V из памяти, прибавляют к своим состояниям взвешенную таким образом через внимание память). После обновления состояния с учётом памяти специалисты вычисляют свои обычные функции (LSTM/GRU или FFN для трансформера).

Заодно как бонус получаем линейную вычислительную сложность от числа специалистов.

Неожиданная приятность другого рода в том, что эта работа ссылается на наш с Мишей Бурцевым Memory Transformer (https://arxiv.org/abs/2006.11527). У нас тоже в одном из вариантов в трансформере был боттлнек в виде памяти небольшого размера.

И дальше собственно эксперименты.

На задачке определения равносторонних треугольников (когда на картинке есть три кластера точек) трансформеры с общим рабочим пространством (особенно HSW) сходятся быстрее обычных и дают более высокое качество. Sparse transformer, кстати, оказывается здесь особенно плох.

Задачки на object tracking из CATER dataset тоже показывают, что трансформеры с SW лучше обычных.

На задачах relational reasoning из CLEVR, где по картинке и текстовому вопросу надо ответить про определённые свойства объектов и их отношения с другими, трансформеры с SW сходятся сильно быстрее обычных и дают более высокое итоговое качество.

На задачах physical reasoning и всё тех же отскакивающих мячиков, что и раньше (надо предсказать будущие кадры видео), RIM + SW лучше всех и это преимущество растёт и с ростом горизонта предсказания, и с ростом числа мячиков.

Была ещё задачка на RL из Starcraft 2, где надо предсказывать будущие наблюдения агентов, RIM + SW тоже бьём и обычный LSTM, и RIM. Что интересно, обычный RIM с попарными взаимодействиями между специалистами, здесь даже хуже LSTM.

В общем хорошая тема, ждём развитий.