Extractive Structures Learned in Pretraining Enable Generalization on Finetuned Facts

TL;DR

This paper introduces extractive structures as a framework to understand how pretrained language models generalize facts after finetuning, revealing mechanisms of fact storage and inference that operate across different layers.

Contribution

The paper proposes the concept of extractive structures to explain how components in language models coordinate to enable fact-based generalization, supported by empirical evidence across multiple models.

Findings

Extractive structures are learned during pretraining when facts are encountered before their implications.

Transfer of extractive structures allows counterfactual reasoning about facts.

Fact learning occurs at both early and late layers, enabling different types of generalization.

Abstract

Pretrained language models (LMs) can generalize to implications of facts that they are finetuned on. For example, if finetuned on ``John Doe lives in Tokyo," LMs can correctly answer ``What language do the people in John Doe's city speak?'' with ``Japanese''. However, little is known about the mechanisms that enable this generalization or how they are learned during pretraining. We introduce extractive structures as a framework for describing how components in LMs (e.g., MLPs or attention heads) coordinate to enable this generalization. The structures consist of informative components that store training facts as weight changes, and upstream and downstream extractive components that query and process the stored information to produce the correct implication. We hypothesize that extractive structures are learned during pretraining when encountering implications of previously known facts. This yields two predictions: a data ordering effect where extractive structures can be learned only if facts precede their implications, and a weight grafting effect where extractive structures can be transferred to predict counterfactual implications. We empirically demonstrate these phenomena in the OLMo-7b, Llama 3-8b, Gemma 2-9b, and Qwen 2-7b models. Of independent interest, our results also indicate that fact learning can occur at both early and late layers, which lead to different forms of generalization.