Add layer norm

README.md

@@ -234,7 +234,8 @@ python examples/variational_autoencoder.py -m
 ```
 
 ## Invariant MNIST (fully connected)
-Some comparisons using parameter scans maxabs normalization as default
+Some comparisons using parameter scans maxabs normalization as default. piecewise polynomial cases use 2 segments. I only
+did one run each.
 ```
 python3 examples/invariant_mnist.py -m mlp.n=2,3,4,5,6 mlp.hidden.width=128 mlp.layer_type=polynomial optimizer=sophia
 ```

high_order_layers_torch/layers.py

@@ -2,6 +2,7 @@
 
 import torch.nn as nn
 from torch.nn import Linear
+import torch.nn.functional as F
 
 from .FunctionalConvolution import *
 from .FunctionalConvolutionTranspose import *
@@ -17,6 +18,28 @@
 )
 
 
+class LazyLayerNormLastDim(nn.Module):
+    """
+    Lazily initialize the layer norm to the last dimension of the input
+    variable. Assumes dimension remains constant.
+    """
+
+    def __init__(self, bias=True):
+        super().__init__()
+        self.weight = None
+        self.bias = bias
+
+    def forward(self, input):
+        if self.weight is None:
+            ndim = input.shape[-1]
+            self.weight = nn.Parameter(torch.ones(ndim)).to(input.device)
+            self.bias = (
+                nn.Parameter(torch.zeros(ndim)).to(input.device) if self.bias else None
+            )
+
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+
 class MaxAbsNormalizationLast(nn.Module):
     """
     Normalize the last dimension of the input variable
@@ -101,7 +124,9 @@ def forward(self, x):
 
 normalization_layers = {
     "max_abs": MaxAbsNormalization,
+    "max_center": MaxCenterNormalization,
     "l2": L2Normalization,
+    "layer_norm": LazyLayerNormLastDim,
 }