Embedding Anomalies under Lion/Tiger Optimizer Training and Countermeasures

Lion/Tiger优化器训练下的Embedding异常和对策 苏剑林 · 2023-08-28 · Machine Learning problem gradient optimizer

Ever since I proposed the Tiger optimizer in Tiger: An Optimizer Taken to the Extreme in "Stinginess", Tiger has become my "go-to" optimizer for training models. Recently I've tried applying Tiger to the pretraining of a 7-billion-parameter model, and the early results look decent, tentatively confirming that Tiger can indeed scale up. However, when inspecting the trained model weights, I found that the Embedding layer had developed some abnormal values — some Embedding components reached the level of $\pm 100$.

After analysis, I found that this phenomenon does not occur with Adam; it's a problem specific to optimizers like Tiger or Lion that use a sign function $\text{sign}$. At the end of this post I offer two reference solutions. This post documents my analysis process for readers' reference.

The Phenomenon

In what follows, our analysis will use the Tiger optimizer as the example, but the analysis and conclusions apply equally to Lion. More...

First, here's what I observed:

1. The Embedding components of some tokens became $\pm 100$;
2. A small additional set of tokens have Embedding components tending toward $\pm 100$;
3. These tokens all appear to be fairly low-frequency tokens;
4. The overall maximum of the Embedding matrix is exactly 100, and the minimum is exactly -100;
5. Apart from the Embedding layer, no other weights exhibit this problem;
6. The model's overall behavior (e.g., training loss, generation tests) is normal.

Some readers might ask: if the model's performance is normal, why worry about it? In my view, there are at least two reasons. First, if one later wants to fine-tune the model, some low-frequency tokens might become high-frequency again, and if their Embeddings are ruined, fine-tuning may not be able to rescue them. Second, some capabilities aren't reflected in the loss at all — for example, in a Chinese-English pretrained model, a small amount of multilingual data mixed into the training corpus often gives the model a certain degree of multilingual ability. This ability clearly depends on the quality of the Embeddings for low-frequency tokens, and if it gets collateral damage from the optimizer, that would be a real loss.

Of course, no matter which optimizer you use, it's always possible for a model to collapse partway through training — that's not surprising in itself, and often hard to fully diagnose. But what's especially intriguing here is how regularly it "collapses" — landing exactly at the neat value $\pm 100$ — which compelled me to dig further into the underlying cause.

Reasoning

Based on the observations above, it seems that these anomalous values only appear in the "Embeddings of low-frequency tokens," which immediately reminded me of the issue discussed in Keras Implementation of Two Optimizers: Lookahead and LazyOptimizer, namely that optimizers with momentum can cause over-optimization of the Embedding layer.

Specifically, once a token has appeared even once, the momentum corresponding to that token's Embedding gets updated to a nonzero value (assuming its gradient isn't exactly zero). Then, in subsequent updates, even if the current batch doesn't contain that token (so its gradient is zero), the Embedding will still be updated because the momentum is nonzero. This is the over-optimization problem for low-frequency tokens. This issue can arise in any optimizer with momentum, including Adam and Tiger. However, with Adam this effect may not be noticeable, because Adam's update magnitude is proportional to the momentum itself: if a token doesn't reappear for a long time, its momentum decays exponentially, quickly approaching zero — in other words, the update magnitude also quickly approaches zero, and the over-updating effect soon vanishes.

With Tiger, though, things are a bit different. Tiger's update is proportional to the sign function $\text{sign}(\boldsymbol{m}_t)$ of the momentum. Even though the momentum $\boldsymbol{m}_t$ decays exponentially, the sign function does not — before $\boldsymbol{m}_t$ becomes exactly 0 due to rounding error, $\text{sign}(\boldsymbol{m}_t)$ stays fixed at $\pm 1$, meaning the update magnitude remains constant the whole time. This makes Tiger's Embedding over-updating problem much more severe. To make matters worse, once a token's Embedding has been pushed in some direction due to over-updating, its gradient may adapt and reinforce that trend — that is, the next time the token appears, its gradient tends to point in the same direction rather than the opposite one. This causes the Embedding to keep over-updating in the same direction for a long time, eventually producing the anomalous values.

The Calculation

So why exactly do the anomalous values end up at $\pm 100$? This is where weight decay comes into play. Tiger's overall update formula is:

\begin{equation}\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta_t \left[\text{sign}(\boldsymbol{m}_t) + \lambda \boldsymbol{\theta}_{t-1}\right]\end{equation}

That is, besides the sign function of the momentum, there's also a weight decay term. In the anomalous experiment mentioned at the start of this post, the decay rate $\lambda$ was set to 0.01.

It's not hard to see that if $\text{sign}(\boldsymbol{m}_t)$ remains constant over a long period, the above iteration formula has a fixed point, occurring when $\text{sign}(\boldsymbol{m}_t) + \lambda \boldsymbol{\theta}^*=\boldsymbol{0}$, i.e.,

\begin{equation}\boldsymbol{\theta}^* = -\frac{\text{sign}(\boldsymbol{m}_t)}{\lambda}\end{equation}

This corresponds exactly to a vector whose elements are $\pm 100$, which explains why the anomalous value turns out to be $\pm 100$. If you're interested, you can also assume that $\eta_t$ is constant, in which case you can directly derive a closed-form expression for $\boldsymbol{\theta}_t$ and further analyze things like convergence speed. I won't pursue that further here.

Countermeasures

Since the problem arises from the over-updating of low-frequency token Embeddings, one natural solution is, as suggested in Keras Implementation of Two Optimizers: Lookahead and LazyOptimizer, to make the Embedding update "lazy" — i.e., only update an Embedding when its corresponding token has actually appeared. If you can obtain the full set of input token IDs, you can directly update only those tokens' Embeddings; if not, you can check whether the gradient norm of an Embedding is nonzero to determine whether it needs updating.

From a more general perspective, this problem is a shared weakness of the Lion/Tiger optimizers when dealing with parameters that have sparse gradients — not limited to the Embedding layer. So another approach to solving the problem is to make the Embedding gradients no longer sparse. To this end, we can consider Tied Embeddings, i.e., sharing the input and output Embeddings. Since the output side reuses the entire Embedding matrix, the whole Embedding matrix ends up with nonzero gradients, preventing $\boldsymbol{m}_t$ from staying constant for long stretches. Of course, Tied Embeddings may introduce their own set of issues; for corresponding solutions, see Revisiting Shared Output Embeddings in Language Models. In my experiments, using a Tied Embedding scheme where the model's feature channels are split in half and swapped resolved the above problem, and the results seemed to be even slightly better than with Untied Embeddings.

Finally, I also consulted the authors of the Lion optimizer about this issue, and their reply was that they had noticed this problem before too. Their solution is to mix optimizers — for instance, using Adam for the Embedding layer and Lion/Tiger for the other layers. Well, that's a solution I hadn't thought of; it doesn't feel especially elegant, but it does work, so readers can choose for themselves.

Summary

This post introduced the Embedding anomaly phenomenon observed when training with the Lion/Tiger optimizers, analyzed the underlying cause, and finally offered a couple of reference solutions.