Geometric Considerations for Normalization Layers in Equivariant Neural Networks
Abstract
In recent years, neural networks that incorporate physical symmetry into their architecture have become indispensable tools for overcoming the scarcity of molecular and material data.
However, despite its importance in deep learning, the design and selection of the normalization layer has often been treated as a side issue.
In this study, we shed light on normalization and discuss how the naive application of batch normalization (BatchNorm) to materials science may face unique geometric challenges.
By conducting case studies of the existing literature, we observe that the normalization layer can not only break desired symmetries in the architecture, but may also introduce undesired symmetries that remove task-relevant information from the representation.
To address these issues, we review alternative normalization layers under a consistent framework and lay out the conditions for making a proper choice, geometric-match .
Overall, our survey provides a useful summary and guidelines of normalization layers for practitioners and highlights open issues for further research and development of deep learning architectures for materials design.
Introduction
Although deep neural networks have been demonstrated to be highly capable and useful tools, they are often very difficult to optimize, costing as many as millions of dollars and several weeks to train.
The technique of batch normalization has proven to be a critically important method for stabilizing optimization techniques, allowing for the efficient training of very deep neural networks [1] [2].
BatchNorm was motivated by the need to train deeper spatial convolutional neural networks, which have since become a relatively mature tool.
More recently, practitioners have developed novel architectures, such as graph and group-equivariant networks, which are effective at processing molecular and materials data for a variety of prediction tasks.
However, in the course of attempting to solve these tasks, it has been empirically shown that naive applications of BatchNorm-based techniques are often insufficient for the training of these models [3] [4] [5] [6].
Moreover, when a practitioner develops a novel normalization technique to enable the training of their network, it is typically the case that this technique is designed specifically for the corresponding task and without consideration of similar problems in other fields of study.
This has led to many similarities and redundancies in the literature, where two or more groups have developed similar or even identical techniques, presumably without any knowledge of each other.
To facilitate the development of novel neural networks for geometric and graph-based tasks, it is imperative to create a broader understanding of normalization techniques that will allow for less complicated and expensive optimization of such networks.
Additionally, providing a more unified characterization of existing normalization techniques will reduce the burden of the practitioner in finding and selecting an appropriate technique for their task, as well as prevent redundant techniques from being rediscovered for the same class of problems.
In this context, our work aims to:
Summarize existing normalization techniques, with a particular focus on those used in the training of recent state-of-the-art group equivariant neural networks 2.
Identify similarities and redundancies in the mechanics of such techniques 3.
Explain the existence of these similarities from the perspective of satisfying geometric constraints 4.
Demonstrate that it is a necessary (but not sufficient) condition that normalization techniques must obey the same geometric constraints as the task
When is BatchNorm Insufficient?
To demonstrate the insufficiency of BatchNorm as a normalization technique for graph and groupequivariant networks, we first examine four specific cases in which BatchNorm is insufficient for the training of a particular network.
Subsequently, we create a generalization of the conditions in each example, which captures a necessary (but not sufficient) condition for a normalization technique to succeed on these architectures.
Case 1.
BatchNorm introduces undesired molecular correlations.
The authors of GemNet, a geometric message passing neural network intended for molecular simulations, state that in their network, applying “batch normalization introduces correlations between separate molecules and atoms”, and further demonstrate that applying traditional normalization methods from spatial convolutional network literature reduces network performance [3].
In other words, when applied to
GemNet, BatchNorm would result in the network learning correlations between the behaviors of separate molecules and atoms, even though in reality their behaviors are disjoint.
Case 2.
BatchNorm breaks rotational equivariance.
In “3D-Rotation-Equivariant Quaternion
Neural Networks”, an architecture intended for the processing of geometric data, the authors create a revised version of BatchNorm which is explicitly intended to be rotationally equivariant, respecting the symmetries of their task [4].
Case 3.
BatchNorm is not compatible with polar coordinates.
In “Rotation equivariant vector field networks”, also an architecture intended for processing geometric data, the authors state that, “BN should only normalize magnitudes of the vectors to unit standard deviation.
It would not make sense to normalize the angles, since their values are already bounded and changing their distribution would alter important information about relative and global orientations.”
Instead, they use an altered form of BatchNorm that satisfies these conditions [5].
Case 4.
Mean-subtraction can remove task-relevant information.
The authors of SetNorm, a work that develops an architecture intended for making predictions on set data, note that normalization layers for architectures that operate on sets should (a) maintain permutation equivariance and (b) not deteriorate performance due to removal of information that is useful for prediction [6].
They further note that batch norm is unsuited to their task as it can violate this second condition, and additionally may introduce dependence between inputs.
This is another example of a violation of the constraints supplied by the task.
In each of these cases, there exists a geometric constraint specified by the task, and applying the vanilla version of BatchNorm violates that constraint, leading to a performance decrease.
To formalize this phenomenon, we must first introduce a few definitions:
Definition 1.
(Equivariance and Invariance)
If G is a group that acts on R N , a function ψ : R N → R N is G-equivariant if ψ ( Ax ) =
Aψ ( x ) for any A ∈ G and any x ∈ R N .
If instead, ψ ( Ax ) = ψ ( x ) , then the function is is G-invariant .
Definition 2.
(Task)
A task is a function ψ ⋆ ( X ) =
y ⋆ that takes as input a representation X of the input data, and returns a target object y ⋆ .
The target object can be a classification target, i.e. y ⋆ ∈ { 0 , ..., k } , or a regression target y ⋆ ∈ R d y .
The objective of deep learning is to find a function that well-approximates this task.
Definition 3.
(Geometric Match and Mismatch)
If a neural network is trained to approximate a task, and that task is equivariant or invariant with respect to a group G , then the network is geometrically matched to the task if:
All of its layers are also equivariant or invariant with respect to that group.
None of its layers introduce invariances which are not respected by the task.
If either condition is violated, then we say that the network is geometrically mismatched to the task.
Geometric matching is a fundamentally important principle in deep learning.
For example, the utility of spatial convolutional neural networks in image classification is enabled by the fact that such networks are invariant to the translation of objects within an image, which is geometrically matched to the task of image classification.
Thus, instead of having to learn this symmetry, the network is fundamentally constrained to obey it.
Normalization Beyond Language and Vision Models
When confronted with a situation in which BatchNorm (or similar variants borrowed from language or vision literature, such as InstanceNorm and LayerNorm) do not work, practitioners developing applications typically create an alternative normalization method for that specific task.
Below, in
Table 1, we summarize a representative set of normalization layers developed for applications to graph and geometric data.
The notation is as follows.
A represents the ‘activations’, or the outputs from the activation function following a convolutional layer, which have shape b
×
x × v , where b is the batch size, x is the number of nodes in a graph or pixels in an ‘image’, and v is the number of features per node, or channels per pixel.
µ and σ represent the mean and standard deviation of the activations, respectively, and the subscript describes which subset of the tensor dimensions these statistics are computed over.
g represents a subset of the feature dimension v (a ‘group’).
| X | represents the size of a graph in the dataset, and N represents the average number of nearest neighbors, where the average is computed across the tensor dimensions represented in the subscript.
RMS represents the root mean square,
MLP is a multilayer perceptron, and m is the ‘hidden message’ computed in the internal layers of a graph neural network.
s is a scaling hyperparameter, and α, β, γ , and Γ are learnable parameters.
For
EvNorm and Vector Neuron, the normalization layer is broken into two components - one that acts on the geometric component of the data, and one that acts on the invariant component of the data.
Language and Vision
Graph and Geometric
BN [1]
A− µ { b,x } σ { b,x }
Quaternion [4]
A RMS { x,v }
LayerNorm [7]
A− µ { x,v } σ { x,v }
Vector Perceptron / EGAN [8] [9]
A RMS { x,v }
InstanceNorm [10]
A− µ { x
[Equation]
x }
GraphNorm [11]
A− α ⊙
[Equation]
x } σ
[Equation]
x
[Equation]
GroupNorm [12]
A−
[Equation]
σ
[Equation]
PairNorm [13] s · (cid:18) A− µ { x
[Equation]
v } σ
[Equation]
x
[Equation]
(cid:19)
FilterResponseNorm [14]
A RMS { x }
MeanSubtractionNorm [15]
A − µ { b,v }
VarianceNorm [2]
A σ { b,x }
MessageNorm [16]
MLP (cid:18)
A + s · ∥A { x
[Equation]
∥ 2 m
[Equation]
x } ∥ m { x }∥ 2 (cid:19)
WeightNorm [17] g W ||W||
DiffGroupNorm [18]
A + s (cid:80) g BN ( softmax g ( A Γ) ⊙ A )
SpectralNorm [19] max h : h ̸ =
[Equation]
Wh ∥ 2 ∥ h ∥ 2 · W
SetNorm [6]
A− µ { b,v } σ { b,v }
Scaled Weight [20] g
[Equation]
EvNorm [21] (cid:32) ∥A Positional ∥ 2 − µ {
[Equation]
,v } σ { x,v } ,
A Invariant ∥A Invariant ∥ 2+ 1 β (cid:33)
Vector Neuron [
[Table]
[Equation]
[Table]
[Equation]
[Table]
These techniques can first be divided into two categories - those that follow the approach of the original BatchNorm, and subtract a mean and divide by a standard deviation computed over some subset of the tensor dimensions, and those that perform a rescaling of the activations based on other data-specific quantities.
There are several trends we can observe from the above examples.
Many variations on the original
BatchNorm were in fact created to solve a geometric mismatch with a specific task; for instance,
InstanceNorm was created for the task of image style transfer, and is invariant to the contrast of content images, which is a symmetry respected by the task, while BatchNorm is not.
Similarly,
SetNorm computes statistics over different tensor dimensions than BatchNorm, because the latter may map two samples whose activations differ in only a scale and bias to the same output, introducing an invariance that the task does not respect - a geometric mismatch that can hurt predictive ability.
3.1 Normalization for Geometric Data
In the particular case of geometric data, such as point clouds or molecules represented by their atomic coordinates, there are several works which introduce normalization techniques designed to be rotationally equivariant, so that two data points representing the same input at different rotations will be mapped to the same output modulo rotation.
Among these are the techniques developed in “Rotation equivariant vector field networks”, “3D-Rotation-Equivariant Quaternion Neural Networks”, and “Learning From Protein Structure with Geometric Vector Perceptrons” [5] [4] [9].
Computing either variances or root-mean-squares over the feature dimension and scaling the activations by these quantities is a rotationally equivariant operation which still provides the benefits of normalization.
The methods introduced in ‘Vector Perceptrons’ and ‘Quaternion Neural Networks’ are identical in form, although each method was developed independently and apparently without mutual knowledge of the others’ work, and despite any differences in the task.
Similarly, the method introduced in ‘Vector Field Networks’ is practically identical to VarianceNorm [2], modulo differences in the datatype, as VarianceNorm was developed for image data.
There have additionally been multiple geometric networks that utilize ‘edge-based’ information for normalization (referring to the implicit or explicit graph structure in the data), which again share commonalities.
For instance, the normalization techniques used in ‘SchNet’ [24] and ‘Allegro’ [25] both scale activations by a quantity derived from the connectivity of each molecular or material graph; in the former, the quantity is computed for each atom in a molecule or material, whereas in the latter, the quantity is computed across all atoms in all molecules in the entire dataset.
Both networks involve the individual node (atom) level prediction of quantities which are explicitly invariant to the number of neighbors of each node, and this normalization technique aims to preserve this invariance property.  Empirical Analysis
To demonstrate the importance of normalization techniques on the performance of equivariant neural networks, we present a series of graphs that depict training and testing metrics for three popular models used in chemistry and materials science:
Allegro [25], GemNet [3], and SchNet [24].
For each network, we train a variant with and without the normalization techniques described by the authors in their respective papers.
To test the effects of normalization on Allegro, we utilize the official PyTorch / e3nn implementation available on Github [28], and train on the ‘Uracil’ data included within the MD17 dataset [27], using the hyperparameters specified in the tutorial written by the authors. For SchNet, we utilize an open-source implementation [29] based on the architecture described in [24], with a batch size of 8, 3 hidden layers, which each have 3 interaction layers of sizes 128, 128, and 4 respectively. We train SchNet on the ‘CHAMPS Scalar Coupling’ dataset [30], using Adam with an initial learning rate of 0.001 and exponential weight decay with coefficient 0.99.
For GemNet, we utilize the official
PyTorch implementation available on Github [31], and train it on the ‘COLL’ dataset [32] included in the repository, using the default hyperparameters defined by the authors.
In each case, we train one network with the normalization methods implemented in the respective repository, and one without, using identical hyperparameter settings.
Figure 1, above, depicts the log of the train and test losses for the Allegro network trained with and without normalization.
For the network that has no normalization, we observe ‘gradient explosion’ [33], and losses that are initially 11 orders of magnitude greater than its normalized counterpart.
Though the unnormalized loss decreases, it is unable to achieve a value comparable to the normalized architecture.
Figure 2, below, depicts the train and test losses for the SchNet and GemNet networks, again trained with and without normalization.
Although the effects are not as dramatic, we still observe that applying normalization techniques yields an order of magnitude decrease in the loss.
Conclusion
Normalization techniques have been proven to be a powerful tool for enabling the training of larger networks, and yet a large proportion of existing deep learning applications to chemistry problems completely forego using them [34] [35] [36] [37] [38].
As the datasets involved in molecular and materials prediction tasks grow, so will the networks, and thus the difficulty involved in successfully training an effective neural network to solve these tasks will increase correspondingly.
Although several techniques have been developed for existing geometric datasets, no general principles of normalization exists for these new classes of architectures.
Here, we initiated the discussion by providing the above survey, noting the commonalities between these newly developed normalization techniques that make them effective for their specific tasks, and using the guideline of geometric matching as a necessary condition for maintaining stability during optimization.
As both the interest and the size of the datasets in molecular and materials deep learning grows, we end by posing two important open challenges that will be crucial to its practical implementations: what are the most general conditions under which traditional normalization techniques fail, and what is the most general formulation of a normalization layer which may be applied to any graph-based or geometric dataset underlying a molecular or materials prediction task?