Mimetic Initialization of Self-Attention Layers
Asher Trockman J. Zico Kolter
It is notoriously difficult to train Transformers on small datasets;
typically, large pre-trained models are instead used as the starting point. We
explore the weights of such pre-trained Transformers (particularly for vision)
to attempt to find reasons for this discrepancy. Surprisingly, we find that
simply initializing the weights of self-attention layers so that they "look"
more like their pre-trained counterparts allows us to train vanilla
Transformers faster and to higher final accuracies, particularly on vision
tasks such as CIFAR-10 and ImageNet classification, where we see gains in
accuracy of over 5% and 4%, respectively. Our initialization scheme is closed
form, learning-free, and very simple: we set the product of the query and key
weights to be approximately the identity, and the product of the value and
projection weights to approximately the negative identity. As this mimics the
patterns we saw in pre-trained Transformers, we call the technique "mimetic
initialization".