pyg-team/pytorch_geometric

Local multi-headed self-attention

Open

#8,972 opened on 2024年2月26日

GitHub で見る
 (3 comments) (1 reaction) (0 assignees)Python (19,985 stars) (3,514 forks)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
Issue 種別
feature
難度新しい貢献者にとっての実装難度です。1 は非常に小さな変更、5 は専門的な作業です。
3
推定時間経験ある貢献者が調査、実装、テスト、pull request 準備にかけるおおよその時間範囲です。
3-5 days
活動状況Issue が今どれだけ取り組みやすいかを示します。新しい、活発、古い、ブロック中、またはメンテナー入力待ちなどです。
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.