Команда исследователей из ШАДа Яндекса ускорила работу графовых нейросетей
Приближается ICML — престижная международная конференция по машинному обучению, которая в этом году пройдет 6–11 июля в Корее. Статус Spotlight присваивают работам, получившим самые высокие оценки программного комитета. Одну из них написали преподаватели и студенты ШАД.

О чем статья?
Графовые нейросети — популярный класс моделей для работы со сложными взаимосвязанными данными: от транспортных сетей и биологических систем до рекомендательных систем и антифрод-решений.

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

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

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

В результате авторам удалось ускорить выполнение отдельных операций до 10 раз и сократить пиковое потребление памяти до 70 раз.

🔜 Полный текст уже доступен на arXiv, а весь код — на GitHub.

Поздравляем авторов, а особенно студентов — такой результат и такой опыт точно пригодятся им в будущем. 🔥