Lazy loaded image1.4 FFN, Residual Addition, and LN

1.4 FFN, Residual Addition, and LN

1.4.1 Three Core Modules

notion image

Feed Forward Network (FFN)

After passing through MHA, tokens have exchanged some information with each other, but they haven't yet thought deeply about what they learned from other tokens. So when a token gathers information through MHA, it then independently learns and processes this information through a feedforward network.

Residual Addition

In the transformer structure, multiple blocks are usually stacked, which leads to optimization issues as the depth of the neural network increases. Residual connections mitigate the rapid attenuation or exploding of gradients during deep network transmission.
In terms of feature prediction, we can consider a simple case like
Where the function indicates a network layer, which predicts the next layer of features based on the last layer’s features. Usually it is challenging to directly predict the output features . Now, if we add a residual addition to it like below:
This time, the function does not have to predict the features directly based on . Instead, the function only predicts the difference between the last layer’s features and the target features , i.e. . This makes the prediction much simpler! By the way, this “predicting the difference” concept can also be seen in Denoising Diffusion Probabilistic Models (DDPM), which predicts the noise instead of the original at each time step for better model performance.
In practice, this mechanism allows gradients to flow more smoothly, even in very deep networks, thereby avoiding the training difficulties caused by gradient vanishing.

Layer Norm

Normalization helps speed up convergence and alleviates the problem of vanishing or exploding gradients.
In general, Batch Norm is used in CV, Layer Norm is used in NLP. A key difference is what information you need to retain. Here’s an example:
  • In NLP, for the sentences I love LLM, I show speed, and I practice CV, where each sentence undergoes normalization. Suppose each word is a token, and we perform normalization like Batch Norm on [I, I, I], [love, show, practice]... Obviously, applying the Batch Norm in NLP does not make sense.
    • Layer Norm normalizes within a sentence itself, whereas Batch Norm normalizes across sentences. Therefore, Layer Norm allows for preserving context dependencies.
  • In CV, Batch Norm is applied across different channels of an image (such as RGB channels, or latent channels in deeper networks). After normalization, distribution distinction across channels is maintained.
notion image
In the image:
  • H (Height): usually refers to the vertical dimension of the image or feature map.
  • W (Width): usually refers to the horizontal dimension of the image or feature map.
  • C (Channel): usually refers to the color channels of the image (e.g., R, G, B in RGB images) or feature channels (e.g., different feature channels in convolutional neural networks).
  • N (Number of samples): usually refers to the batch size.
Different normalization operations process these dimensions in different ways:
  • Batch Norm: Normalizes along the N and H,W dimensions, typically used in CNN.
  • Layer Norm: Normalizes along the C and H,W dimensions, typically used for input normalization in fully connected layers (MLP).
  • Instance Norm: Normalizes along the H and W dimensions, commonly used in domain adaptation/style transfer tasks.
  • Group Norm: Normalizes along grouped C and H,W dimensions, often used as a substitute for Batch Norm to avoid performance fluctuations due to different batch sizes.
<ins/>

1.4.2 Position and Implementation of Layer Norm

Position of Layer Norm

Post-LayerNorm, Pre-LayerNorm, and Sandwich-LayerNorm
Post-LayerNorm, Pre-LayerNorm, and Sandwich-LayerNorm
  • Sensitivity Analysis of Layer Norm
    • notion image
    • Post-norm: Deep layers are prone to unstable training, and the gradient range increases layer by layer.
    • Pre-norm: The gradient range is relatively stable in each layer, training is more stable, but depth is limited.
    • Post-norm w/ warm-up: Effectively control the activation values in each layer to prevent them from being too large. This helps the model learn data features better
    • Sandwich-norm: Can control the activation values as well from being too large, but still suffers from training instability and collapse.
  • Differences between Post-Norm and Pre-Norm:
    • Generally speaking, Post-Norm after the residuals makes features regularization more efficient and improves the model convergence; while Pre-Norm doesn’t regularize some features and doesn’t block the residuals from the model bottom to top, which avoids gradient vanishing or exploding. This is why Pre-Norm is more stable and more commonly used in training LLM. The current popular conclusion is: under the same setting, Pre-Norm architectures are easier to train, but their final performance is not as good as Post-Norm architecture. A L-layer Pre-Norm model’s effective layers are less than a L-layer Post-Norm model.
    • Pre-Norm structure increases model width at the expense of depth, as depth is typically more important than width. So reducing depth leads to worse performance eventually. In contrast, Post-Norm does the opposite. Thus, once Post-Norm is trained well, it performs better. However, it's hard to train at first and requires warm-up.
    • Post-Norm's instability mainly comes from gradient vanishing, as well as being stuck in local optimal during early updates due to a large learning rate.

Implementation of Layer Norm

Normalizes all features of each individual sample, which removes the magnitude differences between different samples while preserving the relative magnitudes among different features within a single sample.
LayerNorm is commonly used in the NLP domain. In this case, the input size is (batch size × sequence length × embedding dimension). The formulas are as follows:
where and are learnable scale and shift parameters, and the is a small numebr for preventing “dividing by 0” in the third equation; the indicates the sequence () and embedding () dimensions.
Note that this operation is technically standardization, but conventionally people call it normalization in deep learning.
  • LayerNorm Code
    • This is an example of applying LN to image data (i.e. 4 dimensions)
<ins/>

1.4.3 Other Normalization

RMS Norm

RMS Norm, short for root mean square layer normalization. Compared to Layer Norm, RMS Norm removes the part of subtracting the mean during calculation, with the idea of treating the mean as 0, leading to faster computational speed and almost the same performance.
One popular interpretation of the Layer Norm is that it recenters and rescales input vectors to achieve invariance to input and weight shifts. The recentering leads to the model's output remaining relatively stable even if the inputs or weights are shifted by a constant amount. The rescaling results in that the model's behavior is not significantly affected by scaling the inputs or weight. The authors of RMS Norm believe that the success of LN lies in the invariance to rescaling, not in the centering of the mean.
RMS measures the square root of the mean of the square of inputs, then scales the input onto a unit sphere with radius of , where the H is the number of features in a sample (across multiple dimensions). By doing this, the output distribution remains unaffected by input and weight rescaling, which helps in stabilizing deep neural network training.
Although the Euclidean norm differs from the RMS by only a factor of and has been successfully applied in some domains, it does not perform well in layer normalization.

Deep Norm

Deep Norm can alleviate the issue of unstable training caused by gradient explosions in deep Transformers. It limits model updates within a reasonable range, making the training process more stable.
The Deep Norm method up-scales the residual connection (alpha > 1) before applying Layer Norm, and additionally, down-scales the model parameters (beta < 1) during initialization.
💡
Causes of Unstable Training:
  1. At the early stages of training, model parameters update too rapidly, causing the model to fall into a bad local optima. This, in turn, increases the input magnitudes to each Layer Norm.
  1. The gradient at the bottom layers of Pre-LN is often greater than that at the top layers.
  1. When using DeepNorm, replace Post-LN with it.
  1. DeepNorm is essentially Layer Norm, but before normalization, it performs an up-scale to residual connection. , where represents modules like Self-Attention or other token-mixers, and is a constant.
  1. torch.nn.init.xavier_normal_(tensor, gain=1) is Xavier normal initialization, where the mean of the parameter is 0, and the standard deviation is as below. Here, fan_in and fan_out refer to the number of input and output elements of the tensor, respectively. This type of initialization is used to ensure the variance of the output remains unchanged from the input.
  1. DeepNorm also down-scales parameters during initialization. Notably, the initialization differs for ffn, v_proj, out_proj, and q_proj, k_proj.
  1. The best values of and are different across architectures.
📌
How does Xavier Initialization, a.k.a. Glorot Initialization, preserve variance?
Assume a dense layer of MLP
  • : input with dim of
  • : weight matrix with dim of
  • : output with dim of
We also assume that we have:
  • i.i.d., with zero mean and variance of
  • is independent, with zero mean and variance of
Then we have
Note that if and are independent.
Next, we have
This means that as long as we set the variance of weight to be , the input and output will keep the same variance. To consider the input and output dimension of the weight, Xavier choose to set
When the , we technically have . Additionally, a hyperparameter of is introduced for researchers to finetune the variance.
<ins/>

1.4.4 Computation and Activation of FFN

A common formula for computing the FFN block with activation function:
The FFN comprises two linear layers with activation in between (typically ReLU) to add nonlinearity and increase model capacity.
Gated Linear Unit (GLU)
The GLU includes two parallel linear projections with one going through a sigmoid gate, essentially acting as a gating mechanism, selecting some elements and ignoring others.
By incorporating the GLU into FFN, we have
SwiGLU and GeGLU are based on Swish and GeLU activation functions, respectively, replacing the sigmoid in GLU. For example: In LLaMA2-7B, the original input dimension for FFN is 4096, and typically the hidden dimension is 4x the input = 16384. In order to keep the parameter count similar (In FFN with SwiGLU there are 3 weight matrices instead of 2), the hidden dimension is reduced — typically to 2/3 of the original. To make it a multiple of 256 and make tensor cores’ life easier, it’s rounded to the nearest multiple, resulting in a final hidden size of 11008.
Swish Activation Function: GeLU Activation Function: , which is an approximate version that’s commenly used
Large language models now typically use SwiGLU to replace the traditional FFN structure

1.4.5 Common Activation Functions

Sigmoid

Drawbacks:
  • When the input is large or very small, the gradient approaches 0, which can easily lead to gradient vanishing.
  • The function output is not centered around 0, leading to always positive activation values. This makes later activation functions work on a limited value space.
    • For example, except for the initial input images, all latents in all layers of the model are positive, and all weights matrices are working on positive values, generating positive outputs. This can push some layers to the saturated region of sigmoid, resulting in small gradients in that layer. Because of the nature of the chain rule, any single small gradients can make the whole gradients propagation very small and not efficient for model updating.
  • The Sigmoid function involves exponential operations, which run relatively slowly on computers and consume more computing resources.
notion image

Tanh

Advantage: tanh is “zero-centered”. Therefore, in practical applications, tanh is somewhat better than sigmoid.
Disadvantages:
  • Still suffers from the problem of gradient saturation
  • Still involves exponential operations
notion image

ReLU

Advantages:
  • ReLU solves the vanishing gradient problem: when the input is positive, the neuron won’t saturate.
  • Because ReLU is linear and non-saturating, it allows for fast convergence in SGD.
  • Low computational complexity: no exponential operations are required.
Disadvantages:
  • Like Sigmoid, its output is not zero-centered.
  • Dead ReLU problem: When the input is negative, the gradient is 0. The neuron and all neurons before it will always get zero gradient from it, preventing their corresponding parameters from being updated.
notion image

Leaky ReLU

Advantages:
  • Solves the problem of dead neurons when ReLU receives negative input.
  • Leaky ReLU is linear and non-saturating, which allows for fast convergence in SGD.
  • No exponential operations required.
Disadvantages:
  • The parameter α in the function must be manually set based on prior knowledge or experimentation (typically set to 0.01).
  • Kind of like linear activation, which can result in poor performance on complex classification tasks.
notion image

Exponential Linear Unit (ELU)

Advantages:
  • ELU tries to bring the output mean of the activation function closer to zero, making the normal gradient closer to the unit natural gradient, which helps speed up learning.
  • ELU does not saturate to a fixed negative value for small inputs, reducing variation and information loss during forward propagation.
Disadvantage:
  • Requires computing exponentials during calculation, resulting in lower computational efficiency.
notion image

Swish(SiLU)

Advantages:
  1. Smooth and non-monotonic
      • Unlike ReLU, which has a sharp corner at zero, Swish is a smooth and continuous function, which benefits gradient propagation.
      • Its non-monotonic nature allows it to retain some negative values in certain regions, giving it greater expressive power.
  1. Proven to perform better in deep models
      • On large-scale tasks like ImageNet, Swish often outperforms ReLU, ELU, and Leaky ReLU in terms of accuracy and convergence.
notion image

SwiGLU

Where:
  • : element-wise multiplication (Hadamard product).
  • W and V are weight matrices.
  • x is the input vector.

Softmax

The Softmax function is often used as the activation function in the output layer of neural networks. It maps the output values of the output layer to the range [0, 1] through the activation function, forming a probability distribution.
In multi-class classification problems, the larger the output of the Softmax activation function, the higher the probability of the corresponding class being the true class.
Prev
1.3 Attention
Next
1.5 Positional Encoding
Loading...
Article List
LLM Learning Roadmap
✨ Awesome-Anything
🖼️ Digital Image Processing
🍃 LLM Components
🌱 LLM Pre-training
☘️ LLM Post-Training
🍀 LLM Popular Models
🪴 LLM Applications
🌿 LLM Optimization
🌾 LLM Compression
🌵 LLM Hands-on Practice
🌴 LLM Must-read Papers
🌳 LLM Q&A
🐝 VLM Image Encoders
📝 MISC.