Research map / Architectures and efficiency / RNNs and normalization
Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes computing per-example gradient norms simultaneously with parameter gradients for linear, embedding and LayerNorm layers, finds LayerNorm GNS tracks total GNS in transformers, and uses a zero-overhead LayerNorm CUDA kernel to guide a batch size schedule that cuts training time 18%.
Problem
Gradient noise scale (GNS) estimation needs per-example gradient norms, which are costly to obtain in 3D tensor regimes (transformers) without affecting training throughput, and the GNS estimate is noisy.
Approach
Compute per-example gradient norms in the same einsum contraction that produces the weight gradient, avoiding the O(T^2) intermediate of prior tricks; observe that LayerNorm-only GNS predicts total model GNS, so a fused LayerNorm kernel can track GNS with zero throughput overhead.
Key results
- LayerNorm GNS has regression slope near 1.4 against total GNS and only overestimates total GNS by under 40% across EMA alpha values.
- Custom LayerNorm kernel computing per-example gradient norms shows practically zero overhead versus PyTorch LayerNorm on an H100.
- Linear batch size schedule guided by LayerNorm GNS reduces training time by 18% on a 111M Chinchilla-optimal model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a GNS-guided linear batch size schedule instead of a fixed batch size | 111M language model · OpenWebText | training time: 18% | Helped |
| Use a linear batch size schedule instead of a fixed batch size | 111M language model · OpenWebText | loss at matched tokens processed | Helped |
| Compute per-example gradient norms during the LayerNorm backward pass with a custom kernel | LayerNorm · LayerNorm forward and backward pass | throughput overhead: zero | Helped |
| Estimate total GNS using LayerNorm GNS | 111M language model · OpenWebText | total GNS prediction: approximately 1.4 | Helped |
| Increase the learning rate (ablation) | 111M language model · OpenWebText, temperature experiment | GNS | Mixed |
| Change the batch size (ablation) | 111M language model · OpenWebText, temperature experiment | GNS | No clear effect |
| Use Flash attention instead of the pure PyTorch attention implementation (ablation) | 111M language model in bfloat16 AMP · OpenWebText | training stability | Hurt |
| Use cosine attention in place of Flash attention (ablation) | Student model · Teacher-student stability simulation | training stability | Helped |
| Use float32 precision in block 1 (the second block) (ablation) | 111M language model · OpenWebText | training stability | Helped |
| Apply spectral normalization to the QKV projection in block 1 (the second block) (ablation) | 111M language model · OpenWebText | training stability | Helped |
| Vary model size among 70M, 111M, and 161M parameters at constant FLOPs (ablation) | Cerebras-GPT training recipe · OpenWebText | loss: 70M, 111M and 161M | Mixed |
| Use only LayerNorm layers for GNS estimation instead of collecting all layer norms (ablation) | 1.3B GPT · OpenWebText | MFU: 40% -> 57% | Helped |
Surprising findings
- Changing batch size does not produce the GNS change predicted by the temperature theory, although learning-rate changes do.
- Flash attention in bfloat16 caused divergence in the 111M model, mitigated by cosine attention or spectral normalization in block 1.
Limitations
- Only Transformers with LayerNorm sublayers were studied; RNNs, CNNs and state-space models were not evaluated.
- The 1.3B-scale result required a constant-factor correction after a bug was discovered, and the experiment was not rerun.
- Simultaneous method is more I/O-expensive for large models with short sequences.
Experimental setup
Datasets: OpenWebText
Benchmarks: OpenWebText
Metrics: gradient noise scale correlation coefficient regression slope tokens saved throughput (MFU) kernel time overhead
Baselines: Backpack per-example gradient method DDP-based GNS estimation PyTorch LayerNorm
Builds on: gradient noise scale Backpack per-example gradient trick Cerebras-GPT nanoGPT
Closest papers by meaning
- Peri-LN: Revisiting Normalization Layer in the Transformer Architecture (ICML 2025 · similarity 0.85)
- HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization (NeurIPS 2025 · similarity 0.84)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.84)
- Subcritical Signal Propagation at Initialization in Normalization-Free Transformers (NeurIPS 2026 · similarity 0.84)
- Learning in Compact Spaces with Approximately Normalized Transformer (NeurIPS 2025 · similarity 0.84)
- Global Convergence in Training Large-Scale Transformers (NeurIPS 2024 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.