Linear knowledge before fine-tuning predicts language model generalization

Knowledge That Is Linear Before Fine-Tuning Generalizes After It

Computation and LanguageArtificial IntelligenceMachine Learning

Summary

Language models remember many facts, but some facts help the model learn new tasks better than others. The authors studied how facts that can be read out easily with simple math (linear methods) before training also help models generalize better after training. They showed this by re-labeling relationships, fine-tuning models on part of the data, and testing if the model could guess new answers. The model’s ability to generalize matched closely with how linearly accessible the original facts were, not just whether the model could remember the facts.

What this means in practice

  • For machine learning engineers: Identify which facts in pretrained language models will improve generalization after task-specific training by measuring their linear accessibility beforehand.
  • For natural language processing practitioners: Better select and design fine-tuning tasks that lead to broader generalization by focusing on facts with linearly readable structure in pretrained models.

Authors

Faridoun Mehri, Hinrich Schütze

Abstract

Not all facts a language model can recall support further learning equally well. We distinguish active facts, which support generalization beyond their own recall, from inert facts, which can be recalled but do not support such generalization in a given task. We study this distinction by constructing new relations from existing ones. Take the relation that maps a place to its continent (Paris -> Europe) and assign each of its answers an arbitrary new label (Europe -> crab, Asia -> zebra). The new relation then maps Paris -> crab. We fine-tune the model on the new relation for a subset of subjects, without telling it the original relation or the relabeling rule, and test whether it assigns the corresponding labels to held-out subjects that share an answer in the original relation. Having learned Paris -> crab, does it answer Munich -> crab? Across eleven language models and 70 relations, how linearly readable the original relation is in pretrained representations strongly predicts this generalization. Within each model, the two correlate across relations at Pearson r = 0.87 to 0.92. This association persists when controlling, separately and jointly, for the number of fine-tuning subjects and answers, the model's accuracy on the original relation, and how often the subjects occur in pretraining text. It is also robust to re-randomizing the label assignment and fine-tuning seed, and to replacing the animal labels with labels from thirteen other semantic categories. When we map a fine-tuned model's wrong label back to the original answer it was assigned to, Munich -> zebra becomes Munich -> Asia. Answers recovered this way are more often among the linear probe's top answers for the subject than among the pretrained model's own. Whether a new relation generalizes depends less on whether its facts can be recalled than on how readily their shared structure can be read out linearly.