Agreed that the huge (in comparison to human short term memory) context size makes updating weights less pressing. Also agreed that training to improve the ability to write memories and documentation by looking at improvements in accuracy on subsequent tasks of a variety of types makes sense. I still think that weight updates have something to recommend them, notably the “grokking” phenomena where further training can reorganize a substantial body of weights.
I have a suggestion on incremental weight updating. I hope it is ok to place it here. If not, can anyone suggest a better location?
One idea for combating catastrophic forgetting that I’d like to add to the mix (alongside other work-in-progress in the field, like the Titans architecture, which addresses a different axis) is to retrain _relevant_ parts of the prior training along with the new continual/incremental training examples, which trickle in sparsely, one at a time.
This is similar to MIR (R. Aljundi et al. NeurIPS 2019 https://arxiv.org/abs/1908.04742) but with a somewhat different problem case and a different choice of trade-off.
This would be a more selective version of rehearsal/replay.
To be more specific, suppose that we have a pre-trained model (I’ll talk about adding some data to pre-training shortly) and we’ve been given a single new example to train on. We may want to crank up the learning rate (since we have one datum, not umpteen) and do the usual gradient descent (more on this later). But we don’t want to forget previously learned skills!
Can we pick out what is most at risk and try to preserve it? Well, the new datum is going to drive some parameter updates. Say we consider the N parameters with the largest changes from the new datum as the riskiest places.
If we could pick out the _old_ training examples most relevant to these N parameters we could retrain with those _old_ training examples as well as with the new datum. We could hopefully converge to parameters that preserve the old skill and yet add the new one, finding weights consistent with _both_.
How might we locate these old training examples? Well, during pre-training there are parameter updates too. For each parameter, there is some pre-training example that produced the maximum update to that parameter. Keep that example during pre-training. More precisely, if we keep the pre-training examples in some known, persistent, order, just keep the _index_ of the pre-training example that produced the maximum update for that parameter and the magnitude of that update. One can extend this by keeping the k-highest update indices.
This is affordable, costing a few bytes per parameter (optionally multiplied by k) for the pointer and the magnitude, a few terabytes on disk for a teraparameter model.
The pre-training compute cost is also small, a numerical compare per parameter per learning step, O(P) cost, the same order as the optimizer’s own state update. These compares only rarely trigger a record write—about k*ln(|corpus|)/|corpus| of the time.
One probably wants to track the update normalized by Adam’s second-moment estimate, typically already maintained per parameter, to ameliorate scale differences across parameters and scale changes during training.
One caveat is that this keeps the magnitude of the update at the time that the model is trained on the example, so this magnitude is stale at the time the new example is examined. However, one wouldn’t want to try to remedy this by recalculating the parameter change due to the old example with the current model, because the current model has already learned from the old example, so a recalculation of the parameter change due to the old example with the current model would understate the parameter learning the old example caused. A well-absorbed sample has a small gradient at the current weights precisely because it was absorbed.
Note that persistent pre-training data is a different problem definition from MIR’s, so this task doesn’t compel reservoir sampling like MIR’s problem definition does. If the incoming data is sparse, e.g. emails from a project, it is reasonable to store them persistently as well, so the equivalent of MIR’s “M” set is the whole corpus—but, as in MIR, budgeting compute is still important.
The main difference from MIR is that MIR’s calculation of the incremental loss imposed on the old training example is exact for the foreseen update, at the cost of an expensive calculation requiring two forward passes through the model per candidate examined. This cost forces MIR to score only a random subset “C” of “M”—which suffices for a small buffer but degrades to near-zero recall of at-risk examples as “M” grows to corpus size.
The per-parameter index replaces ‘score a random subset and hope’ with a direct lookup of the historically implicated examples.
Here, the entire corpus is considered, with O(1) cost to discover each interfered old example, at the cost of approximating the incremental loss by looking just at the most perturbed parameters. In other words, by approximating using the leading term by parameter coordinate of the dot product of the old and new samples’ parameter updates, the hopefully dominant term being the term for the parameter used to retrieve the old sample.
In the spirit of the _signed_ loss calculation in MIR, and selection of their replay set “B” from samples with the largest _increase_ in loss, rather than the largest unsigned change in loss: Since we are concerned with _interference_, we really want to replay old training examples which moved parameters in the _opposite_ direction from the new, sparse training example.
So we should keep _two_ k-highest update indices, one for parameter increases and one for parameter decreases, replaying those with their sign opposite to the new sample’s parameter update. Since examples with the same sign as the new one still carry some risk, e.g. of parameter overshoot, one might replay all k examples from the opposite sign indices and the top one or few from the same sign examples.
A second option to potentially adopt from MIR is using multiple iterations, epoch-ettes, instead of, or in addition to, increased learning rates to emphasize the new sample, together with the retrieved old samples. Two advantages of this are that one can revisit the choice of _which_ old samples to use on each iteration and one follows the nonlinear components of the gradient in parameter space. The disadvantage is that computation cost rises linearly with iteration count, albeit in the current idea the cost of _finding_ each old sample is O(1), while MIR needs 2 forward passes per sample in the “C” set.
To compare with TracIn (Pruthi et al., NeurIPS 2020 https://arxiv.org/abs/2002.08484), TracIn finds training examples which are responsible for large loss changes on a subsequent test case—its best-known application is surfacing mislabeled data. The current idea is intended for cases where fully correct training examples e.g. suggested an overly general deduction. E.g. if, in training, all vehicles happened to be cars and had 4 wheels, and the training led to deducing vehicle->4 wheels, but the new datum is an 18 wheel truck, the deduction can be modified to car->4 wheels, accommodating both the old and new data. TracIn also requires revisiting the whole training corpus, while the current idea retrieves a small set by the parameters a new example perturbs most.
This could be tested on a small model and data set: train on the data set, find or construct an example (like the 18-wheel truck) whose incorporation prompts forgetting, then train on that example and compare forgetting under no replay, random replay, and indexed replay, with equal replay budgets for the latter two.
Thanks to Claude (Anthropic) for extensive discussion — literature pointers, critique of earlier drafts, and several refinements noted in the text.
Agreed that the huge (in comparison to human short term memory) context size makes updating weights less pressing. Also agreed that training to improve the ability to write memories and documentation by looking at improvements in accuracy on subsequent tasks of a variety of types makes sense. I still think that weight updates have something to recommend them, notably the “grokking” phenomena where further training can reorganize a substantial body of weights.
I have a suggestion on incremental weight updating. I hope it is ok to place it here. If not, can anyone suggest a better location?
One idea for combating catastrophic forgetting that I’d like to add to the mix (alongside other work-in-progress in the field, like the Titans architecture, which addresses a different axis) is to retrain _relevant_ parts of the prior training along with the new continual/incremental training examples, which trickle in sparsely, one at a time.
This is similar to MIR (R. Aljundi et al. NeurIPS 2019 https://arxiv.org/abs/1908.04742) but with a somewhat different problem case and a different choice of trade-off.
This would be a more selective version of rehearsal/replay.
To be more specific, suppose that we have a pre-trained model (I’ll talk about adding some data to pre-training shortly) and we’ve been given a single new example to train on. We may want to crank up the learning rate (since we have one datum, not umpteen) and do the usual gradient descent (more on this later). But we don’t want to forget previously learned skills!
Can we pick out what is most at risk and try to preserve it? Well, the new datum is going to drive some parameter updates. Say we consider the N parameters with the largest changes from the new datum as the riskiest places.
If we could pick out the _old_ training examples most relevant to these N parameters we could retrain with those _old_ training examples as well as with the new datum. We could hopefully converge to parameters that preserve the old skill and yet add the new one, finding weights consistent with _both_.
How might we locate these old training examples? Well, during pre-training there are parameter updates too. For each parameter, there is some pre-training example that produced the maximum update to that parameter. Keep that example during pre-training. More precisely, if we keep the pre-training examples in some known, persistent, order, just keep the _index_ of the pre-training example that produced the maximum update for that parameter and the magnitude of that update. One can extend this by keeping the k-highest update indices.
This is affordable, costing a few bytes per parameter (optionally multiplied by k) for the pointer and the magnitude, a few terabytes on disk for a teraparameter model.
The pre-training compute cost is also small, a numerical compare per parameter per learning step, O(P) cost, the same order as the optimizer’s own state update. These compares only rarely trigger a record write—about k*ln(|corpus|)/|corpus| of the time.
One probably wants to track the update normalized by Adam’s second-moment estimate, typically already maintained per parameter, to ameliorate scale differences across parameters and scale changes during training.
One caveat is that this keeps the magnitude of the update at the time that the model is trained on the example, so this magnitude is stale at the time the new example is examined. However, one wouldn’t want to try to remedy this by recalculating the parameter change due to the old example with the current model, because the current model has already learned from the old example, so a recalculation of the parameter change due to the old example with the current model would understate the parameter learning the old example caused. A well-absorbed sample has a small gradient at the current weights precisely because it was absorbed.
Note that persistent pre-training data is a different problem definition from MIR’s, so this task doesn’t compel reservoir sampling like MIR’s problem definition does. If the incoming data is sparse, e.g. emails from a project, it is reasonable to store them persistently as well, so the equivalent of MIR’s “M” set is the whole corpus—but, as in MIR, budgeting compute is still important.
The main difference from MIR is that MIR’s calculation of the incremental loss imposed on the old training example is exact for the foreseen update, at the cost of an expensive calculation requiring two forward passes through the model per candidate examined. This cost forces MIR to score only a random subset “C” of “M”—which suffices for a small buffer but degrades to near-zero recall of at-risk examples as “M” grows to corpus size.
The per-parameter index replaces ‘score a random subset and hope’ with a direct lookup of the historically implicated examples.
Here, the entire corpus is considered, with O(1) cost to discover each interfered old example, at the cost of approximating the incremental loss by looking just at the most perturbed parameters. In other words, by approximating using the leading term by parameter coordinate of the dot product of the old and new samples’ parameter updates, the hopefully dominant term being the term for the parameter used to retrieve the old sample.
In the spirit of the _signed_ loss calculation in MIR, and selection of their replay set “B” from samples with the largest _increase_ in loss, rather than the largest unsigned change in loss: Since we are concerned with _interference_, we really want to replay old training examples which moved parameters in the _opposite_ direction from the new, sparse training example.
So we should keep _two_ k-highest update indices, one for parameter increases and one for parameter decreases, replaying those with their sign opposite to the new sample’s parameter update. Since examples with the same sign as the new one still carry some risk, e.g. of parameter overshoot, one might replay all k examples from the opposite sign indices and the top one or few from the same sign examples.
A second option to potentially adopt from MIR is using multiple iterations, epoch-ettes, instead of, or in addition to, increased learning rates to emphasize the new sample, together with the retrieved old samples. Two advantages of this are that one can revisit the choice of _which_ old samples to use on each iteration and one follows the nonlinear components of the gradient in parameter space. The disadvantage is that computation cost rises linearly with iteration count, albeit in the current idea the cost of _finding_ each old sample is O(1), while MIR needs 2 forward passes per sample in the “C” set.
To compare with TracIn (Pruthi et al., NeurIPS 2020 https://arxiv.org/abs/2002.08484), TracIn finds training examples which are responsible for large loss changes on a subsequent test case—its best-known application is surfacing mislabeled data. The current idea is intended for cases where fully correct training examples e.g. suggested an overly general deduction. E.g. if, in training, all vehicles happened to be cars and had 4 wheels, and the training led to deducing vehicle->4 wheels, but the new datum is an 18 wheel truck, the deduction can be modified to car->4 wheels, accommodating both the old and new data. TracIn also requires revisiting the whole training corpus, while the current idea retrieves a small set by the parameters a new example perturbs most.
This could be tested on a small model and data set: train on the data set, find or construct an example (like the 18-wheel truck) whose incorporation prompts forgetting, then train on that example and compare forgetting under no replay, random replay, and indexed replay, with equal replay budgets for the latter two.
Thanks to Claude (Anthropic) for extensive discussion — literature pointers, critique of earlier drafts, and several refinements noted in the text.
What do you think?