Advanced Spatial GNNs (GraphSAGE Variants, PNA)

While foundational spatial GNNs like the original GraphSAGE provide a powerful framework for learning from graph structures by aggregating neighbor information, their effectiveness can sometimes be limited by the choice of aggregation function. Standard aggregators like mean, max, or sum pooling each have inherent biases. Mean pooling can obscure information from particularly influential neighbors, max pooling focuses only on the strongest signals, and sum pooling can be heavily influenced by node degree, potentially leading to instability or scaling issues. Advanced spatial GNNs address these limitations by introducing more sophisticated aggregation mechanisms or combining multiple strategies.

GraphSAGE Variants and Enhancements

The GraphSAGE framework itself is flexible, allowing for various aggregation functions. Research and practice have investigated variants of the initial mean, max-pooling, and LSTM aggregators:

1. Attention-based Aggregation: While Graph Attention Networks (GAT), covered earlier, represent a distinct architecture, incorporating attention mechanisms within a GraphSAGE-like aggregation step is a common enhancement. This allows the model to learn the importance of different neighbors dynamically, weighting their contributions during aggregation instead of treating them uniformly (like mean) or based solely on feature magnitude (like max).

2. Combining Aggregators: Instead of choosing a single aggregator, some approaches combine the outputs of multiple aggregators (e.g., concatenating the results of mean and max pooling). This allows the model to potentially capture different aspects of the neighborhood distribution simultaneously.

3. Normalization Techniques: Applying normalization techniques like Layer Normalization or Batch Normalization after aggregation but before the non-linear activation and update step can improve training stability and performance, especially in deeper models or graphs with diverse neighborhood sizes.

These variants often aim to improve the representational capacity of the aggregation step, making the resulting node embeddings more informative. However, a more structured approach to combining multiple aggregation perspectives led to the development of Principal Neighbourhood Aggregation.

Principal Neighbourhood Aggregation (PNA)

Principal Neighbourhood Aggregation (PNA) directly tackles the shortcomings of single, simple aggregators. The core idea is that no single aggregator (mean, max, sum, standard deviation) is sufficient to capture the full information contained within a node's neighborhood distribution. PNA proposes combining multiple aggregators simultaneously and, importantly, incorporating degree information explicitly as a scaling factor.

Motivation

Consider the limitations:

• Mean/Max Aggregation: These are insensitive to the number of neighbors (degree). A mean or max over 2 neighbors might yield the same result as over 200, losing potentially significant structural information.
• Sum Aggregation: While sensitive to degree, it can explode for high-degree nodes and vanish for low-degree nodes, making training difficult. It also doesn't capture the distribution of features well (e.g., variance).

PNA aims to create a more comprehensive aggregation function by:

1. Using a combination of multiple aggregators to capture different statistical properties of the neighborhood.
2. Using degree-based scalers to modulate the aggregated information, making the network aware of the neighborhood size.

PNA Mechanism

In a PNA layer, the aggregation step for updating node $i$'s representation $h_i^{(k)}$ to $h_i^{(k+1)}$ involves computing:

$$h_i^{(k+1)} = U \left( h_i^{(k)}, \bigoplus_{\text{agg} \in A, s \in S} s(d_i) \cdot \text{agg}(\{ m_{ji}^{(k)} \mid j \in \mathcal{N}(i) \}) \right)$$

Where:

• $h_i^{(k)}$ is the feature vector of node $i$ at layer $k$.
• $m_{ji}^{(k)}$ is the message from neighbor $j$ to node $i$ at layer $k$, often derived from $h_j^{(k)}$ (e.g., $m_{ji}^{(k)} = \text{MLP}(h_j^{(k)})$).
• $A$ is a set of aggregation functions. Common choices include:
  • Mean ($\mu$): Average of neighbor messages.
  • Standard Deviation ($\sigma$): Spread of neighbor messages.
  • Maximum ($\max$): Element-wise maximum message.
  • Minimum ($\min$): Element-wise minimum message.
• $d_i = |\mathcal{N}(i)|$ is the in-degree of node $i$.
• $S$ is a set of scaling functions applied based on the degree $d_i$. Typical scalers include:
  • Identity: $s(d_i) = 1$.
  • Logarithmic Scaler (Amplification): $s(d_i) = \log(d_i + 1)$. Amplifies signals for higher-degree nodes.
  • Inverse Logarithmic Scaler (Dampening): $s(d_i) = 1 / \log(d_i + 1)$. Dampens signals for higher-degree nodes, relative to lower-degree ones.
• $\bigoplus$ represents a combination operator, usually concatenation, that joins the results from all aggregator-scaler pairs.
• $U$ is an update function, typically an MLP, that combines the node's previous representation with the aggregated and scaled neighborhood information.

Flow of Principal Neighbourhood Aggregation (PNA). Input messages $m_{ji}$ from neighbors are processed by multiple aggregators (mean, std, max, min, etc.). Each aggregated result is then scaled using different functions of the node's degree $d_i$ (identity, log, inv-log, etc.). These scaled results are combined (e.g., concatenated) and then passed to an update function $U$ along with the node's previous representation $h_i^{(k)}$ to produce the new representation $h_i^{(k+1)}$.

Advantages

PNA often demonstrates strong performance on various graph machine learning benchmarks, particularly those requiring sensitivity to fine-grained structural differences or involving graphs with widely varying node degrees. By explicitly combining multiple statistical moments of the neighborhood feature distribution and scaling them by degree, PNA layers can generate richer and more discriminative node representations compared to simpler spatial methods.

Implementing PNA requires calculating node degrees and applying the multiple aggregations and scalings. Libraries like PyTorch Geometric provide optimized implementations (e.g., PNAConv). The increased computational cost per layer due to multiple aggregations is a trade-off for its enhanced expressive capability.

Comparing Advanced Spatial Approaches

• GraphSAGE Variants: Often simpler to implement and computationally less intensive than PNA. Attention-based aggregation within GraphSAGE can provide dynamic weighting but doesn't explicitly incorporate degree scaling or multiple fixed aggregators in the same structured way as PNA.
• PNA: Provides a principled way to combine multiple fixed aggregators and degree information. It can be more powerful on complex graph structures but comes with higher computational overhead per layer. Its effectiveness relies on the chosen set of aggregators and scalers being suitable for the task and data.

Choosing between these approaches depends on the specific problem, the characteristics of the graph data (especially degree distribution), and the available computational resources. For tasks where distinguishing complex local structures is important, PNA is a strong candidate. Simpler GraphSAGE variants might suffice or even be preferable for very large graphs where computational cost is a primary constraint, or when combined with effective sampling strategies (discussed in Chapter 3).