adding adjacency matrix normalization

Author: flyIchtus

models/src/anemoi/models/layers/block.py

@@ -12,7 +12,7 @@
 import os
 from abc import ABC
 from abc import abstractmethod
-from typing import Optional
+from typing import Literal, Optional
 from typing import Union
 
 import einops
@@ -443,6 +443,7 @@ def __init__(
         edge_dim: int,
         bias: bool = True,
         qk_norm: bool = False,
+        adj_norm: Literal['sym','rw'] | None = None,
         update_src_nodes: bool = False,
         layer_kernels: DotDict,
         graph_attention_backend: str = "triton",
@@ -464,6 +465,8 @@ def __init__(
             Add bias or not
         qk_norm : bool, by default False
             Normalize query and key
+        adj_norm : Literal['sym','rw'] | None
+            Normalize adjacency aggregation: D^-1A ('rw') or D^{-1/2}AD^{-1/2} ('sym')
         update_src_nodes: bool, by default False
             Update src if src and dst nodes are given
         layer_kernels : DotDict
@@ -479,6 +482,7 @@ def __init__(
         self.out_channels_conv = out_channels // num_heads
         self.num_heads = num_heads
         self.qk_norm = qk_norm
+        self.adj_norm = adj_norm
 
         Linear = layer_kernels.Linear
         LayerNorm = layer_kernels.LayerNorm
@@ -513,7 +517,7 @@ def __init__(
             self.conv = GraphTransformerFunction.apply
         else:
             LOGGER.warning(f"{self.__class__.__name__} using pyg graph attention backend, consider using 'triton'.")
-            self.conv = GraphTransformerConv(out_channels=self.out_channels_conv)
+            self.conv = GraphTransformerConv(out_channels=self.out_channels_conv,adj_norm=self.adj_norm)
 
     def run_node_dst_mlp(self, x, **layer_kwargs):
         return self.node_dst_mlp(self.layer_norm_mlp_dst(x, **layer_kwargs))

models/src/anemoi/models/layers/conv.py

@@ -8,7 +8,7 @@
 # nor does it submit to any jurisdiction.
 
 
-from typing import Optional
+from typing import Literal, Optional
 
 import torch
 from torch import Tensor
@@ -18,6 +18,7 @@
 from torch_geometric.typing import OptPairTensor
 from torch_geometric.typing import OptTensor
 from torch_geometric.typing import Size
+from torch_geometric.utils import degree
 from torch_geometric.utils import scatter
 from torch_geometric.utils import softmax
 
@@ -83,17 +84,23 @@ class GraphTransformerConv(MessagePassing):
 
     Adapted from 'Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification'
     (https://arxiv.org/abs/2009.03509)
+    
+    Edge normalization taken from  'Semi-Supervised Classification with Graph Convolutional Networks'(https://arxiv.org/abs/1609.02907)
+    Code inspired from https://pytorch-geometric.readthedocs.io/en/2.6.0/_modules/torch_geometric/utils/laplacian.html#get_laplacian
     """
 
     def __init__(
         self,
         out_channels: int,
         dropout: float = 0.0,
+        adj_norm: Literal["sym", "rw"] | None = None,
         **kwargs,
     ):
         kwargs.setdefault("aggr", "add")
         super().__init__(node_dim=0, **kwargs)
 
+        self.adj_norm = adj_norm
+
         self.out_channels = out_channels
         self.dropout = dropout
 
@@ -108,6 +115,19 @@ def forward(
     ):
         dim_size = query.shape[0]
         heads = query.shape[1]
+        
+        edge_weights = torch.ones(edge_index.size(1), dtype = query.dtype, device = query.device)
+
+        if self.adj_norm is not None:
+            row, col = edge_index
+            deg = degree(col,dtype=query.dtype)
+            
+            if self.adj_norm=="sym":
+                deg_inv_sqrt = deg.pow_(-0.5)
+                deg_inv_sqrt.masked_fill_(deg_inv_sqrt==float('inf'),0)
+                edge_weights = (deg_inv_sqrt[row] * deg_inv_sqrt[col])
+            elif self.adj_norm=='rw':
+                edge_weights = (deg.pow_(-1.0)[row] * edge_weights)
 
         out = self.propagate(
             edge_index=edge_index,
@@ -118,8 +138,9 @@ def forward(
             query=query,
             key=key,
             value=value,
+            edge_weights = edge_weights.repeat(1,heads)
         )
-
+        
         return out
 
     def message(
@@ -128,6 +149,7 @@ def message(
         query_i: Tensor,
         key_j: Tensor,
         value_j: Tensor,
+        edge_weights: Tensor,
         edge_attr: OptTensor,
         index: Tensor,
         ptr: OptTensor,
@@ -141,4 +163,4 @@ def message(
         alpha = softmax(alpha, index, ptr, size_i)
         alpha = dropout(alpha, p=self.dropout, training=self.training)
 
-        return (value_j + edge_attr) * alpha.view(-1, heads, 1)
+        return edge_weights.view(-1,heads,1) * (value_j + edge_attr) * alpha.view(-1, heads, 1)

models/src/anemoi/models/schemas/processor.py

@@ -42,6 +42,8 @@ class GraphTransformerProcessorSchema(TransformerModelComponent):
     "Number of chunks to divide the layer into. Default to 2."
     qk_norm: bool = Field(example=False)
     "Normalize the query and key vectors. Default to False."
+    adj_norm: Literal['sym','rw'] | None =Field(example=None)
+    "Normalize adjacency aggregation: D^-1A ('rw') or D^{-1/2}AD^{-1/2} ('sym')"
 
     @model_validator(mode="after")
     def check_valid_extras(self) -> Any:

models/tests/layers/block/test_block_graphtransformer.py

@@ -68,6 +68,7 @@ def block(init_proc):
         bias=bias,
         update_src_nodes=False,
         qk_norm=qk_norm,
+        adj_norm=None,
         graph_attention_backend=graph_attention_backend,
     )
 

training/src/anemoi/training/config/model/graphtransformer.yaml

@@ -32,6 +32,7 @@ processor:
   qk_norm: False
   cpu_offload: ${model.cpu_offload}
   layer_kernels: ${model.layer_kernels}
+  adj_norm: null
   graph_attention_backend: "triton"  # Options: "triton", "pyg"
 
 encoder: