Demystifying Graph Attention Networks (GAT)

A step-by-step TensorFlow implementation

Source: https://arxiv.org/pdf/1710.10903

If you are new to Graph Deep Learning, standard neural networks might feel familiar; you pass data through dense layers through convolutions, and sequences through recurrent units. But what happens when your data points are connected in an intricate web, like social networks, molecules, or time-series data treated as interconnected nodes?

Enter Graph Attention Networks (GAT). GATs allow nodes in a graph to look at their neighbours, identify the most important ones, and assign them higher priority (attention) when updating their own information.

In this post, we will break down a complete, production-ready GAT Autoencoder implementation in TensorFlow/Keras. We will go function by function, ensuring you understand exactly what happens under the hood.

The Architecture Blueprint

Before diving into code, let’s understand the two main components we are building:

1. The Custom GATLayer: A custom Keras layer that calculates how nodes should pay attention to each other and aggregate their features.
2. The GAT Autoencoder Model: A network that compresses the graph structure down to a core bottleneck and reconstructs it back to its original layout

Part 1: Breaking Down the Custom GATLayer

Here is the backbone of our network. Let’s dissect it block by block.

1. The class setup and constructor(init)

@tf.keras.utils.register_keras_serializable(package="CustomLayers", name="GATLayer")
class GATLayer(tf.keras.layers.Layer):
    def __init__(self, units, activation="tanh", kernel_regularizer=None, **kwargs):
        super(GATLayer, self).__init__(**kwargs)
        self.units = units
        self.activation_name = activation
        self.activation = tf.keras.activations.get(activation)
        self.kernel_regularizer = tf.keras.regularizers.get(kernel_regularizer)

The @register_keras_serializable Decorator: This is a crucial piece of infrastructure. It tells Keras, “Hey, remember this class name! If I save this model and load it tomorrow, don’t throw an error, look up this exact Python class.”

init: This initialize our hyperparameters. Units define our output feature dimension size. activation dictates our non-linear thresholding function, and kernel_regularizer prevents our network weights from growing too large (overfitting).

2. Weight Initialisation (build)

    def build(self, input_shape):
        feature_dim = input_shape[-1]

        # Linear feature transformation weight matrix
        self.W = self.add_weight(
            shape=(feature_dim, self.units),
            initializer="glorot_uniform",
            regularizer=self.kernel_regularizer,
            name="linear_transform_W"
        )

        # Attention parameter vector (concatenation mechanism)
        self.a = self.add_weight(
            shape=(2 * self.units, 1),
            initializer="glorot_uniform",
            regularizer=self.kernel_regularizer,
            name="attention_vector_a"
        )
        super(GATLayer, self).build(input_shape)

The build method runs once when the layer encounters its first data input shape. Here, we create our trainable parameters.

Weight Matrix W: Transforms our input node features into a new, optimised mathematical space.

Attention Vector a: A single-column vector used to score how well two nodes relate to one another. Because it scores a pair of nodes together, its size is exactly twice the output dimension size (2 * self.units).

3. The Core Logic Operations (call)

This is where the GAT mathematical magic happens. Let’s split its internal steps up to see how the tensor changes shapes.

Step 3.1 Linear Transformation

    def call(self, X):
        batch_size = tf.shape(X)[0]
        num_nodes = tf.shape(X)[1]

        # 1. Linear Transformation: H = X * W -> Shape: (batch, num_nodes, units)
        h = tf.linalg.matmul(X, self.W)

We project our input X by multiplying it with our weight matrix W. Now every node is represented by a vector of length units.

Step 3.2: Pairing Every Node Combination

  # 2. Construct attention matrix matching every node pair combination
  h_i = tf.repeat(h, repeats=num_nodes, axis=1)
  h_j = tf.tile(h, [1, num_nodes, 1])
  embeddings = tf.concat([h_i, h_j], axis=-1)

To calculate attention across all possible node combinations, we need to compare each node with every other node.

We use tf.repeat and tf.tile to align every possible target node next to every source node.

We use tf.concat to glue their vectors together side-by-side.

Step 3.3: Calculating Attention Scores

  # 3. Compute raw scores and apply LeakyReLU: (batch, num_nodes * num_nodes, 1)
  scores = tf.linalg.matmul(embeddings, self.a)
  scores = tf.nn.leaky_relu(scores, alpha=0.2)
  attention_matrix = tf.reshape(scores, (batch_size, num_nodes, num_nodes))

We multiply our concatenated node pairs by our attention vector a to produce a single raw numerical value.

We pass it through a LeakyReLU activation function to add non-linearity.

We reshape this long vector into a neat square grid (num_nodes x number_nodes).

Step 3.4: Applying Softmax

  # 4. Softmax activation across target neighborhoods to normalize coefficients
  attention_weights = tf.nn.softmax(attention_matrix, axis=-1)

Raw attention scores are hard to interpret. We apply a softmax operation across rows so that the attention coefficients for each node’s neighbourhood add up to exactly 1.0. Think of these as a percentage of importance.

Step 3.5: Neighbourhood Aggregation

We perform a matrix multiplication between our normalised attention_weights and our node features h. Each node updates its own state by collecting a weighted sum of its neighbours’ features. Finally, we run it through our chosen activation function.

4. Saving and Loading Infrastructure (get_config & from_config)

    def get_config(self):
        config = super(GATLayer, self).get_config()
        config.update({
            "units": self.units,
            "activation": self.activation_name,
            "kernel_regularizer": tf.keras.regularizers.serialize(self.kernel_regularizer)
        })
        return config

    @classmethod
    def from_config(cls, config):
        if config.get("kernel_regularizer") is not None:
            config["kernel_regularizer"] = tf.keras.regularizers.deserialize(config["kernel_regularizer"])
        return cls(**config)

get_config: Converts custom settings into a standard dictionary structure during export.

from_config: Reverse the process, safely unpacking serialised sub-configurations (like the regularizer definitions) when the custom layers come back to life.

Part 2: Building the Complete GAT Autoencoder

An Autoencoder takes complex information, compresses it into a tight bottleneck space, and attempts to reconstruct the original data layout exactly. If it can do this, it means it has mastered learning the core relationships within the data.

def build_gat_autoencoder_gpu(window_size: int,
                              n_features: int = 1,
                              enc_units_1: int = 64,
                              enc_units_2: int = 32,
                              dropout_rate: float = 0.2,
                              l2_reg: float = 1e-4) -> tf.keras.Model:

    inp = Input(shape=(window_size, n_features), name="input")

Input: This network models sequential patterns or windowed structures as graphs. Each position in the window_size serves as a node, and n_feautres represents the properties tracked at that specific point.

The Encoder Phase (Compression)

    # GAT Layer 1
    x = GATLayer(units=enc_units_1, activation="tanh", kernel_regularizer=l2(l2_reg), name="enc_gat_1")(inp)
    x = Dropout(dropout_rate, name="enc_drop_1")(x)

    # GAT Layer 2 (Bottleneck Core)
    bottleneck = GATLayer(units=enc_units_2, activation="tanh", kernel_regularizer=l2(l2_reg), name="enc_gat_2")(x)

enc_gat_1: Mixes features across nodes into 64 deep representation vectors.

Dropout: Randomly silences a portion of connections during training. This forces the layer to stay resilient and avoid relying heavily on individual node paths.

enc_gat_2: Compress our features down to a small 32-unit bottleneck representation.

The Decoder Phase (Reconstruction)

    # GAT Layer 3
    x = GATLayer(units=enc_units_2, activation="tanh", kernel_regularizer=l2(l2_reg), name="dec_gat_1")(bottleneck)
    x = Dropout(dropout_rate, name="dec_drop_1")(x)

    # GAT Layer 4
    x = GATLayer(units=enc_units_1, activation="tanh", kernel_regularizer=l2(l2_reg), name="dec_gat_2")(x)

    # Reconstruction projection block 
    out = TimeDistributed(Dense(n_features), name="reconstruction")(x)

The decoder layers expand our graph representations back to their original sizes step by step.

TimeDistributed(Dense(n_features)): This wraps a standard dense layer around every individual node independently, transforming the hidden internal features back to match the original n_features target design.

Compiling the System

    model = tf.keras.Model(inputs=inp, outputs=out, name="GAT_autoencoder_gpu")
    model.compile(optimizer=Adam(learning_rate=1e-3), loss="huber")
    return model

We bind our input and output definitions together into a single structural tf.keras.Model. We use the Adam optimiser to smoothly update learning weights and evaluate training quality using the Huber loss function, which minimises mean-squared error while remaining highly robust to extreme outliers.

Wrap Up

You have now built a custom Graph Attention Network architecture from the ground up! By dynamically computing attention coefficients, your model learns precisely which neighbours matter most. Because we included proper serialisation methods, you can seamlessly save, deploy, and load this model across production environments without running into loading errors.

Resources:

1. https://arxiv.org/pdf/1710.10903