Repository: pyg-team/pytorch_geometric
Pull request #7465 (2023-05-30 09:57:05)

Add `Performer` to `GPSConv`
### Description

* Most code is adopted and simplified from https://github.com/lucidrains/performer-pytorch/blob/main/performer_pytorch/performer_pytorch.py and https://github.com/rampasek/GraphGPS/blob/main/graphgps/layer/gps_layer.py Should we add some comments to cite? Checked that both of them have the MIT License.
* There is a projection matrix redrawer that redraws the projection matrix in the performer. It seems to be not required so I just put it into the `examples/graph_gps.py`. Not sure where should I put it.
* There are also some other params and kernels to be added to the performer, such as softmax kernel etc. For now I added some TODOs.

### ZINC Example
```python
python3 examples/graph_gps.py --attn_type performer
```
I directly replace the multihead attention with performer fast attention with the `ReLU` kernel. The performance is similar (there are some randomness for each run and for the graph classification task) Following is the `test_mae` plot for each epoch.

![training](https://github.com/pyg-team/pytorch_geometric/assets/21955420/50aac085-460c-4293-b74c-966d4857a05e)




codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7465?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7465](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7465?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (af025c7) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/af4360a2ed5490e17b12d06f46bca13e6bbf46e1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (af4360a) will **decrease** coverage by `0.23%`.
> The diff coverage is `98.93%`.

```diff
@@            Coverage Diff             @@
##           master    #7465      +/-   ##
==========================================
- Coverage   91.68%   91.45%   -0.23%     
==========================================
  Files         445      445              
  Lines       24699    24787      +88     
==========================================
+ Hits        22645    22670      +25     
- Misses       2054     2117      +63     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7465?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/conv/gps\_conv.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7465?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvZ3BzX2NvbnYucHk=) | `99.37% <98.93%> (-0.63%)` | :arrow_down: |

... and [17 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7465/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

zechengz: @wsad1 Current plan is to add this `performer` and I am also working on another one called `bigbird`. So, with normal multihead attention we at least have three different global attentions. I also observed that there are also relative a lot of attention methods other than `performer` and `bigbird` that can be useful to add to PyG. So creating the folder `torch/nn/attention` makes sense to me. Thanks!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 61fb5eb54..ae0f0b87c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added Performer to `GPSConv` and remove `attn_dropout` argument from `GPSConv` ([#7465](https://github.com/pyg-team/pytorch_geometric/pull/7465))
 - Enabled `LinkNeighborLoader` to return number of sampled nodes and edges per hop ([#7516](https://github.com/pyg-team/pytorch_geometric/pull/7516))
 - Added the `HM` personalized fashion recommendation dataset ([#7515](https://github.com/pyg-team/pytorch_geometric/pull/7515))
 - Added the `GraphMixer` model ([#7501](https://github.com/pyg-team/pytorch_geometric/pull/7501))
diff --git a/examples/graph_gps.py b/examples/graph_gps.py
index d2a606ef8..4f83934e5 100644
--- a/examples/graph_gps.py
+++ b/examples/graph_gps.py
@@ -1,4 +1,6 @@
+import argparse
 import os.path as osp
+from typing import Any, Dict, Optional
 
 import torch
 from torch.nn import (
@@ -15,6 +17,7 @@ import torch_geometric.transforms as T
 from torch_geometric.datasets import ZINC
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import GINEConv, GPSConv, global_add_pool
+from torch_geometric.nn.attention import PerformerAttention
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ZINC-PE')
 transform = T.AddRandomWalkPE(walk_length=20, attr_name='pe')
@@ -26,9 +29,16 @@ train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
 val_loader = DataLoader(val_dataset, batch_size=64)
 test_loader = DataLoader(test_dataset, batch_size=64)
 
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--attn_type', default='multihead',
+    help="Global attention type such as 'multihead' or 'performer'.")
+args = parser.parse_args()
+
 
 class GPS(torch.nn.Module):
-    def __init__(self, channels: int, pe_dim: int, num_layers: int):
+    def __init__(self, channels: int, pe_dim: int, num_layers: int,
+                 attn_type: str, attn_kwargs: Dict[str, Any]):
         super().__init__()
 
         self.node_emb = Embedding(28, channels - pe_dim)
@@ -43,7 +53,8 @@ class GPS(torch.nn.Module):
                 ReLU(),
                 Linear(channels, channels),
             )
-            conv = GPSConv(channels, GINEConv(nn), heads=4, attn_dropout=0.5)
+            conv = GPSConv(channels, GINEConv(nn), heads=4,
+                           attn_type=attn_type, attn_kwargs=attn_kwargs)
             self.convs.append(conv)
 
         self.mlp = Sequential(
@@ -53,6 +64,9 @@ class GPS(torch.nn.Module):
             ReLU(),
             Linear(channels // 4, 1),
         )
+        self.redraw_projection = RedrawProjection(
+            self.convs,
+            redraw_interval=1000 if attn_type == 'performer' else None)
 
     def forward(self, x, pe, edge_index, edge_attr, batch):
         x_pe = self.pe_norm(pe)
@@ -65,8 +79,32 @@ class GPS(torch.nn.Module):
         return self.mlp(x)
 
 
+class RedrawProjection:
+    def __init__(self, model: torch.nn.Module,
+                 redraw_interval: Optional[int] = None):
+        self.model = model
+        self.redraw_interval = redraw_interval
+        self.num_last_redraw = 0
+
+    def redraw_projections(self):
+        if not self.model.training or self.redraw_interval is None:
+            return
+        if self.num_last_redraw >= self.redraw_interval:
+            fast_attentions = [
+                module for module in self.model.modules()
+                if isinstance(module, PerformerAttention)
+            ]
+            for fast_attention in fast_attentions:
+                fast_attention.redraw_projection_matrix()
+            self.num_last_redraw = 0
+            return
+        self.num_last_redraw += 1
+
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = GPS(channels=64, pe_dim=8, num_layers=10).to(device)
+attn_kwargs = {'dropout': 0.5}
+model = GPS(channels=64, pe_dim=8, num_layers=10, attn_type=args.attn_type,
+            attn_kwargs=attn_kwargs).to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
 scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=20,
                               min_lr=0.00001)
@@ -79,6 +117,7 @@ def train():
     for data in train_loader:
         data = data.to(device)
         optimizer.zero_grad()
+        model.redraw_projection.redraw_projections()
         out = model(data.x, data.pe, data.edge_index, data.edge_attr,
                     data.batch)
         loss = (out.squeeze() - data.y).abs().mean()
diff --git a/torch_geometric/nn/attention/__init__.py b/torch_geometric/nn/attention/__init__.py
new file mode 100644
index 000000000..947d58501
--- /dev/null
+++ b/torch_geometric/nn/attention/__init__.py
@@ -0,0 +1,3 @@
+from .performer import PerformerAttention
+
+__all__ = ['PerformerAttention']
diff --git a/torch_geometric/nn/attention/performer.py b/torch_geometric/nn/attention/performer.py
new file mode 100644
index 000000000..b82b51f8e
--- /dev/null
+++ b/torch_geometric/nn/attention/performer.py
@@ -0,0 +1,199 @@
+import math
+from typing import Callable, Optional
+
+import torch
+from torch import Tensor
+
+
+def _orthogonal_matrix(dim: int) -> Tensor:
+    r"""Get an orthogonal matrix by applying QR decomposition."""
+    # Random matrix from normal distribution
+    mat = torch.randn((dim, dim))
+    # QR decomposition to two orthogonal matrices
+    q, _ = torch.linalg.qr(mat.cpu(), mode='reduced')
+    return q.t()
+
+
+def orthogonal_matrix(num_rows: int, num_cols: int) -> Tensor:
+    r"""Generate an orthogonal matrix with `num_rows` rows
+    and `num_cols` columns.
+    """
+    num_full_blocks = int(num_rows / num_cols)
+    blocks = []
+    for _ in range(num_full_blocks):
+        q = _orthogonal_matrix(num_cols)
+        blocks.append(q)
+    remain_rows = num_rows - num_full_blocks * num_cols
+    if remain_rows > 0:
+        q = _orthogonal_matrix(num_cols)
+        blocks.append(q[:remain_rows])
+    mat = torch.cat(blocks)
+    # multiplier = torch.randn((num_rows, num_cols)).norm(dim=1)
+    # scaler = torch.diag(multiplier)
+    # mat = scaler @ mat
+    return mat
+
+
+def linear_attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+    r"""Efficient attention mechanism from the
+    `"Rethinking Attention with Performers"
+    <https://arxiv.org/abs/2009.14794>`_ paper.
+
+    .. math::
+        \mathbf{\hat{D}}^{-1}(\mathbf{Q}'((\mathbf{K}')^{\top} \mathbf{V}))
+
+    """
+    k_contract = k.sum(dim=-2)
+    D_inv = 1.0 / torch.einsum('...Lr,...r->...L', q, k_contract)
+    kv = torch.einsum('...Lr,...Ld->...rd', k, v)
+    qkv = torch.einsum('...Lr,...rd->...Ld', q, kv)
+    out = torch.einsum('...L,...Ld->...Ld', D_inv, qkv)
+    return out
+
+
+def generalized_kernel(x: Tensor, mat: Tensor,
+                       kernel: Callable = torch.nn.ReLU(),
+                       epsilon: float = 0.001) -> Tensor:
+    r"""Apply generalized kernelizable attention with
+    kernel functions such as the ReLU.
+    """
+    num_batches, num_heads, *_ = x.shape
+    # Expand projection matrix to number of batches and number of heads
+    projection = mat.expand(num_batches, num_heads, *mat.shape)
+    # "Inner" product x with projection matrix
+    x = torch.einsum('...id,...jd->...ij', x, projection)
+    out = kernel(x) + epsilon
+    return out
+
+
+class PerformerProjection(torch.nn.Module):
+    r"""The fast attention that uses a projection matrix
+    from the `"Rethinking Attention with Performers"
+    <https://arxiv.org/abs/2009.14794>`_ paper. This class
+    projects :math:`\mathbf{Q}` and :math:`\mathbf{K}` matrices
+    with specified kernel.
+
+    Args:
+        num_cols (int): Projection matrix number of columns.
+        kernel (Callable, optional): Kernels for generalized attention.
+            If not specified, `ReLU` kernel will be used.
+            (default: :obj:`torch.nn.ReLU()`)
+    """
+    def __init__(self, num_cols: int, kernel: Callable = torch.nn.ReLU()):
+        super().__init__()
+        num_rows = int(num_cols * math.log(num_cols))
+        self.num_rows = num_rows
+        self.num_cols = num_cols
+        # Generate an orthogonal projection matrix
+        # with the shape (num_rows, num_cols)
+        projection_matrix = orthogonal_matrix(self.num_rows, self.num_cols)
+        self.register_buffer('projection_matrix', projection_matrix)
+        assert kernel is not None
+        self.kernel = kernel
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        q = generalized_kernel(q, self.projection_matrix, self.kernel)
+        k = generalized_kernel(k, self.projection_matrix, self.kernel)
+        out = linear_attention(q, k, v)
+        return out
+
+
+class PerformerAttention(torch.nn.Module):
+    r"""The linear scaled attention mechanism from the
+    `"Rethinking Attention with Performers"
+    <https://arxiv.org/abs/2009.14794>`_ paper.
+
+    Args:
+        channels (int): Size of each input sample.
+        heads (int, optional): Number of parallel attention heads.
+        head_channels (int, optional): Size of each attention head.
+            (default: :obj:`64.`)
+        kernel (Callable, optional): Kernels for generalized attention.
+            If not specified, `ReLU` kernel will be used.
+            (default: :obj:`torch.nn.ReLU()`)
+        qkv_bias (bool, optional): If specified, add bias to query, key
+            and value in the self attention. (default: :obj:`False`)
+        attn_out_bias (bool, optional): If specified, add bias to the
+            attention output. (default: :obj:`True`)
+        dropout (float, optional): Dropout probability of the final
+            attention output. (default: :obj:`0.0`)
+
+    """
+    def __init__(
+        self,
+        channels: int,
+        heads: int,
+        head_channels: int = 64,
+        kernel: Callable = torch.nn.ReLU(),
+        qkv_bias: bool = False,
+        attn_out_bias: bool = True,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        assert channels % heads == 0
+        if head_channels is None:
+            head_channels = channels // heads
+
+        self.heads = heads
+        self.head_channels = head_channels
+        self.kernel = kernel
+        self.fast_attn = PerformerProjection(head_channels, kernel)
+
+        inner_channels = head_channels * heads
+        self.q = torch.nn.Linear(channels, inner_channels, bias=qkv_bias)
+        self.k = torch.nn.Linear(channels, inner_channels, bias=qkv_bias)
+        self.v = torch.nn.Linear(channels, inner_channels, bias=qkv_bias)
+        self.attn_out = torch.nn.Linear(inner_channels, channels,
+                                        bias=attn_out_bias)
+        self.dropout = torch.nn.Dropout(dropout)
+
+    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
+        r"""
+
+        Args:
+            x (torch.Tensor): Node feature tensor
+                :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with
+                batch-size :math:`B`, (maximum) number of nodes :math:`N` for
+                each graph, and feature dimension :math:`F`.
+            mask (torch.Tensor, optional): Mask matrix
+                :math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}` indicating
+                the valid nodes for each graph. (default: :obj:`None`)
+        """
+        B, N, *_ = x.shape
+        q, k, v = self.q(x), self.k(x), self.v(x)
+        # Reshape and permute q, k and v to proper shape
+        # (B, N, num_heads * head_channels) to (b, num_heads, n, head_channels)
+        q, k, v = map(
+            lambda t: t.reshape(B, N, self.heads, self.head_channels).permute(
+                0, 2, 1, 3), (q, k, v))
+        if mask is not None:
+            mask = mask[:, None, :, None]
+            v.masked_fill_(~mask, 0.)
+        out = self.fast_attn(q, k, v)
+        out = out.permute(0, 2, 1, 3).reshape(B, N, -1)
+        out = self.attn_out(out)
+        out = self.dropout(out)
+        return out
+
+    @torch.no_grad()
+    def redraw_projection_matrix(self):
+        r"""As described in the paper, periodically redraw
+        examples to improve overall approximation of attention."""
+        num_rows = self.fast_attn.num_rows
+        num_cols = self.fast_attn.num_cols
+        projection_matrix = orthogonal_matrix(num_rows, num_cols)
+        self.fast_attn.projection_matrix.copy_(projection_matrix)
+        del projection_matrix
+
+    def _reset_parameters(self):
+        self.q.reset_parameters()
+        self.k.reset_parameters()
+        self.v.reset_parameters()
+        self.attn_out.reset_parameters()
+        self.redraw_projection_matrix()
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}('
+                f'heads={self.heads}, '
+                f'head_channels={self.head_channels} '
+                f'kernel={self.kernel})')
diff --git a/torch_geometric/nn/conv/gps_conv.py b/torch_geometric/nn/conv/gps_conv.py
index 1b5787d44..c967e10a8 100644
--- a/torch_geometric/nn/conv/gps_conv.py
+++ b/torch_geometric/nn/conv/gps_conv.py
@@ -6,6 +6,7 @@ import torch.nn.functional as F
 from torch import Tensor
 from torch.nn import Dropout, Linear, Sequential
 
+from torch_geometric.nn.attention import PerformerAttention
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.inits import reset
 from torch_geometric.nn.resolver import (
@@ -43,8 +44,6 @@ class GPSConv(torch.nn.Module):
             (default: :obj:`1`)
         dropout (float, optional): Dropout probability of intermediate
             embeddings. (default: :obj:`0.`)
-        attn_dropout (float, optional): Dropout probability of the normalized
-            attention coefficients. (default: :obj:`0`)
         act (str or Callable, optional): The non-linear activation function to
             use. (default: :obj:`"relu"`)
         act_kwargs (Dict[str, Any], optional): Arguments passed to the
@@ -55,6 +54,10 @@ class GPSConv(torch.nn.Module):
         norm_kwargs (Dict[str, Any], optional): Arguments passed to the
             respective normalization function defined by :obj:`norm`.
             (default: :obj:`None`)
+        attn_type (str): Global attention type, :obj:`multihead` or
+            :obj:`performer`. (default: :obj:`multihead`)
+        attn_kwargs (Dict[str, Any], optional): Arguments passed to the
+            attention layer. (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -62,11 +65,12 @@ class GPSConv(torch.nn.Module):
         conv: Optional[MessagePassing],
         heads: int = 1,
         dropout: float = 0.0,
-        attn_dropout: float = 0.0,
         act: str = 'relu',
         act_kwargs: Optional[Dict[str, Any]] = None,
         norm: Optional[str] = 'batch_norm',
         norm_kwargs: Optional[Dict[str, Any]] = None,
+        attn_type: str = 'multihead',
+        attn_kwargs: Optional[Dict[str, Any]] = None,
     ):
         super().__init__()
 
@@ -74,13 +78,25 @@ class GPSConv(torch.nn.Module):
         self.conv = conv
         self.heads = heads
         self.dropout = dropout
-
-        self.attn = torch.nn.MultiheadAttention(
-            channels,
-            heads,
-            dropout=attn_dropout,
-            batch_first=True,
-        )
+        self.attn_type = attn_type
+
+        attn_kwargs = attn_kwargs or {}
+        if attn_type == 'multihead':
+            self.attn = torch.nn.MultiheadAttention(
+                channels,
+                heads,
+                batch_first=True,
+                **attn_kwargs,
+            )
+        elif attn_type == 'performer':
+            self.attn = PerformerAttention(
+                channels=channels,
+                heads=heads,
+                **attn_kwargs,
+            )
+        else:
+            # TODO: Support BigBird
+            raise ValueError(f'{attn_type} is not supported')
 
         self.mlp = Sequential(
             Linear(channels, channels * 2),
@@ -135,7 +151,13 @@ class GPSConv(torch.nn.Module):
 
         # Global attention transformer-style model.
         h, mask = to_dense_batch(x, batch)
-        h, _ = self.attn(h, h, h, key_padding_mask=~mask, need_weights=False)
+
+        if isinstance(self.attn, torch.nn.MultiheadAttention):
+            h, _ = self.attn(h, h, h, key_padding_mask=~mask,
+                             need_weights=False)
+        elif isinstance(self.attn, PerformerAttention):
+            h = self.attn(h, mask=mask)
+
         h = h[mask]
         h = F.dropout(h, p=self.dropout, training=self.training)
         h = h + x  # Residual connection.
@@ -159,4 +181,5 @@ class GPSConv(torch.nn.Module):
 
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}({self.channels}, '
-                f'conv={self.conv}, heads={self.heads})')
+                f'conv={self.conv}, heads={self.heads}, '
+                f'attn_type={self.attn_type})')


Test patch:
diff --git a/test/nn/attention/test_performer.py b/test/nn/attention/test_performer.py
new file mode 100644
index 000000000..c0b45940e
--- /dev/null
+++ b/test/nn/attention/test_performer.py
@@ -0,0 +1,13 @@
+import torch
+
+from torch_geometric.nn.attention import PerformerAttention
+
+
+def test_performer_attention():
+    x = torch.randn(1, 4, 16)
+    mask = torch.ones([1, 4], dtype=torch.bool)
+    attn = PerformerAttention(channels=16, heads=4)
+    out = attn(x, mask)
+    assert out.shape == (1, 4, 16)
+    assert str(attn) == ('PerformerAttention(heads=4, '
+                         'head_channels=64 kernel=ReLU())')
diff --git a/test/nn/conv/test_gps_conv.py b/test/nn/conv/test_gps_conv.py
index eda1a4889..1e063f2e4 100644
--- a/test/nn/conv/test_gps_conv.py
+++ b/test/nn/conv/test_gps_conv.py
@@ -6,18 +6,20 @@ from torch_geometric.nn import GPSConv, SAGEConv
 from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import to_torch_csc_tensor
 
+x = torch.randn(4, 16)
+edge_index = torch.tensor([[0, 1, 2, 3], [1, 0, 3, 2]])
+batch = torch.tensor([0, 0, 1, 1])
+adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
-@pytest.mark.parametrize('norm', [None, 'batch_norm', 'layer_norm'])
-def test_gps_conv(norm):
-    x = torch.randn(4, 16)
-    edge_index = torch.tensor([[0, 1, 2, 3], [1, 0, 3, 2]])
-    batch = torch.tensor([0, 0, 1, 1])
-    adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
-    conv = GPSConv(16, conv=SAGEConv(16, 16), heads=4, norm=norm)
+@pytest.mark.parametrize('attn_type', ['multihead', 'performer'])
+@pytest.mark.parametrize('norm', [None, 'batch_norm', 'layer_norm'])
+def test_gps_conv(norm, attn_type):
+    conv = GPSConv(16, conv=SAGEConv(16, 16), heads=4, norm=norm,
+                   attn_type=attn_type)
     conv.reset_parameters()
-    assert str(conv) == ('GPSConv(16, conv=SAGEConv(16, 16, aggr=mean), '
-                         'heads=4)')
+    assert str(conv) == (f'GPSConv(16, conv=SAGEConv(16, 16, aggr=mean), '
+                         f'heads=4, attn_type={attn_type})')
 
     out = conv(x, edge_index)
     assert out.size() == (4, 16)