Reawakening Knowledge:
Anticipatory Recovery from Catastrophic Interference via Structured Training

In this first experiment, we have $T=100$ documents, and we do cyclic fine-tuning of a pre-trained Pythia-1B model (Biderman et al., 2023) on the documents for $E=5$ epochs in the same ordering. Both the documents and the ordering are sampled at random beforehand, but kept fixed during the sequential fine-tuning process. We refer to these $T$ documents as $x_1,\cdots ,x_T$. At the start, we fine-tune on $x_1$ for $M=10$ gradient steps, leading to a significant decrease in the model’s loss on $x_1$. As we move away from $x_1$ and fine-tune on other documents, we naturally observe catastrophic interference: the model’s loss on $x_1$ gradually increases until we finish fine-tuning on all other documents and return to $x_1$. As we iterate through the same document sequence for a second time, we would normally expect the loss on $x_1$ to increase monotonically after the initial decrease. However, Figure 1(a) shows that the loss on $x_1$ peaks around $x_{60}$ and then starts to decrease. Before we return to $x_1$, the model has recovered more than half of its initial forgetting during the second epoch. We refer to this counterintuitive decrease in loss as the anticipatory recovery phenomenon. In Figure 1(b), we plot the losses for all the documents and re-align them so that $0$ on the x-axis refers to the loss on each document $t$ immediately before training on it for the first time. The figure confirms that the anticipatory recovery phenomenon exists for not only $x_1$ but all documents. On the other hand, when we randomly shuffle the document order within each epoch (except $x_1$ is always the first document), we do not observe such anticipatory recovery behavior, and the loss on $x_1$ keeps increasing before we return to it every time.

To quantify the strength of the anticipatory recovery phenomenon, we define the recovery score as the proportion of the initial forgetting during the current epoch that the model recovers before returning to the same document. Mathematically, let the mean (over $t$) of the maximum loss on each document $x_t$ between the $n^{\text{th}}$ and $(n+1{)}^{\text{th}}$ time we train on that document be $l_{\text{max}}\left(n\right)$, right before the $(n+1{)}^{\text{th}}$ time we train on it be $l_{\text{before}}\left(n\right)$, and right after the $(n+1{)}^{\text{th}}$ time we train on it be $l_{\text{after}}\left(n\right)$. Then we define the recovery score (RS) for epoch $n$ to be¹¹1In some cases, a randomly initialized model will produce loss curves that decrease throughout the epoch, because its knowledge is so poor that it enjoys positive generalization among all documents. This yields a misleadingly large recovery score under this definition. We do not include such cases in our experiments so do not bother with more nuanced recovery scores. $RS\left(n\right)=\frac{l_{\text{max}}\left(n\right)-l_{\text{before}}\left(n\right)}{l_{\text{max}}\left(n\right)-l_{\text{\%}after}(n-1)}.$ In the following subsections we compute the recovery scores for different model sizes and training hyperparameters (Figures 2 and 5) to investigate their effects on the anticipatory recovery phenomenon.

To study how the model size affects the amount of anticipatory recovery, we repeat this experiment with pre-trained Pythia models (Biderman et al., 2023) of sizes 160M, 410M, 1.4B, and 2.8B. We plot the average loss curves as well as the recovery score for epoch 4 in Figure 2(a) . We observe that larger models clearly demonstrate stronger anticipatory recovery. The sharp increase of average recovery score from the 160M model to the 410M model indicates that anticipatory recovery is an emergent behavior.

In this section we discuss the effect of other training hyperparameters on the anticipatory recovery phenomenon. We also include additional experiment details in Appendix 7.1 and additional results in Appendix 8.

To examine the generality of the anticipatory recovery phenomenon, in this subsection we explore the sequential cyclic training setting on two tasks in computer vision: causal image modeling and image classification. More detailed experiment setup is available in Appendix 7.

We compare the performance of training with fixed ordering and random shuffling of each epoch in the prequential evaluation setting with $T=100$ documents. Prequential evaluation (Cai et al., 2021) measures online performance by evaluating the model on the document it is about to be trained on, which is equivalent to evaluating the training loss of each batch. For the random shuffling condition, document 1 is always trained first in each epoch, and the other 99 documents are presented in a different random order in each epoch. In this experiment we set $C=512$, $M=10$, $E=5$ and the error bars are based on 10 different seeds.

The resulting loss curves are plotted in Figure 1(c), and the average loss throughout each run is summarized in Table 1. We exclude epoch 1 when computing the average loss since all documents are new to the model for both the cyclic training and random shuffling condition. We observe that training with fixed ordering is superior to random shuffling in the prequential evaluation setting across all 4 pre-trained Pythia models of different sizes, due to the structure in the data stream. The results suggest practical potential of structured training.

We first explore the similarities of gradient information between documents during the training process. Our goal is to test the hypothesis that anticipatory recovery is mediated by increased similarity between gradients of proximal documents in our training sequence.

We do cyclic training for 4 epochs and compute the gradient of each document at the attention layer of transformer block 12 at the conclusion of training. In Figure 8(a), we plot the cosine similarities between these gradient vectors of each document. Results show that the gradients have mostly positive cosine similarities (except for the last document, on which the model has just been trained). To our surprise, the gradient similarities are highest near the center of the heat map rather than peaking along the diagonal. That is, the gradient similarity between documents $x_{t-1}$ and $x_t$ depends on where we are in the cycle. This result suggests an additional layer to the anticipatory recovery phenomenon: Recovery for document $x_t$ is greatest from training on document $x_{t-1}$, but the strength of the potential facilitation between $x_{t-1}$ and $x_t$ is actually greatest after we train for another $b$ documents (for some small number $b$). We verify this by computing the pairwise recovery: we take the model checkpoint after 4 epochs of cyclic training, and then for each pair of documents $(x_i,x_j)$, we do $M$ gradient updates on $x_i$ and compute the difference in the loss of $x_j$ before and after these gradient updates. We plot these pairwise loss recoveries in Figure 8(b). Results confirm that the amount of recovery on document $x_j$ is highest when the model checkpoint is taken from roughly $b$ documents before or after document $x_j$ in cyclic training and then fine-tuned on a proximal document $x_i$ in the sequence. The fact that this pairwise loss recovery matrix is roughly symmetric also suggests that the anticipatory recovery phenomenon approximately exhibits task symmetry: gradient updates on document $x_t$ decrease the loss for document $x_{t+k}$ for small integers $k$, and vice versa. We provide additional visualizations for $T=50,100,200$ in Section 9 and a more detailed description for this phenomenon.

We explore the structure of model weights along the optimization trajectory of cyclic training. We flatten and concatenate the model weight vectors after fine-tuning on each document. However, the cosine similarities between these raw model weight vectors are all very close to 1 without obvious structure, due to the proximity of model weights along the same optimization trajectory and numerical instability in huge weight vectors. To resolve these issues, we instead explore the structure of “weight residuals.” We compute the weight residuals by subtracting the average of weights in a window of length $T$ centered at the current document from the current weight, i.e. $w_{\text{res}}\left(t\right)=w\left(t\right)-\frac{1}{T}{\sum }_{n=t-T/2}^{t+T/2}w\left(n\right)$. This removes the shared components along the optimization trajectory and allows us to focus on the model weight updates for each document. Figure 8(c) visualizes a heat map of the cosine similarity between each pair of weight residuals from the second epoch to the fourth epoch. The visualization shows a cyclic structure in the weight residuals, as equidistant bright stripes that align with the training epochs. Furthermore, each stripe spans several documents, suggesting the similarity of weight residuals for proximal documents.

In addition to cosine similarities between weights, we explore visualizing the weights in a lower-dimensional space with Principle Component Analysis (PCA). We compute the top three PCs of the flattened last-layer weight vector (the output word embedding layer) for the Pythia-1B model, and plot its trajectory in Figure 9. The plot exhibits a clear helical structure that gradually converges. We believe this is highly relevant to anticipatory recovery: right before revisiting a task, the projected model weights in the helix move closer to the point corresponding to the previous appearance of that task, leading to anticipatory recovery on the loss of that task. As we go on with cyclic training, the model also exhibits less forgetting and gradually converges to a solution that achieves low loss on all tasks. It is important to note that the helical structure of the weight trajectory is not an obviously necessary consequence of the cyclical training. Cyclical training could be expected to yield a repeating pattern, but the facts that the tasks come to be organized in a circle that respects their ordering and that the trajectory goes through one full revolution per epoch (rather than some other arc length) are nontrivial and seem to be essential for anticipatory recovery.

In addition to gradients and weights, we visualize the trajectory of activations on a single document during the course of cyclic training. We do cyclic training for three epochs and save model checkpoints after fine-tuning on each document. We then compute the model activations before the output word embedding layer for document $x_1$ on each model checkpoint, and plot the cosine similarities between the flattened activation vectors in Figure 8(d). From the plot we can clearly observe the blocked pattern wherein the similarity between the activations become progressively higher across each epoch of cyclic training. This pattern suggests that every time we train on document $x_i$, the internal representation of $x_i$ in the model is more resistant to gradient updates on other documents $x_j$.

To further understand the essential mechanism that yields anticipatory recovery, we design a minimalist “toy” simulation experiment. In this toy simulation, each task (formerly, document) $i\in \{1,\cdots ,T\}$ is described by a single data point, $x_1,\cdots ,x_T\in R^N$. We assume a learnable linear embedding $P\in R^{M\times N}$ that projects each $x_i$ into an $M$-dimensional embedding space. We also assume a learnable vector $w$ and task-specific mappings $f_i$, where $f_i\left(w\right)$ is the target for task $i$ in the same embedding space. We require each $f_i$ to be invertible as a simplifying assumption.

We define the loss for task $i$ as ${\ell }_i(P,w)=\frac{1}{2}\parallel P{x}_i-f_i\left(w\right){\parallel }_2^2$. Just as when training a deep net, we assume here that representation learning occurs slowly, and that one training step for task $i$ involves a single gradient update of $P$ with step size $\alpha$:

In contrast, at each training step, $w$, analogous to the fast-adapting weights in a neural network, can be rapidly tuned to solve for task $i$, yielding the loss minimizer conditional on $P$:

As in our main experiments, we sequentially optimize each ${\ell }_i$ as we iterate through the sequence of tasks. In each training step, we first update $P$ and then solve for $w$ given the update to $P$. Updating $P$ approximately reduces the distance between $P{x}_{i+1}$ and $f_{i+1}\left(f_i^{-1}\right(P{x}_i\left)\right)$, which entails reducing the upper bound of $||f_{i+1}^{-1}\left(P{x}_{i+1}\right)-f_i^{-1}\left(P{x}_i\right)||$, assuming that each $f_i^{-1}$ is Lipschitz continuous. As a result of this optimization objective, the model will evolve along the optimization trajectory such that the $f_i^{-1}\left(P{x}_i\right)$ for all tasks $i$ gradually form a circular pattern. This gives an intuitive explanation on the anticipatory recovery phenomenon, since updaing $w$ according to Equation 2 will also bring it closer to $f_{i+1}^{-1}\left(P{x}_{i+1}\right)$, thus reducing the loss on task $i+1$ and exhibits anticipatory recovery.

We experimented with two very simple choices of $f_i$: $f_i\left(w\right)=w$ and $f_i\left(w\right)=y_i-w$ for some task-dependent targets $y_i$. We follow the same order over tasks—$1,\cdots ,T$—for multiple epochs of training. The resulting loss curves are shown in Figure 10, which exhibits very similar anticipatory recovery trajectory as the full-blown LLM experiment. Visualizations of the 2-dimensional PCA embeddings for $f_i^{-1}\left(P{x}_i\right)$ in the second experiment are shown in Figure 11, which confirms our analysis that they gradually self-organize into a cyclic structure.

There are two potential reasons large overparameterized networks might produce the anticipatory recovery in a way analogous to the toy simulation. First, for larger networks, it is more likely that the network can develop task-specific parameters that quickly adapt to and memorize new input data, corresponding to Equation 2. And when the fast memorization is achieved, the gradient descent dynamics of the slow weights push the representations of the two adjacent tasks ($P{x}_i$ and $P{x}_{i+1}$) closer when $f$ is an identity function, according to Equation 1. This effect can be seen in earlier LLM experiments (Figure 2), where larger models achieve significantly lower losses within a few gradient update steps. Second, larger networks have more learning capacity to map the features of two adjacent tasks closer. In our linear projection model, anticipatory recovery keeps growing over many epochs, whereas the anticipatory effect is already at the strongest within 2 or 3 epochs in LLM experiments. Moreover, all data points are randomly generated in the toy model, which makes it easier to separate and map their representations according to a temporal structure than real-world data. In contrast, real-world data could require more representation capacity since data points are noisy and correlated.

In this section, we visualized model weight dynamics with heatmaps and we showed model activations and gradients during cyclic training. We discussed the special temporal structure that is exhibited in these heat maps. We also plotted the pairwise degree of recovery for fine-tuning on document $i$ and evaluating on document $j$, as well as the change of distance between fine-tuned model weights on different tasks. The results suggest that after we train on a document, the model’s representation of that document becomes less sensitive to gradient updates on other documents. Finally, we showed a simple toy experiment that demonstrates a similar anticipatory recovery phenomenon in its loss curve, and discuss its connections to neural network training dynamics through the lens of task-specific and task-general parameters. Overall, these results shed some light on the dynamics of cyclic training.

We use the Huggingface Transformers Library (Wolf et al., 2020) for fine-tuning the LLMs. The learning rate $0.001$ for vanilla gradient descent and $0.00001$ for Adam. For all experiments we run 3 to 5 trials with different random seeds, except the results in Figure 1 which are based on 20 seeds. The shaded area in the figures denotes standard deviation among trials and documents (for shift-averaged loss curves).

The images are resized to 256x256 followed by a center crop of 224x224. We use the Adam optimizer with learning rate $0.0001$ for $M=10$ gradient steps on each batch of images.

The shift-averaged loss curves plotted in Figure 1(b) are calculated by replicating Figure 1(a) on each document in the training sequence, re-aligning these curves so that 0 on the x-axis always represents the moment before the first occurrence of the focal document, and average them. For example, if the length of the sequence is 50, then for training epoch 0.5 on the x-axis, we take the loss of document 1 after training on document 25; the loss of document 2 after training on document 26; …; the loss of document 50 after training on document 24 of the next epoch; and average these losses. Subsequent figures (Figure 2-7, 12(a)-16) are plotted with the same approach.

For Figure 10(a), we pick $f_i\left(w\right)=w$, and each data point $x_i$ and $w$ is a vector of length $N=M=1000$. We have $T=25$ data points and use the vanilla gradient descent optimizer with learning rate 0.01. The projection matrix is initialized with every entry sampled independently from $N(0,1/N^2)$. Each entry of the data points $x_i$ and $w$ is sampled independently from $Unif(-1,1)$. For Figure 10(b), we pick $f_i\left(w\right)=y_i-w$, $N=M=100$, $T=25$, and learning rate 0.01. Each entry of $y_i$ is also sampled independently from $Unif(-1,1)$.

Each experiment presented in the paper is run with one NVIDIA A100 GPU, 2 CPUs, and 32GB of RAM. The training time highly depends on the hyperparameter choices, especially model size and number of gradient steps. The longest fine-tuning experiment with 20 gradient steps per episode on the Pythia-1B model takes roughly 30 minutes under this setup. The minimal compute resource requirement needed to reproduce the experiments with a Pythia-1B model is one GPU with 16GB of GPU memory.

Throughout the paper we have been focusing on the setting where the document ordering is sampled once and stay fixed for all epochs. What if we only fix the first and the last document, and shuffle the documents in between? We experimented with shuffling the documents from document $x_2$ through $x_N$ for $N\in \{8,16,24\}$ every epoch. In Figure 12(a) we plot the loss curves for document $x_1$. From the loss curves we can observe that even when $N=24$ we can still observe some anticipatory recovery, suggesting that the order of the tasks between two consecutive repetitions of the $x_{25}$ and $x_1$ can be arbitrary for us to observe recovery on $x_1$.

In Figure 5(b) we observe that there is no anticipation when we take only one gradient descent step on each document with learning rate 0.001. We experimented with one-step gradient descent using a higher learning rate, 0.01. We plot the resulting average loss curves under the same training setup in Figure 12(b). We observe that, with a larger learning rate, slight anticipation is still observed for $1$ gradient step.

In addition to varying the model width and model depth in Figure 3, we also experimented with varying the number of attention heads $h\in \{2,4,8,16\}$ while keeping model width to be 2048 and model depth to be 16. Loss curves on document $x_1$ are shown in Figure 12(c). The results suggest that the number of attention heads does not have a big effect on cyclic training in our setting.

Here we compare the performance of the initialization scheme used by Biderman et al. (2023) (also used for all randomly initialized models in the main text) and a simple initialization scheme that samples the weight matrices from an isotropic Gaussian distribution with $\sigma =0.02$. The loss curves for document 1 under these two initializations of the Pythia-1B model are plotted in Figure 12(d). We observe that Pythia’s initialization scheme achieves much better average loss and also exhibits stronger anticipatory recovery. This demonstrates the importance of LLM initializations. The result is consistent with our observations in section 3.3 that the model’s ability to fit on each task is correlated with the amount of anticipatory recovery.

In addition to comparing pre-trained models and randomly initialized models in Figure 2, we further study the effect of model pre-training by examining model checkpoints with different numbers of pre-training steps. We took Pythia models pre-trained for 6K, 12K, 24K, 48K, and 96K steps and plot the shift-averaged loss curves for cyclic fine-tuning in Figure 13. We found that more pre-training does give rise to higher anticipatory recovery. As we summarize at the end of section 3.3, we hypothesize this result fits a broader pattern in which the strength of the anticipatory recovery effect is related to how well the model can fit each successive training task. Models with more pre-training steps are more capable of fitting each successive training task, and therefore exhibit higher anticipatory recovery.

For experiments in the main paper we used a constant learning rate during the fine-tuning process. To study whether anticipatory recovery occurs in typical LLM optimization schemes, we experimented with cosine learning rate scheduling on the Pythia-1B model with 10 epochs (minimum learning rate = 0, maximum number of epochs = 10), and plot the results in Figure 14. We show that the model also exhibits the anticipatory recovery effect in other learning rate schedules.

To examine how the anticipatory recovery effect may generalize to other forms of structured training, we experimented with a new setting where only the first 20 documents are kept fixed in each epoch, and a random number of other documents (between 20 and 100) are inserted after the first 20. These random padding documents appear only once in the entire sequence. This new setting generalizes cyclic training in that (1) rather than having the same documents in every epoch, we insert random other documents between every repetition (2) epochs can have different lengths. The resulting loss curve is plotted in Figure 15. We still observe anticipatory recovery for documents 2 through 20 in this setting, suggesting that anticipatory recovery exists as long as there is a repeating sub-sequence in the data stream.

In Figure 16(a) we experimented with inversely scaling the learning rate with the number of gradient steps. We use a learning rate of $0.01$ for $M=1$, learning rate $0.002$ for $M=5$, learning rate $0.001$ for $M=10$, and learning rate $0.0005$ for $M=20$. The results suggest that the anticipation effect is stronger when the same total update is divided among more gradient steps.

In Figure 16(b) we experimented with different number of gradient steps $M\in \{10,20,40\}$ for context length 1024. The results confirm that longer context length is not a fundamental limitation to anticipatory recovery, and we can achieve the same recovery score as a smaller context length with more gradient steps.

To evaluate how the anticipatory recovery phenomenon generalizes across different LLM architectures, we experimented with GPT-2 architecture (Radford et al., 2019), specifically the GPT2-large pre-trained model (812M parameters) on the CNN/Daily Mail dataset. The loss curve for document 1 is plotted in Figure 17. The model consistently observed anticipatory recovery. Note that GPT-2 is the predecessor of many modern LLMs and therefore the results further suggest that the anticipatory recovery phenomenon is prevalent among more recent LLM architectures.

To evaluate how the anticipatory recovery phenomenon generalizes across different natural language datasets, we experimented with the wikitext-103 dataset (Merity et al., 2017), which contains over 100 million tokens from articles on Wikipedia. Since Wikipedia data is part of the pre-training dataset of Pythia, we only experiment with randomly initialized models. The loss curve for document 1 is plotted in Figure 18. The result suggest that the anticipatory recovery phenomenon is generalizable to different data sources.

We plot the magnitude of the difference between the $x_1$ activations of model checkpoints saved at consecutive episodes throughout four epochs of cyclic training of a Pythia-410M model in Figure 19(a), and observe a clear stepwise pattern. In contrast, the magnitude of model weight updates (Figure 19(b)) decreases monotonically over the training episodes and do not exhibit this stepwise pattern. This result is consistent with the pattern we observe in section 4.3.

Similar to Figure 8(b), we plot the pairwise loss recoveries for each pair of documents $(x_i,x_j)$ in longer document sequences, where $T=50,100,200$ respectively, in Figure 20. We use the 1B model and default hyperparameters. We observe that, as we increase the length of the document sequence, the highlight area near the center of the matrix is separated into two blobs, one in the top-left corner and the other in the bottom-right corner. We also observe a "boundary" on the sides of the matrix where there is little or no recovery. The width of this "boundary" stays relatively constant across different lengths of document sequences and is around 10 to 15 documents. This confirms our observation in the main text that the recovery on document $x_j$ when fine-tuning on a proximal document $x_i$ is highest when the model checkpoint is taken from document $x_{j\pm b}$ where $b$ is a small number relative to the length of the document sequence.