Bounded Hyperbolic Tangent: A Stable and Efficient Alternative to Pre-Layer Normalization in Large Language Models

AI summary

Overview

Research area: Natural Language Processing — Transformer architecture and normalization for large language model pretraining.

Technical level: Advanced. The paper combines a Chebyshev-inequality-based probabilistic input bound, a Jacobian comparison against RMSNorm, and a finite-depth variance propagation theorem, alongside large-scale pretraining experiments.

Scope: The paper proposes Bounded Hyperbolic Tanh (BHyT), a drop-in replacement for Pre-Layer Normalization (Pre-LN) that couples a bounded tanh transformation with data-driven input bounding to jointly improve training stability and throughput in LLM pretraining.

What This Paper Is About

Pre-Layer Normalization (Pre-LN), usually instantiated as RMSNorm, is the default normalization in modern LLMs, but it carries two problems: repeated per-token statistic computation adds latency and memory overhead, and a "curse of depth" causes hidden-state magnitude and variance to grow across layers, which can push the block Jacobian toward identity and make deeper layers behave like costly identity mappings. Efficiency-oriented alternatives such as Dynamic Tanh (DyT) remove that overhead but do not explicitly control residual-stream growth, so they remain fragile at depth. BHyT is proposed to address both problems at once by bounding the pre-tanh input with per-token statistics and by computing exact statistics only once per Transformer block.

Key Contributions

1. BHyT, a drop-in Pre-LN replacement. It combines a bounded tanh(·) transformation with data-driven input bounding, keeping the pre-tanh argument in a non-saturating range with high probability, thereby controlling depth-wise activation growth.

2. A finite-depth variance analysis. The paper shows that for a network of depth L, if the hyperparameters satisfy λ/κ < 1/√L, the output variance under BHyT is strictly smaller than under LayerNorm Scaling (LNS) at every layer 1 ≤ ℓ ≤ L. The paper notes that λ = 1 and κ = 10 satisfies this condition for depths L < 100.

3. A block-level variance approximation. BHyT computes the input variance exactly once at the first Pre-LN site, reuses it, and replaces the second exact variance computation (after attention) with a model-dependent approximation, reducing repeated reduction operations and memory movement.

4. Evaluation across three axes. Stability, performance, and efficiency are measured on Llama-374M and Llama-1B pretrained on C4, plus a 20B-token Llama-3B comparison against Peri-LN as the strongest stability-oriented baseline. Code is released at https://github.com/MLAI-Yonsei/BHyT.

Main Findings

• Stability: On Llama-1B after 20B-token C4 pretraining, BHyT keeps layer-wise activation magnitude and variance substantially more controlled across depth than RMSNorm and DyT, and comparable to stability-oriented normalization methods. RMSNorm shows clear growth in both statistics with depth, and DyT shows stronger depth-wise growth than BHyT even under its selected configuration.

• Theoretical guarantee holds at typical depths: The finite-depth bound (λ/κ < 1/√L implies smaller variance than LNS at every layer) is satisfiable with λ = 1 and κ = 10 for networks with depth L < 100, which the paper describes as typical model depths.

• Pretraining-only performance: BHyT achieves the lowest PT train loss, PT eval loss, and perplexity at both scales. On Llama-374M: train loss 3.519, eval loss 3.207, eval PPL 24.714, average downstream accuracy 40.31 ±0.11, outperforming the second-best method by 0.3 points. On Llama-1B: train loss 2.828, eval loss 2.802, eval PPL 16.470, average accuracy 43.42 ±0.25, outperforming the second-best method by 0.6 points.

• Post-SFT performance: After fine-tuning on Lima1K for 15 epochs, BHyT achieves the lowest C4 PPL, lowest WikiText-2 PPL, and highest average accuracy at both scales. On Llama-374M: C4 PPL 27.40 ±0.13, WikiText-2 PPL 38.39 ±0.39, average accuracy 41.20 ±0.17 (0.2 points above second best). On Llama-1B: C4 PPL 19.37 ±0.15, WikiText-2 PPL 23.84 ±0.29, average accuracy 46.07 ±0.16 (1.4 points above second best).

• Larger-scale comparison against Peri-LN: On Llama-3B trained for 20B tokens, pretraining-only BHyT records train loss 2.756 and eval loss 2.760 with average accuracy 44.86 ±0.11, versus Peri-LN at 2.811, 2.812, and 42.08 ±0.25; the paper states BHyT outperforms Peri-LN by 3.4 percentage points in average downstream accuracy in this setting. After SFT, BHyT reaches C4 PPL 17.43 ±0.15, WikiText-2 PPL 21.33 ±0.20, and average accuracy 46.45 ±0.28, versus Peri-LN at 18.72, 23.04, and 44.24 ±0.20; the paper states a 2.5 percentage-point average accuracy advantage.

• Training throughput: On Llama-1B, BHyT reaches 83.0K training tokens/s, outperforming RMSNorm by 1.63% in training throughput. Across the paper, BHyT achieves an average of 1.6% faster training and an average of 1.77% higher token generation throughput compared to RMSNorm.

• Generation throughput: On Llama-1B with input length 512 and maximum new token length 512, BHyT achieves 1,199.9 ±3.0 tokens/s, exceeding RMSNorm (1181.1, -1.6%), LNS (1150.7, -4.1%), and Peri-LN (984.4, -18.0%). DyT remains faster at 1,352.0 ±11.9 tokens/s (+12.7%) because it avoids statistical computation entirely. At a maximum new token length of 128, BHyT records 1176.7 ±39.1 tokens/s.

• Variance approximation quality: Figure 2 compares approximated against actual activation variances for the second BHyT_MLP layer at layers 2, 5, 8, 11, 13, and 15 in Llama-1B, using 100 randomly sampled inputs from the C4 training corpus, with Pearson correlation coefficient r reported. The approximate variance closely tracks the empirical variance of intermediate hidden states; results for all 16 layers are in Appendix C.11.

Methodology in Plain English

The researchers start from the observation that a bounded activation like tanh can keep signal magnitudes in check, but only if its input is not allowed to drift into the saturated region where gradients vanish. To control that input, they use Chebyshev's inequality, a distribution-agnostic bound stating that for any random variable with finite mean and variance, most coordinates fall within κ standard deviations of the mean with probability at least 1 − κ⁻². This yields a scaling rule: divide the input by κ times its standard deviation, so the pre-tanh value stays inside a chosen range [−λ, λ] with high probability. For a target probability p, κ is set to (1−p)^(−1/2); the paper notes that p = 0.99 gives κ = 10.

The idealized version, BHyT*, rescales each input using its exact mean and standard deviation. Because computing those statistics at every normalization site would erase the efficiency benefit, the practical BHyT uses an RMSNorm-style zero-mean approximation at the first site in each block (before attention) and computes the input variance exactly there. For the second site (before the MLP), it avoids a second exact computation. Instead, it approximates the variance of the post-attention residual as the sum of the already-computed input variance and an estimate of the attention output variance. That estimate is derived under an idealized attention assumption (approximately uniform attention weights over a long sequence, zero-mean Gaussian activations and projections, and a near-linear tanh regime) and depends only on model-level quantities: the value and output projection weights, the sequence length, and the BHyT parameters. At inference these weights are fixed, so the term can be precomputed or cached; during training it is refreshed periodically rather than recomputed every step.

The paper also derives a Jacobian bound showing the BHyT Jacobian is a diagonally scaled version of the RMSNorm Jacobian, with ‖J^BHyT(x)‖₂ ≤ (λ/κ)‖J^RMS(x)‖₂, and then proves the finite-depth variance result against LNS. Empirically they pretrain Llama-3.2-style models (Llama-374M on 7.5B tokens, Llama-1B on 20B tokens, and Llama-3B on 20B tokens) on C4 using LlamaFactory, roughly 20 training tokens per parameter, then fine-tune on Lima1K for 15 epochs with five seeds. Downstream accuracy is measured in the 3-shot setting with lm-evaluation-harness on ARC-C, ARC-E, PIQA, HellaSwag, OpenBookQA (OBQA), Winogrande, MMLU, and BoolQ.

Why This Matters

Impact on research: BHyT reframes the normalization debate. Rather than choosing between stability-focused methods that add computation (Peri-LN, LNS) and efficiency-focused methods that abandon statistic computation (DyT), it argues that a bounded activation becomes a viable Pre-LN replacement only when its input range is tied to per-token statistics. This connects the empirical "curse of depth" literature to a concrete design principle and supplies a finite-depth variance guarantee that DyT does not offer.

Real-world applications:

• Large-scale LLM pretraining pipelines, where normalization reductions accumulate latency across depth and control of activation growth affects whether deep layers remain useful.
• Inference serving, where generation throughput directly determines cost per token; BHyT reports 1,199.9 tokens/s on Llama-1B versus 1181.1 for RMSNorm at the 512-token setting.
• Fine-tuning and transfer workflows, since the paper evaluates post-SFT accuracy plus C4 and WikiText-2 perplexity to check for forgetting and generalization beyond the pretraining distribution.
• Architecture research and ablation infrastructure, because BHyT is presented as a drop-in replacement with released code.

Industry relevance: The gains reported are small — 1.63% training throughput and 1.77% average generation throughput over RMSNorm — but they apply to the normalization layer that every Transformer block executes twice, so they scale with depth, model count, and total serving volume. The claim that matters for practitioners is the combination: better stability at depth at no throughput penalty relative to RMSNorm, and materially higher throughput than Peri-LN.

Future Directions

• Scaling beyond Llama-3B. The largest model the paper trains is Llama-3B on 20B tokens; whether the variance bound and stability advantage hold at substantially larger scales and longer token budgets is not established here.

• Sharper attention-variance approximations. The approximation in Theorem 3.4 rests on Assumption 3.3 (approximately uniform attention, zero-mean Gaussian activations and projections, near-linear tanh regime). Testing how the approximation degrades when these assumptions break — for instance with sparse or highly peaked attention — is an open question.

• Selection of λ and κ in practice. The paper notes the condition λ/κ < 1/√L is met by λ = 1, κ = 10 for depths under 100. How these should be chosen for deeper or shallower models, and whether the periodic refresh rate of the approximation term matters, are not resolved.

• Systematic comparison with concurrent normalization-free alternatives. The paper positions BHyT against Derf (an erf-based replacement for tanh in DyT) and SeeDNorm (dynamic RMSNorm rescaling) as concurrent work but does not report head-to-head experiments with them.

Target Audience

Researchers and engineers working on Transformer architecture, LLM pretraining infrastructure, and inference efficiency. It is most useful to readers with prior familiarity with Pre-LN, RMSNorm, DyT, and layer-wise variance behavior, who want a stability argument backed by a probabilistic input bound and a block-level efficiency mechanism rather than an empirical-only normalization swap. Practitioners evaluating normalization alternatives for a production pretraining run will find the throughput and downstream tables directly relevant; readers seeking results at frontier model scales will find the largest experiments are Llama-3B on 20B tokens.

Authors’ abstract

Pre-Layer Normalization (Pre-LN) is the de facto choice for large language models (LLMs) and is crucial for stable pretraining and effective transfer learning. However, Pre-LN incurs repeated statistical-computation overhead and remains vulnerable to the curse of depth, where hidden-state magnitudes and variances grow as the number of layers increases, destabilizing training. Efficiency-oriented normalization-free methods such as Dynamic Tanh (DyT) improve throughput but remain fragile at depth. To jointly address stability and efficiency, we propose Bounded Hyperbolic Tanh (BHyT), a drop-in replacement for Pre-LN. BHyT combines a tanh nonlinearity with explicit, data-driven input bounding to keep activations within a non-saturating range. It prevents depth-wise growth in activation magnitude and variance and provides a theoretical stability guarantee. For efficiency, BHyT computes exact statistics once per block and replaces a second normalization with a lightweight variance approximation. Empirically, BHyT demonstrates improved stability and efficiency during pretraining, achieving an average of 1.6\% faster training and an average of 1.77\% higher token generation throughput compared to RMSNorm, while maintaining strong pretraining-only and post-SFT performance across language understanding and reasoning benchmarks\footnote{Code is available at: https://github.com/MLAI-Yonsei/BHyT}.