LayerNorm in ViT Models

A typical ViT applies LayerNorm to the CLS embedding before projection:

x = transformer(x)     # [B, N+1, D]
x = x[:, 0, :]         # [B, D]
x = layer_norm(x)      # [B, D]
x = x @ proj

LayerNorm operates independently on each sample across its $D$ feature dimensions.

Suppose

$$x \in \mathbb{R}^{B\times D}.$$

Each batch element $b$ contains one CLS embedding:

$$x_b = [x_{b,1}, x_{b,2}, \dots, x_{b,D}].$$

LayerNorm computes the mean and variance of those $D$ values:

$$\mu_b = \frac{1}{D}\sum_{i=1}^{D}x_{b,i}$$

and

$$\sigma_b^2 = \frac{1}{D}\sum_{i=1}^{D}(x_{b,i}-\mu_b)^2.$$

Each feature is then normalized:

$$\hat{x}_{b,i} = \frac{x_{b,i}-\mu_b} {\sqrt{\sigma_b^2+\epsilon}}.$$

Finally, LayerNorm applies learned parameters $\gamma,\beta\in\mathbb{R}^{D}$:

$$y_{b,i} = \gamma_i\hat{x}_{b,i}+\beta_i.$$

LayerNorm does not normalize across the batch dimension $B$. It normalizes across the embedding dimension $D$, independently for each CLS token.

For example, let

$$B=2,\qquad D=4$$

and

$$x= \begin{bmatrix} 1 & 2 & 3 & 4\\ 10 & 20 & 30 & 40 \end{bmatrix}.$$

The two rows are normalized independently. Their means are

$$\mu_1 = 2.5,\qquad \mu_2 = 25.$$

Ignoring $\epsilon$, $\gamma$, and $\beta$ for simplicity, both rows become approximately

$$[-1.342,-0.447,0.447,1.342].$$

When $\epsilon$ is negligible, the normalized values are invariant to a positive uniform rescaling and a uniform offset of an embedding. The learned $\gamma$ and $\beta$ parameters can subsequently apply a separate scale and bias to each feature.

This illustrates an important property of LayerNorm: it removes the overall offset and scale of each embedding, while preserving the relative pattern among its dimensions.

PyTorch behavior

CLIP and ViT implementations commonly define the layer as

ln_post = nn.LayerNorm(D)

For an input with shape [B, D], PyTorch treats the final dimension as the normalized shape. Conceptually, the operation is

for b in range(B):
    mean = x[b].mean()
    var = x[b].var(unbiased=False)

    x_norm[b] = (x[b] - mean) / sqrt(var + eps)

    output[b] = gamma * x_norm[b] + beta

Both gamma and beta have shape [D].

In the above code snippet, gamma * x_norm[b], * means element-wise multiplication. All three vectors have shape [D]:

$$y_{b,i}=\gamma_i\hat{x}_{b,i}+\beta_i.$$

Each feature has its own learned scale $\gamma_i$ and bias $\beta_i$. Matrix multiplication would use @ instead.

Why is this useful before the CLIP projection?

After the transformer, the CLS token may have arbitrary magnitude and offset:

$$x_{\text{CLS}}\in\mathbb{R}^{D}.$$

LayerNorm normalizes this representation and applies its learned affine transformation before projection:

$$x_{\text{CLS}} \rightarrow \mathrm{LN}(x_{\text{CLS}}) \rightarrow \mathrm{LN}(x_{\text{CLS}})W_{\text{proj}}.$$

If

$$W_{\text{proj}}\in\mathbb{R}^{D\times D_{\text{embed}}},$$

then

$$[B,D]@[D,D_{\text{embed}}] \rightarrow [B,D_{\text{embed}}].$$

LayerNorm(D) does not normalize dimension $D$ globally across the tensor. It independently normalizes each vector in the final dimension:

           normalize this way →
x = [
      [ d1, d2, d3, ... dD ],    ← image 1 CLS token
      [ d1, d2, d3, ... dD ],    ← image 2 CLS token
      [ d1, d2, d3, ... dD ],    ← image 3 CLS token
      ...
    ]

LayerNorm before CLS extraction

Applying the same LayerNorm(D) before extracting the CLS token

x = transformer(x)     # [B, N+1, D]
x = layer_norm(x)      # [B, N+1, D]

independently normalizes every token’s $D$-dimensional embedding. A tensor with shape

$$(B,N+1,D)$$

therefore contains $B(N+1)$ separate vectors to normalize. For a standard LayerNorm(D),

layer_norm(x)[:, 0, :]

and

layer_norm(x[:, 0, :])

produce the same normalized CLS embeddings. This equivalence helps clarify LayerNorm placement within ViT architectures.