Глобальная память для виртуальных узлов: новый метод повышает точность графовых нейросетей
Виртуальные узлы давно используются в графовых нейросетях как простое решение для глобальной коммуникации между узлами. Однако стандартная схема узел-виртуальный узел-узел сжимает весь граф в однородное состояние и рассылает его одинаково всем узлам. Это создает узкое место, ограничивающее выразительность модели.
Новая научная работа, размещенная на arXiv, предлагает способ обойти это ограничение без использования self-attention. Исследователи опираются на анализ Two-Radius и выделяют два ключевых требования к глобальной памяти: она должна быть факторизована на независимо записываемые и читаемые состояния, а также сохранять кратность элементов.
Для выполнения первого требования авторы используют адресуемые слоты cross-attention. Такая архитектура позволяет каждому узлу обращаться к нужному фрагменту памяти, а не получать один и тот же вектор. Это уже улучшает гибкость модели, но, как показано в статье, адресуемости самой по себе недостаточно.
Проблема в том, что softmax attention инвариантен к равномерному дублированию: если умножить вход на повторяющиеся элементы, внимание не изменится. Авторы решают это, вставляя каждый слот запроса как частный ключ и значение. Это восстанавливает отброшенную нормализационную массу и дает инъективное представление мультимножества на ограниченных цветовых доменах.
Благодаря такой конструкции модель способна реализовать 1-WL уточнение, что гарантирует различение неразличимых в обычном подходе структур. На практике предложенный метод показывает улучшения на задачах с учетом кратности, подсчете мотивов и предсказании наборов связей.
Вычислительная сложность алгоритма составляет O(nMd) — линейно по числу узлов и слотов памяти, что делает его применимым к графам среднего размера. Эксперименты подтверждают, что адресуемая и сохраняющая кратность глобальная память справляется с задачами, где стандартные виртуальные узлы дают сбой.
