Output embeddings cause catastrophic forgetting in large language models
A Deafening Silence: Catastrophic Forgetting Lives in the Output Embeddings of Tokens the Data Never Speaks
Computation and LanguageArtificial Intelligence
Summary
Large language models often forget what they learned when trained again on new data, especially for words not seen in the new training. The authors found this forgetting happens mostly in the model’s output part related to rare words. They discovered this is due to how the optimizer (Adam) updates the model during training. By making a simple change to the optimizer settings for that output step, they greatly reduce forgetting without hurting learning on new data. This fix works across many model sizes and types.
What this means in practice
- •For machine learning engineers: Reduce forgetting in large language models during continual training without extra data by adjusting optimizer settings for output tokens.
- •For ai model deployment teams: Improve stability and retention of model knowledge when updating deployed language models with new, limited datasets.
Authors
Jonghyun Han, Younghoon Song, Jongyoul Park
Abstract
Continual pre-training and fine-tuning in Large Language Models (LLMs) inevitably induce catastrophic forgetting, typically mitigated by replay using often-inaccessible original data. In this data-free regime, we analyze where forgetting occurs and why. Systematic parameter freezing across five settings up to 1.4B reveals that forgetting concentrates selectively in the output embeddings of tokens rarely seen in the new corpus, whereas the same sqrt(v-hat) band of the body is inert and new learning resides elsewhere. This localization is governed by the vocabulary deficiency of the corpus rather than the training mode, allowing pre-retraining risk ranking from token counts alone within a fixed base model. Mechanistically, absent tokens receive persistent one-sided softmax gradients that Adam's second-moment (sqrt(v-hat)) normalization amplifies into full-sized updates. We therefore propose an intervention: raising Adam's epsilon exclusively for the output projection during training. Across eight settings spanning 160M to 12B parameters and four model families, this removes 39.4% to 67.9% of forgetting across all seven stable configurations without degrading target learning or requiring per-model tuning. The defense combines additively or better with replay (79.8% on Qwen/Korean) and rescues released-head LoRA from a 23-fold forgetting surge. Because post-hoc editing of the drifted rows recovers under 5% of forgetting, the intervention must operate during training. Our findings indicate that a single-line optimizer adjustment may serve as the primary defense against catastrophic forgetting where the corpus starves the vocabulary.