Spotlight на ICML: ускоряем графовые нейросети

Начинается ежегодный марафон по чтению статей с ICML 2026, которая в этом году пройдет 6-11 июля. Приятно, что попадается много инженерных статей с практическими результатами.

Среди таких – работа Яндекса и ШАДа. Она получила статус Spotlight: это статьи, которым программный комитет дал самые высокие оценки. Таких работ на конференции всего 536, 2.2% от всех поданных заявок.

В статье исследователи из Yandex Research, команд ML-инфраструктуры и инференса и студенты ШАД работали над ускорением Graph Neural Networks. Дело в том, что вычисления на графах неэффективно ложаться на парадигму GPU-вычислений, данные лежат в памяти неравномерно, и поэтому видеокарта значительную часть времени ждет загрузки данных вместо того, чтобы выполнять вычисления, которые представляют собой достаточно простые операции. Как раз из-за этого скорость работы ограничена чтением и записью данных, и грамотная работа с GPU-памятью — это главный источник ускорения таких операций.

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

Они разработали новые готовые GPU-кернелы для популярных семейств графовых нейросетей, и добились потрясающего ускорения в 3–10 раз в зависимости от архитектуры. Потребление памяти в отдельных сценариях сокращается более чем на порядок.

Поздравляем команду с заслуженным Spotlight и прекрасной работой!

Статья уже доступна на arXiv, а код выложили вот здесь.