Repository: pyg-team/pytorch_geometric

Issue:
GAT model - last layer incorrect?
### 🐛 Describe the bug

In GAT paper last layer averages attention heads instead of concatenation, which corresponds to `concat=False` in `GATConv` in PyTorchGeometric. This is also the case in all examples.

However, in `GAT` model the last layer is still concatenated, as only the convolution is passed to `BasicGNN`. This means that using `GAT`, we actually concatenate the attention heads in the last layer, which is incorrect, or at least an unexpected behavior.

If I'm correct, I see two ways of fixing this:
- leave as-is, but add a disclaimer to the docs, possibly with example how to go from the current output `(num_nodes, K * out_channels)` (for `K` attention heads) to `(num_nodes, out_channels)`, where attention heads are averaged
- override the `.forward()` method in `GAT` and add `concat_last_layer=False` option:
  - if `True`, run as-is
  - if `False`, override the `.forward()` method, use the parent's `.forward()` up to `N-1` layer, and run the last layer separately, using concat=False`

### Environment

* PyG version: 2.0.4
* PyTorch version: 1.10
* OS: Windows 10
* Python version: 3.9
* CUDA/cuDNN version: 11.3
* How you installed PyTorch and PyG (`conda`, `pip`, source): `pip`
* Any other relevant information (*e.g.*, version of `torch-scatter`):



Discussion:
That is interesting, I think you are right that we need to fix the `concat` option in the last layer for `GAT`. In my understanding, we only need to fix this in case `GAT` makes use of the `out_channels` argument, i.e. the last `GATConv` layer actually maps to `out_channels`. In that case, it might be the easiest to just fix `concat=False`.
I think something like this would work:
```
class GAT(BasicGNN):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        kwargs = copy.copy(kwargs)
        kwargs["concat"] = False
        if out_channels is not None and jk is None:
            self.convs[-1] = self.init_conv(hidden_channels, out_channels, **kwargs)
        else:
            self.convs[-1] = self.init_conv(hidden_channels, hidden_channels, **kwargs)
```

So we always make sure that the last `GATConv` layer uses `concat=False`. This fixes the parameter, but according to the paper this makes sense, and previous layers can still use either concatenation or summation.

If this seems correct, I can make a PR.

Resolving patch:
diff --git a/torch_geometric/nn/models/basic_gnn.py b/torch_geometric/nn/models/basic_gnn.py
--- a/torch_geometric/nn/models/basic_gnn.py
+++ b/torch_geometric/nn/models/basic_gnn.py
@@ -86,6 +86,7 @@ def __init__(
             self.convs.append(
                 self.init_conv(hidden_channels, hidden_channels, **kwargs))
         if out_channels is not None and jk is None:
+            self._is_conv_to_out = True
             self.convs.append(
                 self.init_conv(hidden_channels, out_channels, **kwargs))
         else:
@@ -282,14 +283,24 @@ class GAT(BasicGNN):
     def init_conv(self, in_channels: int, out_channels: int,
                   **kwargs) -> MessagePassing:
 
-        kwargs = copy.copy(kwargs)
-        if 'heads' in kwargs and out_channels % kwargs['heads'] != 0:
-            kwargs['heads'] = 1
-        if 'concat' not in kwargs or kwargs['concat']:
-            out_channels = out_channels // kwargs.get('heads', 1)
+        heads = kwargs.pop('heads', 1)
+        concat = kwargs.pop('concat', True)
 
-        return GATConv(in_channels, out_channels, dropout=self.dropout,
-                       **kwargs)
+        # Do not use concatenation in case the layer `GATConv` layer maps to
+        # the desired output channels (out_channels != None and jk != None):
+        if getattr(self, '_is_conv_to_out', False):
+            concat = False
+
+        if concat and out_channels % heads != 0:
+            raise ValueError(f"Ensure that the number of output channels of "
+                             f"'GATConv' (got '{out_channels}') is divisible "
+                             f"by the number of heads (got '{heads}')")
+
+        if concat:
+            out_channels = out_channels // heads
+
+        return GATConv(in_channels, out_channels, heads=heads, concat=concat,
+                       dropout=self.dropout, **kwargs)
 
 
 class PNA(BasicGNN):


Test patch: