nGPT — Normalized Transformer on the Hypersphere

§ 1 · Premise

Normalization as a structural property

Standard transformers normalize at three specific places per block: a Pre-Norm before attention, a Pre-Norm before the FFN, and a final norm before unembedding. Between those normalizations, activation norms drift — sometimes growing across depth, sometimes shrinking. The model is asked to learn weights that produce well-conditioned next-block inputs from poorly-conditioned previous-block outputs, and the per-row norms of the projection matrices $W_q, W_k, W_v, W_o, W_u, W_v$ are themselves unconstrained.

Loshchilov et al.’s 2024 proposal (NVIDIA; ICLR 2025) treats the unbounded vector norms as the source of the problem and the introduced LayerNorm/RMSNorm modules as the symptomatic patch. Their inversion: constrain every vector in the network — token embeddings, attention K/Q/V projections, attention output, FFN intermediate, FFN output, even the rows of every weight matrix — to lie on the unit hypersphere $S^{d-1} \subset \mathbb{R}^d$. The model operates entirely on unit vectors. Norms cannot drift because they are always $1$ by construction.

Two structural consequences follow immediately. First, the residual update can no longer be a plain addition (the sum of two unit vectors leaves the sphere), so the block update is rewritten as a step along the sphere with a learnable step size. Second, every inner product $W \mathbf{x}$ is a cosine similarity rather than an arbitrary dot product — bounded in $[-1, 1]$ — which removes the need for a $1/\sqrt{d_k}$ softmax temperature and replaces it with a learnable scale.

The Section 1 framing in the paper places nGPT in a lineage with prior work on representation learning on the hypersphere (Wang and Isola, ICML 2020) and with the observation by Xiong et al. that aggressive normalization placement — Pre-Norm — was already moving in this direction. nGPT extrapolates the trajectory to its limit: not “where do we normalize” but “is normalization the only invariant the architecture should carry”.

§ 2 · Derivation

The hypersphere update, step by step

Start from a baseline Pre-Norm transformer block. Let $\mathbf{h} \in \mathbb{R}^{d_{\text{model}}}$ denote the residual-stream state at the input to a block. The Pre-Norm update (Eq. 4 and 5 of the paper) is:

The attention and MLP sublayers produce vectors of unconstrained magnitude that are added into the residual.

The sphere update. Loshchilov et al. (§ 2.2.2) observe that on the sphere the natural analogue of linear interpolation between two unit vectors is spherical linear interpolation (SLERP, Shoemake 1985; Eq. 6 in the paper):

where $\theta = \arccos(\mathbf{a} \cdot \mathbf{b})$. For small $\alpha$ (small steps), SLERP is well-approximated by ordinary linear interpolation followed by retraction onto the sphere. Writing $\mathbf{a} \leftarrow \mathbf{a} + \alpha (\mathbf{b} - \mathbf{a})$ and generalizing $\alpha$ from a scalar step size to a per-coordinate learnable vector $\boldsymbol{\alpha} \in \mathbb{R}^{d_{\text{model}}}_{\geq 0}$ (Eq. 8, Eq. 9 in the paper), the nGPT block update becomes (Eq. 10, Eq. 11):

where $\mathbf{h}_A = \mathrm{Norm}(\mathrm{Attn}(\mathbf{h}))$, $\mathbf{h}_M = \mathrm{Norm}(\mathrm{MLP}(\mathbf{h}))$, and $\mathrm{Norm}(\cdot)$ is plain L2 normalization with no learned gain. Each block moves the residual state a learned fraction $\boldsymbol{\alpha}$ toward the (unit-normalized) sublayer suggestion and then retracts back to the sphere. The authors call $\boldsymbol{\alpha}_A, \boldsymbol{\alpha}_M$ eigen learning rates — the diagonal of a variable-metric optimizer acting in the embedding space (paper § 2.2.2 and Appendix A.2). $\boldsymbol{\alpha}_A$ is initialized to $0.05$ on the order of $1/n_{\text{layers}}$ with scale $1/\sqrt{d_{\text{model}}}$.

Weight-row normalization. Each weight matrix $W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ has its rows projected to unit norm along the input dimension after every gradient step (paper § 2.6, step 2). Concretely, this applies to the input embedding $E_{\text{input}}$, output embedding $E_{\text{output}}$, attention projections $W_q, W_k, W_v, W_o$, and MLP projections $W_u, W_v, W_{oMLP}$. After normalization, each row of $W$ is a unit vector in the input space, so the matrix-vector product $W \mathbf{x}$ for unit-norm $\mathbf{x}$ returns a vector whose $i$th component is the cosine similarity between row $i$ of $W$ and $\mathbf{x}$ — bounded in $[-1, 1]$.

Attention with cosine-similarity logits. Because Q and K rows of unit norm produce bounded dot products, the softmax temperature changes. The baseline scales by $1/\sqrt{d_k}$ to keep the pre-softmax variance constant for unit-variance Gaussian queries and keys; for unit-norm Q and K the dot product variance is instead $1/d_k$, so the correct scale to restore variance-one is $\sqrt{d_k}$. The paper introduces a per-head learnable vector $\mathbf{s}_{qk} \in \mathbb{R}^{d_k}$ (Eq. 15, 16) that re-normalizes Q and K after the projection, with initialization equivalent to softmax scale $d_k^{1/4}$:

RoPE is still applied to Q and K before this normalization (§ 2.3.1), preserving the relative-position signal.

MLP rescaling. The SwiGLU intermediates $\mathbf{u}$ and $\boldsymbol{\nu}$ get their own per-coordinate scaling vectors $\mathbf{s}_u, \mathbf{s}_\nu \in \mathbb{R}^{d_{\text{MLP}}}$ (Eq. 20, 21). The rescaling of $\boldsymbol{\nu}$ by $\sqrt{d_{\text{model}}}$ is necessary “to benefit from the non-linearity of SiLU” — without it the gate input would sit in the near-linear regime of SiLU around the origin, collapsing the activation function.

Logit temperature. The output logits $\mathbf{z} = E_{\text{output}} \mathbf{h}$ are also bounded in $[-1, 1]$ per coordinate after unit normalization. A learnable per-vocabulary scale $\mathbf{s}_z \in \mathbb{R}^V$ (Eq. 3) restores the temperature of the cross-entropy softmax.

The full conversion recipe. Section 2.6 enumerates the seven changes from baseline GPT to nGPT: remove all RMSNorm modules; normalize all weight matrices and embeddings along their embedding dimension after each step; replace the residual update with the SLERP-style update on the sphere (Eq. 10, 11); switch the attention softmax scale from $1/\sqrt{d_k}$ to $\sqrt{d_k}$ and rescale Q,K via $\mathbf{s}_{qk}$; rescale MLP intermediates via $\mathbf{s}_u, \mathbf{s}_\nu$; rescale logits via $\mathbf{s}_z$; remove weight decay and learning-rate warmup. The seventh point is non-trivial — the sphere constraint plays the role of weight decay, and the bounded-magnitude initialization removes the need for warmup.

Parameter accounting. RMSNorm modules deleted: $\sim 2 d_{\text{model}} L$ parameters removed. Added: per-block $\boldsymbol{\alpha}_A, \boldsymbol{\alpha}_M$ each in $\mathbb{R}^{d_{\text{model}}}$, per-head $\mathbf{s}_{qk}$ in $\mathbb{R}^{d_k}$, per-MLP $\mathbf{s}_u, \mathbf{s}_\nu$ in $\mathbb{R}^{d_{\text{MLP}}}$, and one global $\mathbf{s}_z$ in $\mathbb{R}^V$. Net change is small — the architecture has fewer multiplicative scales overall, since the per-row matrix scales (the “magnitude” half of weight = direction × magnitude) are removed by the unit-norm constraint.

§ 3 · Reference implementation

A normalized block, sketched

def n_block(h, attn, mlp, alpha_A, alpha_M):
    # h: [B, T, d]  — unit-norm on the hypersphere
    # alpha_A, alpha_M: learnable per-coord eigen learning rates in R^d_{>=0}

    h_A = norm(attn(h))                          # attention output, unit norm
    h   = norm(h + alpha_A * (h_A - h))          # LERP toward h_A + retract to sphere

    h_M = norm(mlp(h))                           # MLP output, unit norm
    h   = norm(h + alpha_M * (h_M - h))          # LERP toward h_M + retract
    return h

def n_attention(h, Wq, Wk, Wv, Wo, s_qk):
    # h is unit-norm. After projection q, k, v are NOT unit-norm; we re-normalize.
    q, k, v = h @ Wq, h @ Wk, h @ Wv             # [B, T, H, d_k]
    q, k = rope(q, k)                            # positional info before re-norm
    q = norm(q) * s_qk                           # learnable QK scale
    k = norm(k) * s_qk
    logits = (q @ k.transpose(-2, -1)) * sqrt(d_k)  # NB: scale is sqrt(d_k), not 1/sqrt
    return (softmax(logits + causal_mask) @ v) @ Wo

After each gradient step, every row of every weight matrix is renormalized to unit length in the input dimension. The forward pass relies on that invariant: all dot products are bounded cosine similarities, and the sphere-update equations are exact only when the inputs are unit vectors.

§ 4 · Empirical evidence

What the paper measured, and what it did not

Headline speedup. Figure 1 of the paper plots validation loss vs. training step for 1B parameter GPT and nGPT models at 4K context on OpenWebText (Gokaslan & Cohen 2019). nGPT reaches the same loss at 20K steps that the baseline reaches at 200K — a ~10× step reduction. Figure 2 extends this across 0.5B and 1B models at context lengths 1K, 4K, 8K. The reported acceleration factor in steps-to-fixed-validation-loss is approximately 4× at 1K context, 10× at 4K, and 20× at 8K. The trend (speedup grows with context length) is consistent across both model sizes.

Wall-clock caveat. The paper is explicit that per-step time is higher for nGPT than for the baseline — 80% slower per step at 4K context and 60% slower at 8K (footnote 3 in § 3.1), before kernel optimization. The 10–20× step speedup therefore translates to a roughly 5–10× wall-clock speedup with unoptimized kernels. The footnote claims this overhead can be reduced; no kernel-optimized comparison is reported in the paper.

Downstream tasks. Figure 3 shows the speedup carries over from validation loss to downstream accuracy on five tasks: ARC-Easy, HellaSwag, WinoGrande, WSC273, LAMBADA-OpenAI seqlen 1024. At equal token budget, nGPT 1B at 4K context achieves higher accuracy on each task than GPT 1B; at equal accuracy, nGPT reaches the target at 4–10× fewer training tokens.

Diagnostics of the sphere constraint. Figure 4 inspects what unconstrained GPT embeddings look like at the end of training and confirms they are far from norm-uniform: the input-embedding norms in 1B GPT span roughly 0.7 to 1.7, with a heavy-tailed eigenvalue spectrum, while nGPT’s are pinned to 1 by construction. Figure 5 plots condition numbers of Q, K, V, and the three MLP matrices across layers; GPT’s attention matrices show condition numbers in the $10^2$–$10^4$ range with high variance, while nGPT’s stay below $10^1$ with no obvious depth trend. The authors interpret this (§ 3.2) as evidence that the sphere constraint prevents the rank-deficiency that GPT’s attention matrices drift toward during training.

Ablations. Section 3.3 and Appendix A.9 report that fixing $\mathbf{s}_{qk}, \mathbf{s}_u, \mathbf{s}_\nu$ as non-learnable and using a single global $\mathbf{s}_z$ gives only “slight degradation in accuracy” — the general per-coordinate form is not strictly necessary. Appendix A.9 also tests omitting the explicit Q,K normalization (since the rows of $W_q, W_k$ are already unit-norm, $\mathbf{q}, \mathbf{k}$ have bounded norm in expectation); this gives a 12% compute saving with “only a minor performance degradation” at training-context length, though it hurts on long-context extrapolation (Appendix A.8).

What is not in the paper. No 7B+ run; no production-scale training. No comparison with non-Adam optimizers. No analysis of how the unit-norm-after-each-step constraint interacts with mixed-precision training or with fully sharded data parallelism. The headline 4–20× speedup is observed on a single dataset (OpenWebText) at two model sizes (0.5B, 1B); the authors note in § 3.1 that “we observe some saturation for the longest runs of nGPT, suggesting that the model capacity is nearly reached for this number of trainable model parameters” — implying the gap may narrow at scale.

Independent reproductions. As of mid-2026, no production open-weight model has shipped the full nGPT recipe. The geometric framing has attracted theoretical follow-up (representation-on-sphere arguments by Wang & Isola, ICML 2020 predate the work and provide motivation), but a public end-to-end reproduction at the 7B scale or larger does not exist. The kernel question remains: standard GEMM kernels do not natively support per-row weight normalization at every step, and the wall-clock speedup depends on whether fused implementations can close the per-step overhead the paper acknowledges. Verdict: clean construction, strong small-scale evidence, no public scale-up demonstration.

@article{arxiv2410_01131,
  title  = {nGPT: Normalized Transformer with Representation Learning on the Hypersphere},
  author = {Ilya Loshchilov and others (NVIDIA)},
  year   = {2024},
  eprint = {2410.01131},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2410.01131}
}

@article{arxiv_2410_01131,
  title         = {nGPT: Normalized Transformer with Representation Learning on the Hypersphere},
  author        = {Ilya Loshchilov et al. (NVIDIA)},
  year          = {2024},
  eprint        = {2410.01131},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2410.01131}
}

{
  "id": "arxiv_2410_01131",
  "type": "article-journal",
  "title": "nGPT: Normalized Transformer with Representation Learning on the Hypersphere",
  "author": [
    {
      "literal": "Ilya Loshchilov et al. (NVIDIA)"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2024
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/2410.01131",
  "number": "2410.01131",
  "source": "arXiv"
}

TY  - JOUR
TI  - nGPT: Normalized Transformer with Representation Learning on the Hypersphere
AU  - Ilya Loshchilov et al. (NVIDIA)
PY  - 2024
JO  - arXiv
AN  - arXiv:2410.01131
UR  - https://arxiv.org/abs/2410.01131
ER  -