Replay improves continual learning for large language models

Replay for LLM Continual Learning with Multi-Teacher On-Policy Distillation

Machine LearningComputation and Language

Summary

When teaching large language models (LLMs) to learn new knowledge from different areas over time, it's hard for them to remember earlier topics if the original experts are not available. The authors compare ways to help the LLM remember by replaying past examples versus storing snapshots of earlier models. They find that replaying past prompts, even with limited budget, helps the model keep prior knowledge nearly as well as storing all snapshots. Adding alignment of internal model details has minimal extra benefit. The replay methods especially improve performance on safety-related tasks without hurting learning of new topics much.

What this means in practice

  • For ai model developers: Maintain performance of large language models on multiple domains over time without requiring access to all original expert models simultaneously.
  • For enterprise ai teams: Use replay methods to efficiently update domain-specific language models in production while preserving safety-related behavior.$Commercial implications: This enables sale of safer, updatable LLMs for applications requiring regulatory compliance and domain adaptability.

Authors

Yiqiao Qiu

Abstract

Multi-teacher on-policy distillation (MOPD) distills domain experts into one student but needs every expert at once. When domains arrive one at a time and earlier experts are gone, the previous-stage student can teach replayed prompts of earlier domains inside every stage. Some pipelines instead store every stage snapshot for a final multi-teacher recovery stage. On five domains with a 1.5B student and five public experts, this paper compares both with joint MOPD. It varies the replay budget and the replay signal. The signal is token-level distillation alone or with alignment of hidden states, attention and value relations. Sequential distillation without replay or recovery is 6.2 points below 208-step joint MOPD on a 14-benchmark average and 3.5 points below on the 11 benchmarks outside safety. Without snapshots, replay toward the previous student closes 3.8 of the 6.2 points when the earlier domains share the replay budget of a stage. With a fixed share of seven replay steps per earlier domain and stage, it closes 5.8 over two training runs. This is level with the 5.9 of recovery from all snapshots, both at its default length and when given as many training steps as the fixed share. Alignment terms on replay steps change the average by about one point or less. If refusal and compliance count equally in safety, shared-budget replay ranks first among the trained configurations, and the fixed share stays at or above recovery. Against the shared budget, the fixed share lowers the score on the newest domain. Recovery lowers this score as well, in five of six cases. The replay steps of the fixed share grow quadratically with the number of stages.