pyg-team/pytorch_geometric

Local multi-headed self-attention

Open

#8,972 opened on Feb 26, 2024

View on GitHub
 (3 comments) (1 reaction) (0 assignees)Python (19,985 stars) (3,514 forks)batch import
featurehelp wanted

Description

🚀 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

Contributor guide

Tech stack
python
Domain
machine learning
Issue type
feature
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
3-5 days
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
stale
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
clear
Prerequisites
PyTorch basicsAttention mechanism understandingPyTorch Geometric conventions
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
50
Research direction
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.