LayerNorm and RMSNorm

ImplementEasy15 min14 checks

layer-normrms-normnormalizationtransformers

Lesson: The Full Transformer — Encoder + DecoderCheat sheetTrack: Deep learning coreTrack: Transformers from scratch

Problem

A transformer normalizes the vector of each token before attention and before the feed forward block. LayerNorm centers each vector and divides it by its standard deviation. RMSNorm skips the centering and divides by the root mean square. Modern LLMs use RMSNorm because it is cheaper and works as well.

Task

Write two functions.

  1. layer_norm(x, gamma, beta, eps=1e-5) returns LayerNorm of x over the last axis.
  2. rms_norm(x, weight, eps=1e-6) returns RMSNorm of x over the last axis.

For each vector x of length D along the last axis:

\mu = \frac{1}{D} \sum_j x_j, \quad \sigma^2 = \frac{1}{D} \sum_j (x_j - \mu)^2, \quad \text{LayerNorm}(x)_j = \gamma_j \frac{x_j - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta_j
\text{RMSNorm}(x)_j = w_j \frac{x_j}{\sqrt{\frac{1}{D} \sum_k x_k^2 + \epsilon}}

The variance is the biased variance. It divides by D, not by D - 1. The eps goes inside the square root in both functions.

Input and output

NameTypeShapeMeaning
xfloat array(..., D)the vectors to normalize, with zero or more leading axes
gammafloat array(D,)the LayerNorm scale of each feature
betafloat array(D,)the LayerNorm shift of each feature
weightfloat array(D,)the RMSNorm scale of each feature
epsfloatthe small number added inside the square root
  • Both functions return a float array with the shape of x.

Constraints

  • x has 1 to 4 axes. 1 <= D <= 1024.
  • A vector can be constant or all zeros. The output must not contain NaN.
  • Do not change the input arrays.
  • Use numpy only.

Examples

pythonx = np.array([[1.0, 2.0, 3.0, 4.0], [2.0, -2.0, 0.0, 4.0]])
gamma = np.array([1.0, 1.0, 2.0, 0.5])
beta = np.array([0.0, 0.5, 0.0, -1.0])
layer_norm(x, gamma, beta)
text[[-1.34163542  0.05278819  0.89442361 -0.32918229]
 [ 0.44721315 -0.84163944 -0.8944263  -0.32918028]]
pythonrms_norm(x, gamma)
text[[ 0.36514835  0.73029669  2.19089008  0.73029669]
 [ 0.81649651 -0.81649651  0.          0.81649651]]

Hints

4 hints

    Solution and notes

    Locked

    Pass every check with Submit to open the reference solution and the notes.

    Run it on your computer

    Clone the repository, install numpy and run the same checks in a terminal:

    python3 scripts/practice/run.py layer-norm-rms-norm my_answer.py
python3 scripts/practice/run.py layer-norm-rms-norm my_answer.py --examples

    Start my_answer.py from starter.py.

    Your code

    Python loads on the first run

    Your code is saved in this browser as you type. On a phone you can read and edit the code, but the checks are easier to use on a desktop.

    Keyboard shortcuts
    Ctrl Enter
    Run the examples
    Ctrl Shift Enter
    Submit
    Tab / Shift Tab
    Indent or outdent the lines
    Esc then Tab
    Move the focus out of the editor
    Alt H
    Show the next hint

    On a Mac, use Cmd in place of Ctrl.

    Select Run to check the examples, or Submit to run every check.