First-order models help predict forgetting during language model fine-tuning
When Can First-Order Models of Fine-Tuning Bound Forgetting?
Artificial Intelligence
Summary
When language models are updated with new information, they can forget facts they previously knew. This paper studies whether an initial test can predict which facts might be forgotten during fine-tuning. The authors show that simple first-order models can reliably estimate whether facts will remain stable or be forgotten if certain mathematical conditions about parameter changes hold. These results hold especially when the model’s response changes less than a certain threshold during training. This helps understand and potentially control forgetting in large language models as they learn new data.
What this means in practice
- •For machine learning engineers: Assess and limit forgetting risk for protected facts during language model fine-tuning by monitoring response coefficient changes.
- •For ai safety teams: Provide guarantees on forgetting probabilities for critical knowledge in language models during ongoing updates under certain response stability conditions.
Authors
Jianchang Su, Wei Zhang
Abstract
Fine-tuning a language model on new data can make it forget facts that it should keep. We ask whether measurements taken at the start of a fine-tuning run can bound, for each protected fact, the probability that the run makes the model forget it. In LoRA fine-tuning with stochastic gradient descent on models from 0.6B to 14B parameters, a first-order response model estimated by finite-difference probes predicts changes of per-fact margins with correlation 0.974-0.998. Predictions of forgetting built on this model nevertheless failed, because forgetting requires parameter changes far outside the region in which the model was validated. The probes can, however, bound the probability that a margin first falls below a boundary near zero: we derive Freedman and Azuma first-passage bounds for a linear surrogate of the margin and test on new runs whether they hold for the model. The bounds contain a term R that measures how much the response coefficients change during the run. The simplified Freedman bound, which sets R = 0, certified most facts but was violated in 14 of 112 conditions, and every fact on which it was violated had R >= a, where a is the distance of the fact's margin to the boundary. The complete Freedman bound certifies only facts with R < a, and it held in every condition. On the violated facts, the spread of the margin across test runs was a median of 14.6 times the prediction of the response model, so the failures are breakdowns of the model, and in our data they occurred only where R >= a. We found this pattern post hoc and tested it in two preregistered confirmatory studies with 43 new conditions: the complete bound held in all of them, and the simplified bound failed there on only 3 facts, each with R >= a. First-order models of fine-tuning can thus bound forgetting on the facts whose response coefficients change by less than their distance to the boundary.