Principal Neighbourhood Aggregation (Corso et al., 2020).
PNA combines multiple aggregation functions with degree-based scalers to create a maximally expressive message passing scheme. By using diverse aggregators (mean, max, sum, std) and scalers (identity, amplification, attenuation), PNA captures a richer set of graph structural properties than any single aggregation method.
Architecture
Node Features [batch, num_nodes, input_dim]
Adjacency [batch, num_nodes, num_nodes]
|
v
+------------------------------------------+
| PNA Layer: |
| 1. Apply all aggregators: |
| [mean(N(v)), max(N(v)), |
| sum(N(v)), std(N(v))] |
| 2. Apply scalers to each: |
| [identity, amplification] |
| 3. Concatenate all combinations |
| 4. Project to output_dim |
+------------------------------------------+
|
v
Node Embeddings [batch, num_nodes, hidden_dim]Usage
model = PNA.build(
input_dim: 16,
hidden_dims: [64, 64],
aggregators: [:mean, :max, :sum, :std],
scalers: [:identity, :amplification],
num_classes: 3
)References
- Corso et al., "Principal Neighbourhood Aggregation for Graph Nets" (NeurIPS 2020)
- https://arxiv.org/abs/2004.05718
Summary
Functions
Build a PNA model.
Get the output size of a PNA model.
Single PNA layer: multiple aggregators + degree scalers + projection.
Types
@type build_opt() :: {:activation, atom()} | {:aggregators, [atom()]} | {:dropout, float()} | {:hidden_dims, [pos_integer()]} | {:input_dim, pos_integer()} | {:num_classes, pos_integer() | nil} | {:pool, atom()} | {:scalers, [atom()]}
Options for build/1.
Functions
Build a PNA model.
Options
:input_dim- Input feature dimension per node (required):hidden_dims- List of hidden dimensions (default: [64, 64]):aggregators- List of aggregation functions (default: [:mean, :max, :sum, :std]):scalers- List of degree scalers (default: [:identity, :amplification]):num_classes- If provided, adds a classification head (default: nil):activation- Activation function (default: :relu):dropout- Dropout rate (default: 0.0):pool- Global pooling for graph classification (default: nil)
Returns
An Axon model with two inputs ("nodes" and "adjacency").
@spec output_size(keyword()) :: pos_integer()
Get the output size of a PNA model.
@spec pna_layer(Axon.t(), Axon.t(), pos_integer(), keyword()) :: Axon.t()
Single PNA layer: multiple aggregators + degree scalers + projection.
Options
:name- Layer name prefix (default: "pna"):aggregators- List of aggregation functions (default: [:mean, :max, :sum, :std]):scalers- List of degree scalers (default: [:identity, :amplification]):activation- Activation function (default: :relu):dropout- Dropout rate (default: 0.0)