PaperlayerView exact revision on arXiv

Source-linked AI summary

SimpleGPT: Improving GPT via A Simple Normalization Strategy

Marco Chen, Xianbiao Qi, Yelin He, Jiaquan Ye, Rong Xiao

arXiv:2602.01212v1cs.LGcs.CLcs.CV

TL;DR

Transformer optimization stability becomes increasingly limiting as models scale, while the link between architectural choices, activation scale, Hessian geometry, and learning rates remains unclear. SimpleGPT introduces SimpleNorm and shows that it supports stable learning rates 3×–10× larger than strong baselines while achieving stronger performance across 1B–8B models.

  • Problem

    The relationship between Transformer architectural heuristics, activation scale, Hessian geometry, and stable learning rates remains theoretically underexplored despite its importance for scaling.

  • Method

    SimpleGPT places normalization immediately after linear mappings and analyzes the resulting activation Hessian to connect reduced curvature with larger stable learning rates.

  • Results

    SimpleGPT consistently outperforms GPT baselines across 1B–8B models and, at 7B scale after 60K steps, lowers LLaMA2 with QKNorm training loss by 0.08, from 2.290 to 2.208.

  • Takeaways & Limitations

    SimpleNorm provides a theoretically motivated architectural route to smoother activation loss landscapes and substantially larger admissible learning rates.

  • Takeaways & Limitations

    The theoretical guarantees assume sufficiently large dimensions and exclude pathological or adversarial alignment between the mapping, derivatives, and leading Hessian eigenspaces.

Abstract

from arXiv · show

In this work, we revisit Transformer optimization through the lens of second-order geometry and establish a direct connection between architectural design, activation scale, the Hessian matrix, and the maximum tolerable learning rate. We introduce a simple normalization strategy, termed SimpleNorm, which stabilizes intermediate activation scales by construction. Then, by analyzing the Hessian of the loss with respect to network activations, we theoretically show that SimpleNorm significantly reduces the spectral norm of the Hessian, thereby permitting larger stable learning rates. We validate our theoretical findings through extensive experiments on large GPT models at parameter scales 1B, 1.4B, 7B and 8B. Empirically, SimpleGPT, our SimpleNorm-based network, tolerates learning rates 3$\times$-10$\times$ larger than standard convention, consistently demonstrates strong optimization stability, and achieves substantially better performance than well-established baselines. Specifically, when training 7B-scale models for 60K steps, SimpleGPT achieves a training loss that is 0.08 lower than that of LLaMA2 with QKNorm, reducing the loss from 2.290 to 2.208. Our source code will be released at https://github.com/Ocram7/SimpleGPT.

1. Introduction

The introduction frames Transformer scalability as an optimization-stability problem and connects architectural design, activation scale, Hessian geometry, and learning-rate limits. It presents SimpleNorm and SimpleGPT as theoretically motivated solutions with larger stable learning rates and empirical gains across model scales.

  • As Transformer depth and width increase, optimization stability increasingly constrains performance and scalability.
  • The authors argue that architectural heuristics should be connected explicitly to classical Hessian-based optimization theory.Existing techniques include normalization placement, residual scaling, and modified nonlinearities.
  • SimpleNorm stabilizes intermediate activation scales and significantly reduces the activation-loss Hessian spectral norm, enabling substantially larger stable learning rates.SimpleGPT is the architecture based on SimpleNorm.
  • SimpleGPT achieves consistent empirical gains across nanoGPT, LLaMA2, and LLaMA3 from 1B to 8B parameters.These experiments test whether the theoretical optimization advantages translate across architectures and scales.
  • 0.08 lower training loss: SimpleGPT reduces LLaMA2 with QKNorm from 2.290 to 2.208 after 60K steps at 7B scale.

2. Related Work

Prior work studies normalization operators, their placement, and normalization-free alternatives as ways to stabilize Transformer training. SimpleGPT distinguishes itself by explicitly connecting local normalization after linear mappings to activation scale, Hessian spectral norm, and optimization smoothness.

  • Normalization methods: Normalization methods stabilize optimization and improve convergence, with BN depending on mini-batch statistics while LN removes batch dependence by computing feature statistics.The supplied passage identifies BN and LN as central normalization approaches, contrasting their dependence on batch statistics.
  • Normalization Placement in Transformers: Transformer normalization placement affects optimization stability: PostNorm follows residual addition, whereas PreNorm precedes attention and MLP sublayers to improve trainability in deep networks.The passage describes a shift from the original Transformer’s PostNorm toward PreNorm in large-scale practice.
  • Normalization Placement in Transformers: DNT organizes InputNorm, PreNorm, MidNorm, PostNorm, and QKNorm using Jacobian- and gradient-stability analysis, combines four placements, and avoids PostNorm because of potential training instabilities.DNT combines InputNorm, PreNorm, MidNorm, and QKNorm while excluding PostNorm.
  • Normalization-free Transformers: Normalization-free Transformer methods replace explicit normalization with bounded point-wise transformations that control activation norms and therefore act as implicit normalization.The passage cites Dynamic Tanh and searched function families such as Derf as examples of this direction.
  • Positioning of our work: SimpleGPT’s key distinction is linking local normalization after linear mappings to stabilized activation scale, reduced Hessian spectral norm, and a smoother optimization landscape.This positioning differs from approaches motivated primarily by heuristic or empirical considerations.

3. Preliminaries

The preliminaries define convex smooth optimization through local curvature and quadratic upper bounds, then connect smoothness to gradient-descent learning-rate stability. They also introduce gradients, Hessians, and Jacobians for a linear projection and its loss.

  • The paper considers unconstrained convex optimization of a differentiable function f: R^d → R.
  • Second-order Taylor expansion describes local curvature, while β-smoothness formalizes the relevant smoothness condition.For convex differentiable functions, β-smoothness is equivalent to a quadratic upper bound.
  • The quadratic upper bound guides learning-rate selection in gradient descent and yields a sufficient condition for monotone objective decrease.This condition characterizes gradient descent’s stability region for convex β-smooth functions.
  • For a linear projection y = W x with loss ℓ, the preliminaries define the loss gradient and Hessian with respect to y, along with the Jacobian and Hessian with respect to x.

4. Methodology

SimpleNorm places normalization immediately after each linear mapping, stabilizing activation scales and making curvature less dependent on weight scale. These properties yield smoother optimization and permit larger, more stable learning rates, while the paper hypothesizes additional expressive benefits from pervasive normalization-induced nonlinearity.

  • SimpleNorm Definition: SimpleNorm applies a normalization operator immediately after each linear mapping, treating the pair as one unified operation.The strategy differs from normalization at residual-block, hidden-state, or parameter-reparameterization levels.
  • SimpleGPT Definition: SimpleGPT inserts SimpleNorm after MLP, attention, output, gating, and memory-related linear projections.Examples include the Wq, Wk, Wv, Wo, W1, and W2 projections.
  • Activation-Scale Stability: SimpleNorm stabilizes intermediate activation norms at Θ(d), preventing drift with depth or weight growth and eliminating a source of activation explosion.The construction normalizes after each linear mapping, with learned per-dimension scaling γ.
  • Hessian Analysis: SimpleNorm decomposes the activation Hessian into Gauss–Newton and curvature terms, with the Gauss–Newton term dominating w.h.p. under high-dimensional, non-pathological conditions.The analysis establishes Hxx = L + C and assumes high effective rank and no pathological alignment.
  • Hessian Analysis: SimpleNorm makes curvature invariant to weight spectral norm, whereas plain linear-projection curvature scales quadratically with weight norm.This removes the dependence of curvature on growing weight scale and produces a smoother loss landscape.
  • Optimization and Expressivity: SimpleNorm enables larger and more stable learning rates by constraining the activation-Hessian spectral norm.The paper also hypothesizes that normalization-induced nonlinearity increases expressive capacity without increasing parameter count.

5. Experiments

Experiments across nanoGPT, Llama2, and Llama3 show that SimpleGPT supports larger learning rates, improves training loss across model scales, and incurs only a modest training-speed slowdown. Its gains remain stable across architectures and training progress.

  • Learning-rate stability: SimpleNorm remains stable at learning rates of 2 × 10−3 and 2 × 10−2, whereas PreNorm shows convergence issues at 2 × 10−3 and PreNorm+QKNorm becomes unstable at 2 × 10−1.Stability is compared while keeping other training settings fixed.
  • Llama2 1B: 0.032 absolute improvement: SimpleGPT 1B reduces training loss from 2.478 to 2.446 versus Llama2 1B with QKNorm.The models are trained for 200K steps on C4, covering approximately 26B training tokens.
  • Llama2 7B: 0.082 improvement: SimpleGPT 7B outperforms Llama2 7B with QKNorm at 60K steps, following gains of 0.062 at 20K and 0.077 at 40K.SimpleGPT 7B uses a 0.001 learning rate, 3× larger than Llama2 7B; the models train on C4 for up to approximately 24B tokens.
  • Llama3 8B: At 8B scale, SimpleGPT uses a 3× larger learning rate than Llama3 8B and achieves substantially lower training loss after 20K steps.Both models use the Llama3 8B architecture, train on C4, and process approximately 8B tokens.
  • nanoGPT 1.4B: 0.043 approximate improvement: SimpleGPT 1.4B outperforms GPT-2 with QKNorm on nanoGPT, while GPT-2 with QKNorm performs nearly identically to the original GPT-2.All models train for 100K steps, or approximately 50B tokens, and SimpleGPT uses a learning rate 3× larger than baseline.

6. Conclusion … C.2. Part II: derivatives w.r.t.W

The paper links Transformer architecture, activation-scale control, Hessian geometry, and optimization stability, while deriving first- and second-order derivatives for SimpleNorm and linear weight mappings. These analyses support larger learning rates through reduced activation Hessian norms and express weight Hessians via Kronecker structure.

  • 6. Conclusion: SimpleNorm reduces the activation Hessian norm, enabling substantially larger admissible learning rates and allowing SimpleGPT learning rates up to 3×-10× larger than strong baselines.The conclusion connects this Hessian reduction directly to optimization stability and learning-rate tolerance.
  • A. Derivatives of ∇xℓand ∇2 xℓfor SimpleNorm: The SimpleNorm activation Hessian decomposes into linear Gauss–Newton and curvature terms induced by normalization.The derivation applies the chain rule through the normalization path and combines the two terms into the Hessian with respect to x.
  • B. Derivatives of ∇W ℓand ∇2 vec(W )ℓfor y = W x: For y = W x, the first-order weight derivative follows from the linear variation dy = dW x and Frobenius-inner-product coefficient matching.The setup treats W ∈ R^d×m, x ∈ R^m, and y ∈ R^d for a scalar loss ℓ(y).
  • B. Derivatives of ∇W ℓand ∇2 vec(W )ℓfor y = W x: The Jacobian of y with respect to vec(W) is x^T⊗I_d, making the dependence on vectorized weights explicit.This representation is used to derive the second-order weight expression.
  • B. Derivatives of ∇W ℓand ∇2 vec(W )ℓfor y = W x: Because W → y is linear, the weight Hessian has no mapping-induced second-order term and factorizes into the input term xx^T and output curvature H_yy.The equivalent block representation uses m × m blocks of size d × d, with entries determined by x_i.
  • C. Derivatives of ∇γℓ, ∇2 γℓ, ∇W ℓand ∇2 vec(W )ℓwhere y = γ ⊙: For the combined scaling-and-normalization mapping, derivatives are obtained through W → z → u → y → ℓ and γ → D → y → ℓ.The intermediate variables include s, P = I − uu^T, and D = Diag(γ).
  • C.2. Part II: derivatives w.r.t.W: The w.r.t. W derivation first obtains the gradient through z and then applies the second-order chain rule, separating Gauss–Newton and normalization-curvature contributions.The curvature contribution is expressed using the bilinear form of the normalization Hessian and its associated matrix form.

D. Proof of Theorem 4.1

The proof decomposes Hxx into Gauss–Newton and curvature terms, bounds the Gauss–Newton contribution under high effective rank and non-pathological alignment, and shows it dominates in high dimensions. The result completes Theorem 4.1.

  • The chain-rule decomposition separates Hxx into a Gauss–Newton term L and a curvature term C.
  • Bounding the Gauss–Newton term: κ = Θ(1) with high probability when the effective rank satisfies reff(W) ≍ cd.This follows for isotropic or subgaussian inputs via concentration for quadratic forms.
  • Bounding the Gauss–Newton term: ∥L∥2 = τκ2∥Hyy∥2 where τ = Θ(1) and κ = Θ(1) [w.h.p].The τ = Θ(1) condition follows from the theorem’s non-pathological alignment assumption.
  • Dominance of L over C: In high dimension, ∥L∥2 ≫ ∥C∥2 with high probability, so the Gauss–Newton term dominates Hxx in typical non-pathological regimes.The comparison assumes τ = Θ(1) and ∥gy∥2/∥Hyy∥2 = O(1).

E. Proof of Theorem 4.2

Under Theorem 4.1’s high-dimensional, non-pathological conditions, the normalized Hessian is dominated by its Gauss–Newton term. The proof then shows that the linear Hessian grows with the weight spectral norm while the normalized Hessian is independent of it, yielding a much larger Lipschitz constant for the linear module.

  • Hessian decomposition: The Hessian with respect to x decomposes into a Gauss–Newton term and a normalization-induced curvature term.The curvature term is denoted C, while the Jacobian contribution is represented through Jy2 x.
  • Dominance under assumptions: Under Theorem 4.1’s stated high-dimensional conditions, the Gauss–Newton term dominates, with ∥L∥2 ≫ ∥C∥2.The assumptions include high effective rank, no pathological alignment, and ∥gy∥2/∥Hyy∥2 = O(1).
  • Weight-scale comparison: The linear Hessian’s spectral norm is directly proportional to ∥W∥2, whereas the normalized Hessian’s spectral norm is independent of ∥W∥2.The comparison assumes W = W1 = W2 and Hy1y1 = Hy2y2 := Hyy.
  • Weight-scale comparison: When α = ∥W∥2 grows to tens or hundreds during training, α^2c_lin ≫ κ^2, so the linear Hessian is much larger with high probability.Here κ = Θ(1) with high probability, while c_lin = Θ(1) under the stated alignment condition.
  • Theorem conclusion: In non-pathological cases, the gradient Lipschitz constant of ∥Hlin_xx∥2 is much larger than that of ∥Hsn_xx∥2.This establishes the claimed comparison and completes the proof of Theorem 4.2.

F. Detailed Experimental Settings · G. Parameters and configurations of SimpleGPT

SimpleGPT applies SimpleNorm across Transformer blocks on Llama2, Llama3, and nanoGPT backbones, with backbone-specific architectures and shared optimization and hardware settings. Configurations span 1B, 1.4B, 7B, and 8B models mapped to these backbones.

  • F. Detailed Experimental Settings: SimpleNorm is applied to every Transformer block except the embedding and classification layers across Llama2, Llama3, and nanoGPT-based SimpleGPT models.Implementations are based on Adam-mini and nanoGPT.
  • F. Detailed Experimental Settings: nanoGPT uses GELU activations, GPT-2-compatible BPE tokenization, a 50,257-token vocabulary, 2,000 warmup steps, and OpenWebText training data.nanoGPT is described as a lightweight GPT-2 implementation.
  • F. Detailed Experimental Settings: Llama2 uses SwiGLU feed-forward activations, RoPE positional encoding, GQA, and a SentencePiece-based BPE tokenizer with a 32K-token vocabulary.The passage associates GQA with reduced inference-time memory and computational cost.
  • F. Detailed Experimental Settings: Llama3 retains Llama2’s dense Transformer design and GQA with eight key-value heads, while using a 128K-token multilingual tiktoken-based vocabulary.The larger vocabulary is described as improving compression rates and language coverage.
  • F. Detailed Experimental Settings: All experiments use AdamW with β1 = 0.9, β2 = 0.95, and weight decay 0.1 unless otherwise stated.Weight decay is adjusted in relation to the larger learning rates enabled by the method.
  • F. Detailed Experimental Settings: Models are trained in PyTorch with bfloat16 precision on A800 GPUs using a cosine learning-rate schedule for every training run.These settings apply across the reported training runs.
  • G. Parameters and configurations of SimpleGPT: Table 1 summarizes model configurations for SimpleGPT at different scales, including the 1B and 7B models based on Llama2.The supplied table passage identifies these scale-to-backbone assignments.
  • G. Parameters and configurations of SimpleGPT: The 8B SimpleGPT model is based on Llama3, whereas the 1.4B model is based on nanoGPT.These assignments complete the backbone mapping given for the listed model scales.

H. More experiments on SimpleGPT 7B · I. More experiments on SimpleGPT 8B

Additional experiments evaluate SimpleGPT 7B and 8B under different learning rates and weight decay values. The reported comparisons use training-loss curves against Llama2 or Llama3 baselines, with and without QKNorm, at learning rates 3e-4 and 3e-3.

  • H. More experiments on SimpleGPT 7B: SimpleGPT 7B is evaluated under varied learning rates and weight decay settings.Results are presented in Figures 8 and 9.
  • H. More experiments on SimpleGPT 7B: At learning rate 3e-4, Figure 8 compares training-loss curves for Llama2 7B, Llama2 7B with QKNorm, and SimpleGPT 7B.The figure also varies weight decay.
  • H. More experiments on SimpleGPT 7B: At learning rate 3e-3, Figure 9 compares training-loss curves for Llama2 7B, Llama2 7B with QKNorm, and SimpleGPT 7B.The figure also varies weight decay.
  • I. More experiments on SimpleGPT 8B: SimpleGPT 8B is evaluated under varied learning rates and weight decay settings.Results are presented in Figures 10 and 11.
  • I. More experiments on SimpleGPT 8B: At learning rate 3e-4, Figure 10 compares training-loss curves for Llama3 8B, Llama3 8B with QKNorm, and SimpleGPT 8B.The figure also varies weight decay.
  • I. More experiments on SimpleGPT 8B: At learning rate 3e-3, Figure 11 compares training-loss curves for Llama3 8B, Llama3 8B with QKNorm, and SimpleGPT 8B.The figure also varies weight decay.

J. More experiments on weight decays

Experiments with two weight decay values show that SimpleGPT 8B consistently outperforms LLaMA2 8B with QKNorm, indicating that SimpleNorm’s benefits are robust to regularization choices.

  • J. More experiments on weight decays: SimpleGPT 8B consistently outperforms LLaMA2 8B with QKNorm across two tested weight decay values.The experiments evaluate robustness to regularization in large-scale training.
  • J. More experiments on weight decays: The results indicate that SimpleNorm’s benefits are not sensitive to the choice of weight decay.Results are shown in Figure 12.
  • J. More experiments on weight decays: Figure 12 captions an overall comparison across Llama3 8B with QKNorm and SimpleGPT 1B under two different weight decay values.The caption identifies the compared models and the regularization settings.
Loading 2602.01212v1…