Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs Сегодня смотрим реальный трансформер для молекул c SE(3)/E(3) эквивариатностью ☯️ Equiformer - разработка группы Atomic Architects из MIT ⚡️ Очень крутые ребята! Работает он на equivariant модулях с irreducible representations от сферических функций (обозначаются как type-L вектора, где L - степень репрезентаций). Такое в целом мы уже видели в других работах, но тут красота статьи в деталях Сначала рассмотрим любимый аттеншен Фичи с source и target ноды подают в линейный слой, а потом складывают и умножают тензорно на эмбеддинги сферической функции относительной позиции атомов. Это операция сохраняет семантические и геометрические признаки. Для аттеншена (весов) юзают только скалярный вектора так как веса аттеншена инвариантны. Вместо дот продакта берут MLP как в GATv2 для более гибких паттернов. Сами фичи после преобразований нод разбивают на скаляры и на другие. По скалярам считают веса, а по другим делают фичи вершин через эквивариантную активацию, тензорное умножение и линейный слой. Потом перезвешивают и все это можно делать как MHA Еще есть другие мелкие детали касательно разных слоев 1) Линейные слои обрабатывают различные type-L вектора отдельно + не используют сдвиг, так как они сломают эквивариантность при L > 0 2) LN нормализуется как RMS L2-нормы type-L вектора. Также нет сдвигов и средних для L > 0 3) Используют более эффективное тензороное умножение, где каждый тип output вектора зависит от одного типа input вектора. Но вообще не понятно, почему на визуализациях type-0 зависит только от type-0 и type-1, но не от type-2 На экспах классно показывает себя на OC20, а вот на QM9 средне, но в целом это маленький датасет, поэтому пофиг 👀 LINK