Skip to content

← Flax step 16 of 100

Medium Primitives

Implement LayerNorm with γ/β

Jump to the task

LayerNorm centers and rescales a feature vector, then applies a learned scale and bias. Here you will implement it directly with Flax Linen parameters.

Normalize across features

For a vector with D features, compute:

mean = sum(x) / D
variance = sum((x - mean)^2) / D
normalized = (x - mean) / sqrt(variance + eps)
output = gamma * normalized + beta

The variance divides by D, not D - 1. gamma and beta each have shape (D,), so every feature has its own learned scale and offset. Initialize gamma to ones and beta to zeros. The initial output is the normalized vector, not generally the original input.

For x = [1, 2, 3, 4], the mean is 2.5 and the variance is 1.25. With eps=1e-5, the initial output is approximately [-1.3416, -0.4472, 0.4472, 1.3416].

What epsilon changes

A positive eps prevents division by zero on a constant vector. All centered values are then zero, so the initial output is zero. A single-feature input behaves the same way.

Before the learned scale and bias, the output has mean zero in exact arithmetic. Its variance is variance / (variance + eps), which is below one when eps is positive. LayerNorm does not normalize to unit Euclidean length. Learned scales and offsets can also change the output mean and variance.

Put epsilon inside the square root as specified. sqrt(variance) + eps also has a nonzero denominator when variance is zero, but it computes a different function and has different derivative behavior.

Last-axis statistics

Use axis=-1, keepdims=True for both reductions. Keeping a length-one final axis lets the statistics broadcast back over the input. For a tensor of shape (batch, sequence, D), this formula normalizes each feature vector separately; it does not pool statistics across examples or sequence positions. The exercise inputs are vectors, but this explains the axis choice.

Unlike BatchNorm's usual training behavior, this calculation uses no batch statistics or running averages. Its learned parameters are the only state.

Problem

Implement MyLayerNorm(nn.Module) with an eps field and a compact __call__(self, x):

  1. Set d = x.shape[-1].
  2. Declare gamma with self.param("gamma", nn.initializers.ones, (d,)) and beta with self.param("beta", nn.initializers.zeros, (d,)).
  3. Compute mu = jnp.mean(x, axis=-1, keepdims=True) and var = jnp.mean((x - mu) ** 2, axis=-1, keepdims=True).
  4. Return gamma * ((x - mu) / jnp.sqrt(var + self.eps)) + beta.

Implement my_layernorm_forward(seed, x, eps) around that module.

The wrapper constructs the module with eps=float(eps), initializes it with jax.random.PRNGKey(int(seed)) and x, then applies it with the full variables mapping and x. These constant parameter initializers make the result independent of the seed. There is no training update.

x is a finite float32 JAX vector of shape (D,), with D >= 1, and values small enough that the intermediate squared values remain finite. eps is a finite positive scalar representable in float32. seed is a nonnegative integer-valued scalar that may arrive as a float. Return a JAX vector of shape (D,), leaving x unchanged.

Loading visualization…