How quickly can a transformer self-heal if you wipe out one of its layers?

Question

Say we have a fully-trained N-layer transformer model (encoder-only, decoder-only, or encoder-decoder), with embedding dimension D, trained for E epochs on S training strings.

Then we take one of those layers, layer L, and reset it back to random weights (all of it, the self-attention weight, the FFN weights, even the layer norm parameters). Assume weights in all other layers are frozen.

What I'm wondering is how quickly it can re-learn. Does it need to see all S x E training examples all over again? Or will it be much quicker, perhaps just needing to see 1% or 10% of the original training data? Or maybe it will go the other way, and refuse to train because the frozen weights around it stop it finding a minima??

I'm also curious if where L is in the stack makes much difference, and if N and D are factors, or the result is about the same however big the transformer is.

(By "re-learn" I mean get back to approximately the same loss/perplexity on validation data that the model had before; I do not mean recreate the exact same weights.)

I'm looking for answers that point to existing papers that have done these experiments, or equally your personal experiments. Or even a "no-one has ever tried this, we don't know", if you think you can say that with confidence.

Aside: I did a test last year of swapping in a different encoder (of different dimension), in an encoder-decoder model, and wiping out the cross-attention weights, but otherwise not touching the rest of the decoder. It acted as if the whole model had been reset to random weights, and seemed to need retraining from scratch, even though 90+% of it was still "trained". But that was only a quick test. Before trying again I'd like to have a better idea of how much retraining time might be needed. (And it struck me that wiping out a layer rather than all cross-attention weights might be a cleaner experiment to learn from.)

Answer 1

If L is close to the end of the model, then it can maybe be thought of as analogous to the conventional practice of replacing the last linear layer or freezing the first L - k layers and only finetuning on the last k. Although the latter doesn't reset the parameters, if your finetuning task is significantly different from your pretraining task, the resulting performance could end up being similar as later layers tend to encode more task-specific information.

I would expect this to use less data as well, as the conventional practice of layer freezing or linear classifier replacement tends to be done in low-resource settings.

I think the keyword you're looking for is layer reinitialization.

RIFLE looks at iterative reinitialization of FC layers during finetuning, to address the issue of FC layers converging too quickly and thus being biased towards pretraining features.

The Impact of Reinitialization on Generalization in Convolutional Neural Networks evaulates several methods for repeatedly reinitializing + retraining subsets of parameters during training. They find that for the same number of iterations, such methods can achieve better test accuracy than training only once.

Both of these are not on transformer models.

This paper uses reinitialization as a method to study the transferability of parameters in pretrained transformers. They reinitialize the parameters for layers that they aren't looking to probe, and study how performance changes during finetuning. They find earlier layers to be most transferrable.

All of these papers are slightly different to what you're looking for. Particularly, many of these settings don't freeze the non-reinitialized parameters. However, I hope this gives you a good starting point into the existing literature.