Where do you insert layer norm in a residual block? After the addition or before the activation function (RELU in this case)?
2 Answers
Two things:
- Layer Norm wasn't invented before ResNet. ResNet still uses the regular Batch Norm.
- The model to use Layer Norm is residual block is ConvNeXt. Based on this line, it applies LayerNorm after the first Conv layer, before the activation layer.
-
$\begingroup$ I am aware of (1), but I should have added that I cannot use batch-norm because of a gradient penalty, which penalizes the gradient with respect to each input independently, rendering batch-norm ineffective. Regarding (2), I am not using any convolutional layers, as I don't work with image data. Do you know of any ResNet implemenatations with MLPs that use layer norm? $\endgroup$ Commented Mar 22, 2023 at 15:42
Usually you insert the normalization layer (be it BatchNorm, LayerNorm or whatever) after the convolutional layer and before the activation layer, i.e. Conv + Norm + ReLU.
The original ResNet applies the addition (skip connection) just before the last ReLU, but this design was revised in a follow-up paper by the same authors. You can also see here for figures and explanations.
Regarding using skip connections with fully-connected layers, I think you can take a look at the transformer architecture. The transformer encoder (and decoder) architecture uses skip connections around fully connected layers. It also uses LayerNorm instead of BatchNorm. However, with this architecture there are again debates where exactly to place the LayerNorm layer. This paper suggests using Pre-LayerNorm instead of Post-LayerNorm, because it stabilizes the training of the Transformer model. You can also see here for figures and explanations.