pyg-team/pytorch_geometric

Local multi-headed self-attention

Open

#8,972 建立於 2024年2月26日

在 GitHub 查看
 (3 留言) (1 反應) (0 負責人)Python (19,985 star) (3,514 fork)batch import
featurehelp wanted

描述

🚀 The feature, motivation and pitch

I am unable to find the clean implementation of local multi-headed self-attention in pytorch geometric. I found three types of multi-head attention, one TransformerConv (https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.TransformerConv.html#torch_geometric.nn.conv.TransformerConv). But this one calculates a linear combination of all features with different attention weights as opposed to dividing features into multiple heads and taking their linear combination: another RGATConv in the similar direction (https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.RGATConv.html). And finally GPSConv (https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.GPSConv.html) that does multi-head attention but is global.

Alternatives

I think it is nice to have the implementation of local self-attention with multiple heads where each head looks into a part of the feature dimension.

Additional context

No response

貢獻者指南

技術棧
python
領域
machine learning
議題類型
feature
難度面向新貢獻者的預計實作難度，1 表示很小改動，5 表示專家級工作。
3
預計時間有經驗貢獻者完成調查、實作、測試並準備 pull request 的粗略時間範圍。
3-5 days
活動狀態議題目前的可參與程度：新鮮、活躍、陳舊、阻塞或等待維護者輸入。
stale
清晰度議題是否清楚說明預期改動、驗收標準和下一步。
clear
前置要求
PyTorch basicsAttention mechanism understandingPyTorch Geometric conventions
新手友善度1-100 的估計分數，表示該議題對首次貢獻者的友善程度。
50
研究方向
Study existing implementations in PyTorch Geometric: TransformerConv (torch geometric/nn/conv/transformer conv.py) performs global linear combination of features, while GPSConv (torch geometric/nn/conv/gps conv.py) does global multi head attention. Design a local version that splits features into heads and computes attention within a local neighborhood, using a message passing scheme. Refer to the RGATConv (torch geometric/nn/conv/rgat conv.py) for relational attention but adapt to local self attention without relations. The implementation should be a new convolution class that takes node features and edge index, and returns updated node features with local multi headed attention.