Near-Zero Token Embeddings Silently Stall Training in Gated Sequence Models

AI summary

Overview

Research area: Machine learning systems and training stability — specifically the numerics of pretraining and continual pretraining for gated/hybrid sequence architectures (Mamba, Gated DeltaNet, xLSTM) and their Transformers counterparts.

Technical level: Advanced. The paper combines a large-scale empirical audit of open-weight checkpoints, a formal two-regime gradient-scaling proposition, IEEE-754 floating-point analysis of gradient clipping, and controlled training experiments.

Scope (one sentence): The paper shows that near-zero input-embedding rows — created by AdamW weight decay on untied embeddings — cause gradient amplification of up to 35 orders of magnitude in gated models, silently zeroing entire batch updates via FP32 clipping overflow, and proposes two zero-overhead fixes (GroupClip for released checkpoints, GateFloor for pretraining).

What This Paper Is About

Many open-weight language models store separate input and output embedding matrices ("untied" weights). When a token never appears in pretraining text, its input row gets zero gradient while AdamW weight decay keeps shrinking it — producing "near-zero tokens" with norms below 10⁻⁶ (down to 10⁻³⁶ in the surveyed checkpoints). In standard Transformers, encountering such a token mid-sequence is harmless. In gated sequence models, it is not: the current-token gate or query vanishes at zero input, the residual stream collapses, Pre-RMSNorm divides by its floor, and gradients compound across depth until FP32 global-norm clipping overflows to infinity and zeroes every gradient in the batch. The paper's goal is to diagnose this failure precisely and to fix it without extra step time or memory.

Key Contributions

1. A prevalence audit across 291 open-weight HuggingFace repositories (280 distinct embedding tables). Reachable near-zero input rows appear in 52.6% of untied gated/hybrid checkpoints (20 of 38) and 30.9% of untied Transformers (47 of 152), but in 0% of weight-tied models (0/80 Transformers, 0/21 gated). The authors trace these rows to untied input tables receiving zero gradient under AdamW weight decay.

2. A mechanism and scaling law for gated architectures. Proposition 1 formalizes a two-regime behavior for the input-embedding gradient as the token norm r → 0: an inverse-norm regime (∝ 1/r) when per-layer gains lift the hidden norm past √(εd) before the first Attention layer, and a depth-saturated regime (ceiling G_K) when Attention rescues the state first. Pure gated stacks (K = ∞) follow 1/r across the tested range down to r ≥ 10⁻³⁵.

3. Identification of the silent failure mode in FP32 gradient clipping. Squaring any gradient entry above 1.84 × 10¹⁹ overflows FP32; in IEEE-754 arithmetic, dividing finite gradients by +∞ yields exact 0.0 with no NaN, so PyTorch, HuggingFace Trainer, and FSDP2 silently discard all current-batch gradients.

4. Two remedies with no step-time or peak-memory overhead. GroupClip decouples input-embedding clipping (using a 64-bit or rescaled norm) from backbone clipping in three lines of code; GateFloor adds a positive gate offset (SiLU(z) + δ, δ ∈ [0.01, 0.1]) for pretraining from scratch.

Main Findings

• Weight tying is the dividing line. 0 of 101 weight-tied checkpoints contain reachable near-zero tokens, versus 67 of 190 untied checkpoints (35.3%). Across all 291 repositories, 92 contain raw rows below 10⁻⁶ and 67 (23.0%) contain reachable near-zero tokens.

• The affected tokens are ordinary strings. Examples include Telugu and Basque subwords in Nemotron, Cyrillic Mongolian words in Falcon-H1, CRLF code fragments such as }); r in Mamba-Codestral, chat tags <|im_user|> and <|im_end|> in Kimi-Linear-48B-Base, and <s> in Nemotron-H-4B-Base. They occupy only 0.001%–0.25% of the vocabulary (median 0.040%) among the 20 affected untied gated models.

• Weight decay explains the decay exactly. Under untied weights with a zero-gradient input row, the norm follows ‖e_v^(T)‖₂ = ‖e_v^(0)‖₂ ∏ₜ(1 − ηₜλ). On OLMoE-1B-7B, a 364M-token sample across all seven subsets found all 251 near-zero rows had zero occurrences (0/251), while observed tokens retained norms in [0.11, 0.60]. In the checkpoint, 250 of 251 rows cluster in [3.23, 3.50] × 10⁻¹² (log-decay ≈ −26.3, plus one outlier at 6.5 × 10⁻⁷), consistent with cumulative decay Σₜ ηₜλ ≈ 26–27 over OLMoE's 5.13T-token schedule (2,000-step warmup, η_max = 4 × 10⁻⁴, λ = 0.1).

• Output rows stay healthy. 92% of 979 sampled reachable near-zero tokens (901 of 979) retain output-head norms at least 0.3× the vocabulary median — which is why standard inference benchmarks never surface the problem.

• Transformers are structurally immune at non-initial positions. On Llama-3.1-8B, mid-sequence gradients stay bounded at 56–98 across r ∈ [10⁻³⁰, 10⁰]; RWKV-4-1.5B stays at 8.5. Layer 1 Attention averages normal-scale value vectors from prior tokens and restores scale in one layer.

• Gated models amplify by over 30 orders of magnitude. In the mid-sequence sweep, ceiling values scale with the number K of gated blocks before Attention: 2.3 × 10² on Bamba-9B-v2 (K = 9), 6.2 × 10² on LFM2.5-1.2B (K = 2), 2.8 × 10¹² on AI21-Jamba-Reasoning-3B (K = 7), 2.5 × 10¹⁵ on Zamba2-1.2B (K = 5), 2.2 × 10¹⁶ on Granite-4.0-H-1B (K = 5), and 1.3 × 10³⁰ on Falcon-Mamba-7B. Per-layer gains fall in γ_ℓ ∈ [10^0.3, 10^3.3].

• Pre-RMSNorm compounds the blowup by 10³× per layer. At ε = 10⁻⁶ the normalization divides by 10⁻³, multiplying backward gradients by 1,000× per layer (316× at ε = 10⁻⁵), compounding past 10³⁰ across depth.

• The deepest case is measured precisely. Nemotron-3-Nano-4B (K = 7 Mamba blocks before Layer 12 Attention, ε = 10⁻⁵) follows ≈ 0.18/r across [10⁻³⁰, 10⁰], peaking at 4.5 × 10³⁴ (2.1 × 10³⁵ on native tokens). Parallel hybrids behave differently: Falcon-H1-1.5B stays near 3.5 × 10⁶ because Attention and Mamba-2 are summed per layer.

• A causal gate-patching experiment rules out the "severed gradient" alternative. Detaching Layer 0's gate backward alone leaves ‖∇ℒ‖∞ at 1.7 × 10³¹ on Nemotron-3-Nano-4B. Opening only Layer 0's gate forward straight-through to "the" — keeping all backward derivatives intact — lifts ‖h_t^(1)‖₂ to 0.61 (above √(εd) ≈ 0.18) and cuts the gradient from 2.1 × 10³⁵ to 6.71 (0.77 when applied to all Mamba gates).

• The overflow threshold is exact and the failure is silent. FP32 max is ≈ 3.40 × 10³⁸; any entry above √FLT_MAX = 1.84 × 10¹⁹ overflows on squaring. On overflow steps in Nemotron-3-Nano-4B, backbone weights still move (‖Δθ_other‖₂ ≈ 0.25) via stale momentum and weight decay, yet fresh-batch displacement is identically zero (0% of weights differ from the g_t = 0 case, versus 42.6% on clean steps). PyTorch clip_grad_norm_, HuggingFace Trainer, and FSDP2 log grad_norm = inf without error; Megatron-Core either aborts or zeroes gradients.

• Corpus trigger rates split cleanly by language. Across 8,000 documents from 15 corpora and 14 models, English web, English SFT, code, and math all give p = 0.0%. Multilingual and template workloads hit routinely: 10.6%–18.8% on Aya, 70.5%–97.4% on Thai, 95.4% on Nemotron-3-Nano-4B Telugu, 100% on Falcon-H1 Mongolian, and 100% at t = 0 for Kimi-Linear-48B-Base's <|im_user|> tag.

• The end-to-end funnel on Telugu CPT is quantified. Over 1,000 steps at 16,384 tokens/step, 0.047% of Nemotron-3-Nano-4B's vocabulary (62 rows) accounts for 32.3 ppm of tokens, triggers 22.5% of batches, and overflows FP32 on 88.1% of hit batches — 19.8% of all steps with 0% fresh-batch update.

• Some workflows are genuinely safe. Rank-16 LoRA on nine checkpoints yields normal adapter gradient norms (3.7–10.3) with 0/9 overflows, because E_in is frozen. Sampling 16,384 Mongolian and Telugu rollout tokens across four models generated 0 near-zero tokens during on-policy RL decoding, though full-parameter RL updates still trigger on prompts containing near-zero template tags.

• GroupClip removes the stalls and heals the rows. On Nemotron-3-Nano-4B Telugu (1,000 steps, 3 seeds), FP32 global clip stalls 198.3 ± 9.9 steps at 13.44 ± 0.17 PPL; GroupClip (FP64) reaches 0.0 stalled steps and 13.04 ± 0.19 PPL, matching Oracle Mask (13.04 ± 0.19) and matching baseline quality by step 800. All 30.7 ± 0.6 encountered near-zero rows self-heal on first contact (3.5 × 10⁻⁴ step), lifting rows from ≤ 2.3 × 10⁻¹⁰ to a median of 0.010. On Falcon-H1-1.5B Mongolian, GroupClip lowers step-100 perplexity from 40.40 to 37.47 and beats an oracle dead-row mask, ultimately reaching 15.58 ± 0.10 versus 15.68 ± 0.11 for FP32 global clip.

• The 300-step ablation confirms the alternatives fall short. FP32 Global Clip (57.0 ± 7.5 stalled), Skip-on-Inf (57.0 ± 7.5), and ZClip (57.0 ± 7.5) all stall identically; FP64 Global Norm still stalls 18.0 ± 1.0 steps; GradientStabilizer avoids stalls but drives v_t = inf. GroupClip (FP64) reaches 0.0 stalled and 15.95 ± 0.25 PPL versus Oracle Mask's 15.94 ± 0.26.

• GateFloor fixes the root cause at pretraining time. On a 370M-parameter 3:1 hybrid pretrained for 1.57B tokens, SiLU(z) + 0.1 reduces r = 10⁻³⁰ embedding gradients by over 28 orders of magnitude at both t = 64 and t = 0, with validation perplexity of 22.28 versus 22.45 for the baseline (single seed).

• The cost is negligible. On H100, GroupClip matches baseline step time and peak memory; the FP64 norm uses 3.1 GB transient memory (21.2 vs. 18.8 ms clip call), removable via a copy-free rescaled variant.

Methodology in Plain English

1. Audit the checkpoints directly. Rather than trusting benchmark scores, the authors used HTTP byte-range requests to read input embedding matrices, output matrices, and tokenizers from 291 HuggingFace repositories, without downloading full models. Any row with norm below 10⁻⁶ was then put through a "round-trip" test (decode the token ID, re-tokenize it alone and inside eight phrase templates) so that unused alignment padding — e.g. 50,257 → 50,304 — and unreachable byte entries were not miscounted. Only rows that survive the round-trip count as "reachable."

2. Test the proposed cause. The weight-decay explanation predicts a specific geometric decay curve. The authors checked it against OLMoE-1B-7B's open corpus (OLMoE-mix-0924), confirming that every near-zero row had zero occurrences in a 364M-token sample and that the observed decay constant matched the checkpoint's actual norms.

3. Sweep the input norm and measure the gradient. They rescaled a single mid-sequence embedding across r ∈ [10⁻³⁰, 10⁰] (and [10⁻³⁵, 10⁰] on three models) and recorded the input-gradient norm, producing the two-regime curves that Proposition 1 predicts. The theory is derived by Taylor-expanding Pre-RMSNorm and the gated block around h = 0.

4. Intervene causally. To show that forward signal loss — not just severed backward paths — drives the blowup, they decoupled forward gate values from backward derivatives, patching only Layer 0's forward gate with a real token's value.

5. Reproduce the numerical failure in isolation. Single-step, fresh-optimizer tests on Falcon-H1-1.5B compared an overflow trigger (G_emb = 1.7 × 10²⁰), a 10⁶ spike, and a clean "the" token across PyTorch, HuggingFace Trainer, FSDP2, and Megatron-Core.

6. Estimate real-world exposure. They sampled 8,000 documents per corpus across 15 corpora and 14 models to get per-document trigger probabilities, then converted these to 128-sequence batch hit rates via 1 − (1 − p)¹²⁸.

7. Run controlled fixes. GroupClip and GateFloor were validated with full continual-pretraining runs (3 seeds, 1,000 steps, 16,384 tokens/step, FP32 master weights, BF16 autocast, AdamW λ = 0.1, FineWeb-2 data, FLORES-200 devtest perplexity) plus a 10-method ablation.

Why This Matters

Impact on research. The paper identifies a failure mode that produces no exception, no NaN loss, and no visible degradation in forward activations — forward activations stay finite throughout. Because the corrupted signal is invisible to standard monitoring (grad_norm = inf is logged, not raised), prior results on multilingual continual pretraining of gated models may have silently absorbed stalled steps. The work also reframes embedding weight decay: what is a mild regularizer in Transformers is a latent hazard in gated architectures, and the paper shows that simply disabling weight decay does not fix an already-released checkpoint, since Step-1 blowup persists.

Real-world applications:

• Multilingual model adaptation. Telugu, Mongolian, Thai

Authors’ abstract

Among 291 surveyed open-weight language model repositories (280 distinct embedding tables), 52.6% of untied gated and hybrid checkpoints (20/38) contain reachable input embeddings with norms below $10^{-6}$, compared with 0% of weight-tied models. Such near-zero rows can arise during pretraining when unseen rare or template tokens receive zero input gradients while AdamW weight decay shrinks untied embeddings. In tested Transformer configurations, non-initial positions or a BOS token prevent this failure mode. In gated sequence models (Mamba, Gated DeltaNet, xLSTM), a near-zero token can amplify input-embedding gradients by more than 30 orders of magnitude (up to 35). At zero input, current-token readout factors (SiLU gates or linear queries) vanish, suppressing historical-state readout. Residual RMS values can then fall below the square root of the RMSNorm floor across consecutive layers, causing Pre-RMSNorm to amplify backward gradients by up to $10^3$ per layer. When a gradient entry exceeds $1.84 \times 10^{19}$, FP32 global-norm clipping in PyTorch, Hugging Face Trainer, and FSDP2 can overflow to infinity and zero the current batch's gradients. In Telugu continual pretraining of Nemotron-3-Nano-4B, this discards updates on approximately 20% of steps and leaves triggered rows frozen. We propose two remedies. For released checkpoints, GroupClip decouples input-embedding clipping using a 64-bit or rescaled norm; its rescaled implementation eliminates update loss and heals triggered rows in one step with negligible measured end-to-end overhead. For pretraining, GateFloor ($\mathrm{SiLU}(z) + δ$, $δ\in [0.01, 0.1]$) keeps the gate positive at zero input, reducing near-zero gradients by 28 orders of magnitude in experiments up to 370M parameters, with perplexity comparable to the baseline.