Skip to content

Normalization Layers

The k3_node.layers.norm module implements 11 specialized normalization layers for graph neural networks, avoiding oversmoothing and gradient instability.

Graph Normalizations

GraphNorm

k3_node.layers.norm.GraphNorm

Bases: Layer

Applies graph normalization over individual graphs as described in the "GraphNorm: A Principled Approach to Accelerating Graph Neural Network Training" <https://arxiv.org/abs/2009.03294>_ paper.

.. math:: \mathbf{x}^{\prime}_i = \frac{\mathbf{x} - \alpha \odot \textrm{E}[\mathbf{x}]} {\sqrt{\textrm{Var}[\mathbf{x} - \alpha \odot \textrm{E}[\mathbf{x}]] + \epsilon}} \odot \gamma + \beta

where :math:\alpha denotes parameters that learn how much information to keep in the mean.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05

GraphSizeNorm

k3_node.layers.norm.GraphSizeNorm

Bases: Layer

Applies Graph Size Normalization over each individual graph in a batch of node features:

.. math:: \mathbf{x}^{\prime}_i = \frac{\mathbf{x}_i}{\sqrt{|\mathcal{V}|}}

PairNorm

k3_node.layers.norm.PairNorm

Bases: Layer

Applies pair normalization over node features as described in the "PairNorm: Tackling Oversmoothing in GNNs" <https://arxiv.org/abs/1909.12223>_ paper.

.. math:: \mathbf{x}i^c &= \mathbf{x}_i - \frac{1}{n} \sum_i \}^n \mathbf{x

\mathbf{x}_i^{\prime} &= s \cdot
\frac{\mathbf{x}_i^c}{\sqrt{\frac{1}{n} \sum_{i=1}^n
{\| \mathbf{x}_i^c \|}^2_2}}

Parameters:

Name Type Description Default
scale float

Scaling factor :math:s of normalization. (default: :obj:1.0)

 1.0
scale_individually bool

If set to :obj:True, will compute the scaling step as :math:\mathbf{x}^{\prime}_i = s \cdot \frac{\mathbf{x}_i^c}{{\| \mathbf{x}_i^c \|}_2}. (default: :obj:False)

 False
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05

DiffGroupNorm

k3_node.layers.norm.DiffGroupNorm

Bases: Layer

The differentiable group normalization layer from the "Towards Deeper Graph Neural Networks with Differentiable Group Normalization" <https://arxiv.org/abs/2006.06972>_ paper, which normalizes node features group-wise via a learnable soft cluster assignment.

.. math::

\mathbf{S} = \text{softmax} (\mathbf{X} \mathbf{W})

where :math:\mathbf{W} \in \mathbb{R}^{F \times G} denotes a trainable weight matrix mapping each node into one of :math:G clusters. Normalization is then performed group-wise via:

.. math::

\mathbf{X}^{\prime} = \mathbf{X} + \lambda \sum_{i = 1}^G
\text{BatchNorm}(\mathbf{S}[:, i] \odot \mathbf{X})

Parameters:

Name Type Description Default
in_channels int

Size of each input sample :math:F.

 required
groups int

The number of groups :math:G.

 required
lamda float

The balancing factor :math:\lambda between input embeddings and normalized embeddings. (default: :obj:0.01)

 0.01
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05
momentum float

The value used for the running mean and running variance computation. (default: :obj:0.1)

 0.1
affine bool

If set to :obj:True, this module has learnable affine parameters :math:\gamma and :math:\beta. (default: :obj:True)

 True
track_running_stats bool

If set to :obj:True, this module tracks the running mean and variance, and when set to :obj:False, this module does not track such statistics and always uses batch statistics in both training and eval modes. (default: :obj:True)

 True

group_distance_ratio(x, y, eps=1e-05) staticmethod

Measures the ratio of inter-group distance over intra-group distance.

MessageNorm

k3_node.layers.norm.MessageNorm

Bases: Layer

Applies message normalization over the aggregated messages as described in the "DeeperGCNs: All You Need to Train Deeper GCNs" <https://arxiv.org/abs/2006.07739>_ paper.

.. math::

\mathbf{x}_i^{\prime} = \mathbf{x}_{i} + s \cdot
{\| \mathbf{x}_i \|}_2 \cdot
\frac{\mathbf{m}_{i}}{{\|\mathbf{m}_i\|}_2}

Parameters:

Name Type Description Default
learn_scale bool

If set to :obj:True, will learn the scaling factor :math:s of message normalization. (default: :obj:False)

 False

MeanSubtractionNorm

k3_node.layers.norm.MeanSubtractionNorm

Bases: Layer

Applies layer normalization by subtracting the mean from the inputs as described in the "Revisiting 'Over-smoothing' in Deep GCNs" <https://arxiv.org/abs/2003.13663>_ paper.

.. math:: \mathbf{x}i = \mathbf{x}_i - \frac{1}{|\mathcal{V}|} \sum_j}} \mathbf{x

Standard & Hetero Normalizations

BatchNorm

k3_node.layers.norm.BatchNorm

Bases: Layer

Applies batch normalization over a batch of features as described in the "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift" <https://arxiv.org/abs/1502.03167>_ paper.

.. math:: \mathbf{x}^{\prime}_i = \frac{\mathbf{x} - \textrm{E}[\mathbf{x}]}{\sqrt{\textrm{Var}[\mathbf{x}] + \epsilon}} \odot \gamma + \beta

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05
momentum float

The value used for the running mean and running variance computation. (default: :obj:0.1)

 0.1
affine bool

If set to :obj:True, this module has learnable affine parameters :math:\gamma and :math:\beta. (default: :obj:True)

 True
track_running_stats bool

If set to :obj:True, this module tracks the running mean and variance, and when set to :obj:False, this module does not track such statistics and always uses batch statistics in both training and eval modes. (default: :obj:True)

 True
allow_single_element bool

If set to :obj:True, batches with only a single element will work as during in evaluation. That is the running mean and variance will be used. Requires :obj:track_running_stats=True. (default: :obj:False)

 False

InstanceNorm

k3_node.layers.norm.InstanceNorm

Bases: Layer

Applies instance normalization over each individual example in a batch of node features as described in the "Instance Normalization: The Missing Ingredient for Fast Stylization" <https://arxiv.org/abs/1607.06450>_ paper.

.. math:: \mathbf{x}^{\prime}_i = \frac{\mathbf{x} - \textrm{E}[\mathbf{x}]}{\sqrt{\textrm{Var}[\mathbf{x}] + \epsilon}} \odot \gamma + \beta

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05
momentum float

The value used for the running mean and running variance computation. (default: :obj:0.1)

 0.1
affine bool

If set to :obj:True, this module has learnable affine parameters :math:\gamma and :math:\beta. (default: :obj:False)

 False
track_running_stats bool

If set to :obj:True, this module tracks the running mean and variance, and when set to :obj:False, this module does not track such statistics and always uses instance statistics in both training and eval modes. (default: :obj:False)

 False

LayerNorm

k3_node.layers.norm.LayerNorm

Bases: Layer

Applies layer normalization over each individual example in a batch of features as described in the "Layer Normalization" <https://arxiv.org/abs/1607.06450>_ paper.

.. math:: \mathbf{x}^{\prime}_i = \frac{\mathbf{x} - \textrm{E}[\mathbf{x}]}{\sqrt{\textrm{Var}[\mathbf{x}] + \epsilon}} \odot \gamma + \beta

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05
affine bool

If set to :obj:True, this module has learnable affine parameters :math:\gamma and :math:\beta. (default: :obj:True)

 True
mode str

The normalization mode to use for layer normalization (:obj:"graph" or :obj:"node"). If :obj:"graph" is used, each graph will be considered as an element to be normalized. If "node" is used, each node will be considered as an element to be normalized. (default: :obj:"graph")

 'graph'

HeteroBatchNorm

k3_node.layers.norm.HeteroBatchNorm

Bases: Layer

Applies batch normalization over a batch of heterogeneous features as described in the "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift" <https://arxiv.org/abs/1502.03167>_ paper. Compared to :class:BatchNorm, :class:HeteroBatchNorm applies normalization individually for each node or edge type.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
num_types int

The number of types.

 required
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05
momentum float

The value used for the running mean and running variance computation. (default: :obj:0.1)

 0.1
affine bool

If set to :obj:True, this module has learnable affine parameters :math:\gamma and :math:\beta. (default: :obj:True)

 True
track_running_stats bool

If set to :obj:True, this module tracks the running mean and variance, and when set to :obj:False, this module does not track such statistics and always uses batch statistics in both training and eval modes. (default: :obj:True)

 True

HeteroLayerNorm

k3_node.layers.norm.HeteroLayerNorm

Bases: Layer

Applies layer normalization over each individual example in a batch of heterogeneous features as described in the "Layer Normalization" <https://arxiv.org/abs/1607.06450>_ paper. Compared to :class:LayerNorm, :class:HeteroLayerNorm applies normalization individually for each node or edge type.

Parameters:

Name Type Description Default
in_channels int

Size of each input sample.

 required
num_types int

The number of types.

 required
eps float

A value added to the denominator for numerical stability. (default: :obj:1e-5)

 1e-05
affine bool

If set to :obj:True, this module has learnable affine parameters :math:\gamma and :math:\beta. (default: :obj:True)

 True
mode str

The normalization mode to use for layer normalization (:obj:"node"). (default: :obj:"node")

 'node'