Taming the Exploding Gradient: How Embedding Variability and Low-Rank Attention Stabilize LLM Training

If you have ever tried to train a deep neural network—particularly a Large Language Model (LLM)—you have likely encountered the nightmare of instability. You stare at the loss curve, watching it go down smoothly, and then suddenly, out of nowhere, it spikes. The loss diverges, the gradients explode, and days of compute time evaporate.

As we push for deeper models to achieve “scaling law” benefits, this instability becomes the primary bottleneck. The deeper the model, the harder it is to keep the signal clean from the input to the output.

In this post, we are diving deep into a fascinating paper titled “Stable Language Model Pre-training by Reducing Embedding Variability”. The authors propose two major contributions that change how we think about training stability:

1. TEV (Token Embedding Variability): A new, computationally cheap metric to monitor stability, replacing expensive gradient variance calculations.
2. MLRA (Multi-head Low-Rank Attention): An architectural change that mathematically guarantees lower variance in the forward pass, preventing the gradients from exploding in the first place.

Let’s unpack the math, the method, and the results.

The Stability Problem in Deep Transformers

To understand the solution, we first need to understand the mechanism of failure. Modern LLMs (like GPT-2, Llama-2, etc.) typically use a Pre-Layer Normalization (Pre-LN) architecture. While Pre-LN helps with optimization compared to Post-LN, it introduces a specific side effect: Gradient Explosion in shallower layers.

In a residual network (like a Transformer), the signal passes through a series of layers. If we denote the token embedding layer as \(\mathbf{E}\), it maps our vocabulary into a vector space.

Here, \(|V|\) is the vocabulary size and \(\mathbf{e}_i\) is the vector for a specific token.

During the backward pass (backpropagation), gradients flow from the final loss back to this embedding layer. Due to the residual connections, the gradients accumulate. The gradient at the very first layer (\(\nabla X_0\)) is the product of gradients through all subsequent layers.

In deep models (e.g., 48, 96, or 100+ layers), this product term grows exponentially. This phenomenon is known as gradient explosion. When gradients explode, the updates to the model weights become massive and erratic, leading to those loss spikes that ruin training runs.

The Cost of Monitoring

Traditionally, researchers monitor gradient variance to detect this instability. If the variance of the gradients is high, the training is unstable. However, calculating the gradient variance at every step is prohibitively expensive—it requires \(O(nd)\) operations for the gradient matrix, which slows down training significantly.

We need a better, cheaper speedometer.

Part 1: Token Embedding Variability (TEV)

The researchers propose that instead of looking at the gradients (which are expensive to compute), we can look at the weights of the embedding layer itself (which we already have).

Specifically, they observed a strong link between noisy gradients and the variability of token embeddings. If the training is unstable, the gradients updating the embedding layer will be erratic. This causes the embedding vectors for different tokens to fluctuate wildly in magnitude and spread.

Defining TEV

The authors introduce Token Embedding Variability (TEV). First, let’s look at the standard deviation of a single token’s embedding vector \(\mathbf{e}_i\).

Here, \(\bar{e}_i\) is the mean of the elements in that token’s vector.

To get a system-wide view, we calculate the mean (\(\mu_{\text{TEV}}\)) and standard deviation (\(\sigma_{\text{TEV}}\)) of this metric across the entire vocabulary \(|V|\).

Why does this work?

In a stable training run, token embeddings shouldn’t behave like outliers. The paper provides an empirical “sanity check” by looking at the mean values of embeddings in pre-trained open-source models (OPT, Pythia, Llama-2, GPT-2).

As shown in Figure 1 above, there is a clear trend: larger, better-performing models exhibit lower TEV. The “thinner” and lower distribution in the larger parameters (like Llama-2 70B or GPT-2 XL) indicates that as models become more stable and capable, the variability of their embedding weights decreases.

This confirms that \(\mu_{\text{TEV}}\) is a valid proxy for stability. If this number spikes, your gradients are likely exploding.

Part 2: The Solution – Multi-head Low-Rank Attention (MLRA)

Diagnosing the problem with TEV is useful, but fixing it is better. The authors propose an architectural change to the Multi-Head Attention mechanism called Multi-head Low-Rank Attention (MLRA).

The Concept

In a standard Transformer, the attention mechanism projects the input \(X\) using weight matrices \(W_Q, W_K, W_V\). These are typically full-rank square matrices.

MLRA proposes factorizing these projection matrices into two smaller, low-rank matrices. Instead of learning one big matrix \(W\), we learn two matrices \(W^U\) (Up-projection) and \(W^D\) (Down-projection), such that:

where \(W^U \in \mathbb{R}^{d_{\text{model}} \times r}\) and \(W^D \in \mathbb{R}^{r \times d_{\text{model}}}\), with \(r\) being the rank (\(r < d_{\text{model}}\)).

The Math: How Factorization Reduces Variance

This is the core innovation. Why does splitting a matrix in two help stability? It comes down to how variance propagates through initialized weights.

Let’s assume we use Kaiming Uniform Initialization, which is standard for these models. The variance of a weight matrix \(W\) under this initialization is \(\frac{1}{3d_{\text{model}}}\).

If we project an input \(X\) (assumed to be normalized) through a standard linear layer, the variance of the output is:

The top equation shows the standard attention variance is 1/3.

Now, look at the bottom equation for the factorized (MLRA) version. When the signal passes through \(W^D\) and then \(W^U\), the variances multiply. Because we are multiplying two matrices that both have small initial variances, the resulting variance is significantly dampened.

The result is 1/9.

By simply factorizing the matrix, the initial variance of the output is reduced by a factor of 3 (from 1/3 to 1/9).

This variance reduction acts as a “damper” on the signal. It prevents the exponential growth of variance as the signal propagates through deep layers, directly counteracting the gradient explosion problem described earlier.

Avoiding the Low-Rank Bottleneck

You might ask: Doesn’t low-rank factorization hurt the model’s expressivity?

This is a valid concern. Previous attempts at low-rank training often failed because the model couldn’t learn complex patterns (the “low-rank bottleneck”).

However, MLRA applies this factorization within the multi-head structure. Each head learns a different subspace. The authors illustrate that while individual matrices might be low-rank, their concatenation (or sum in the residual stream) preserves the full rank of the model.

As shown above, even if vectors are composed of simple basis vectors, combining them can span the full space. By applying factorization per head, MLRA maintains the necessary expressivity while enjoying the stability benefits of low variance.

Experiments and Results

The authors tested this hypothesis by training GPT-2 models from scratch using three configurations:

1. Baseline GPT-2
2. \(\sigma\)Reparam: A state-of-the-art method for stability (Zhai et al., 2023).
3. MLRA (Proposed): The factorized attention method.

They tested at varying depths: 48, 96, and 192 layers.

1. Stability Analysis (Gradient Variance)

The first question is: Does MLRA actually stabilize gradients?

Figure 2 tells a compelling story.

• Left (48 layers): All models are relatively stable, but MLRA (green line) has the lowest gradient variance.
• Middle (96 layers): The baseline GPT-2 (blue) starts to struggle with higher variance spikes. MLRA remains the most stable.
• Right (192 layers): This is the stress test. The baseline GPT-2 actually failed to train (diverged) 5 times and was excluded. MLRA, however, trained smoothly with extremely low gradient variance.

2. Validating TEV as a Proxy

Does the cheap TEV metric actually correlate with the expensive gradient variance metric?

In Figure 3, we see the TEV (top) and Gradient Variance (bottom) plotted over the first 1 billion tokens. The trends are nearly identical. When gradient variance spikes, TEV spikes. This validates that TEV is a reliable, lightweight proxy for monitoring training stability.

3. Downstream Performance (Perplexity)

Stability is great, but does the model actually learn better? The authors measured Zero-Shot Perplexity on standard benchmarks (Lambada, Wikitext, PTB). Lower perplexity is better.

Table 1 highlights the dominance of MLRA:

• Lower TEV: MLRA consistently has the lowest \(\mu_{\text{TEV}}\), confirming it is the most stable.
• Better Performance: Across almost all datasets and depths, MLRA achieves the lowest perplexity.
• Scaling with Depth: The gap widens as the model gets deeper. At 192 layers, MLRA’s perplexity on Wikitext-103 is 44.17 compared to 47.75 for the baseline (at 96 layers) and outperforms the competitor \(\sigma\)Reparam.

Conclusion

Training Deep Transformers is a balancing act. We want depth for performance, but depth invites chaos in the form of exploding gradients.

This paper provides us with two powerful tools to manage this chaos:

1. TEV: A simple metric derived from embedding weights that acts as a “Check Engine” light for training stability, costing almost nothing to compute.
2. MLRA: A principled architectural change that uses matrix factorization to mathematically dampen variance initialization.

By reducing the variance of the forward pass from \(1/3\) to \(1/9\), MLRA prevents the gradients from exploding, allowing us to train significantly deeper models (up to 192 layers in this study) without the fear of divergence.

For students and practitioners, the takeaway is clear: stability isn’t just about hyperparameter tuning (learning rates, batch sizes). It is fundamentally about how variance propagates through your architecture. Sometimes, a simple factorization is all you need to keep the signal clean.