PaperlayerView exact revision on arXiv

Source-linked AI summary

Physics of Language Models: Part 3.1, Knowledge Storage and Extraction

Zeyuan Allen-Zhu, Yuanzhi Li

arXiv:2309.14316v3cs.CLcs.AIcs.LG

TL;DR

The paper investigates whether QA answers reflect genuine extraction from training biographies rather than memorization of similar questions. Using a controlled biography dataset and probing methods, it finds that pretraining augmentation is closely tied to both extraction performance and entity-centered knowledge encoding, motivating earlier QA-like data and data rewriting. The study also notes abnormal mixed-training behavior and a fine-tuning limitation in the tested setting.

  • Problem

    It is unclear whether language models answer factual questions by extracting knowledge from training sources or by recognizing exact or similar questions seen during training.

  • Method

    The authors use a controlled biography dataset, vary pretraining and fine-tuning data, and apply nearly linear probing to examine knowledge representations.

  • Results

    Knowledge augmentation during pretraining improves extraction after fine-tuning and makes knowledge nearly linearly encoded in person-name embeddings, while unaugmented knowledge is distributed across biography tokens.

  • Takeaways & Limitations

    Practitioners should rewrite critical, infrequent pretraining data and introduce more instruction-finetuned data during pretraining rather than postponing it entirely to fine-tuning.

  • Takeaways & Limitations

    Mixed training achieves extraction through somewhat abnormal behavior akin to studying to pass the test, and the reported bidirectional-model result is limited to the tested setting.

Abstract

from arXiv · show

Large language models (LLMs) can store a vast amount of world knowledge, often extractable via question-answering (e.g., "What is Abraham Lincoln's birthday?"). However, do they answer such questions based on exposure to similar questions during training (i.e., cheating), or by genuinely learning to extract knowledge from sources like Wikipedia? In this paper, we investigate this issue using a controlled biography dataset. We find a strong correlation between the model's ability to extract knowledge and various diversity measures of the training data. $\textbf{Essentially}$, for knowledge to be reliably extracted, it must be sufficiently augmented (e.g., through paraphrasing, sentence shuffling, translations) $\textit{during pretraining}$. Without such augmentation, knowledge may be memorized but not extractable, leading to 0% accuracy, regardless of subsequent instruction fine-tuning. To understand why this occurs, we employ (nearly) linear probing to demonstrate a strong connection between the observed correlation and how the model internally encodes knowledge -- whether it is linearly encoded in the hidden embeddings of entity names or distributed across other token embeddings in the training text. This paper provides $\textbf{several key recommendations for LLM pretraining in the industry}$: (1) rewrite the pretraining data -- using small, auxiliary models -- to provide knowledge augmentation, and (2) incorporate more instruction-finetuning data into the pretraining stage before it becomes too late.

1 Introduction

The paper separates memorizing biography text from extracting its factual knowledge through unseen questions. Controlled experiments show that pretraining diversity and augmentation strongly affect extraction, while probing links this effect to how knowledge is encoded.

  • Research question: The study asks whether language models extract facts from training sources or recognize similar questions encountered during training.It uses out-of-distribution questions about individuals whose QAs were not included in training.
  • Mixed training: Mixed training enables models to answer questions about individuals whose QAs were withheld, by combining biographies with QAs for only a fraction of people.The authors describe this as a pretraining process in which QA data helps encode knowledge that generalizes to remaining individuals.
  • Knowledge augmentation: Instruction fine-tuning alone fails to produce reliable knowledge extraction from biography-only pretraining, whereas writing-style variation and sentence shuffling substantially improve accuracy.The failure persists across model size, pretraining duration, and fine-tuning parameters, and is not resolved by minimal LoRA updates.
  • Internal encoding: Knowledge augmentation makes facts nearly linearly encoded in entity-name embeddings; without it, knowledge is distributed across biography tokens and is difficult to extract.The probing results connect training-data diversity with the internal representation used for question answering.
  • Practical implications: The practical recommendations are to rewrite rare but important pretraining data and introduce QA-like instruction data before fine-tuning.Suggested rewriting methods include paraphrasing, sentence shuffling, and translation, potentially using small auxiliary models.

2 Result 0: Our Dataset Families

The paper constructs two controlled biography families, synthetic and close-to-real, with six attributes per person and multiple knowledge-augmentation settings. It evaluates mixed training and pretrain-then-instruction-finetune using autoregressive GPT2/Llama models.

  • Dataset families: The synthetic bioS dataset contains 100,000 individuals with randomly and independently selected attributes and unique full names.Each profile includes six biographical attributes, including a company city dependent on employer headquarters.
  • Dataset families: Each bioS profile is a six-sentence biography whose sentences are selected from approximately 50 templates.The basic bioS single configuration uses one entry per person with a consistent sentence order.
  • Dataset families: The bioS augmentation settings are multiM, fullname, and permute, varying biography entries, replacing pronouns with full names, or shuffling sentences.These settings are introduced to diversify how the same person’s knowledge appears in pretraining data.
  • Dataset families: The bioR dataset contains 100,000 Llama-generated biographies produced from diversified prompts and rejection sampling that ensures all six attributes appear.Its basic configuration likewise uses a single biography per person, with multiM and fullname comparisons.
  • Evaluation and training: Knowledge extraction is evaluated with six question-answer prompts per individual, using exact-match accuracy for generated answers.The questions target birth date, birth city, university, major, company, and company city.
  • Evaluation and training: The study compares pretrain-then-instruction-finetune, mixed training, and LoRA or full finetuning across GPT2 and Llama architectures.Pretraining uses BIO entries, while instruction finetuning uses half the QA data and evaluates on the remainder.

3 Result 1: Mixed Training Enables Knowledge Extraction

Mixed training enables out-of-distribution knowledge extraction by jointly training on biographies and QA tasks, but the model follows an unusual strategy resembling studying from questions before aligning knowledge with biographies.

  • Result 1: Mixed training combines BIO data for all individuals with QA data for a subset and evaluates generation on held-out individuals.The experiments separately monitor BIO and QA first-token accuracy, QA generation accuracy, and train/test partitions.
  • Result 1: Mixed-trained models achieve 86.6% OOD accuracy on bioS single and 77.7% on bioR single.OOD means extracting attributes for individuals whose biographies, but not QA examples, were seen during training.
  • Model’s abnormal learning behavior: The mixed-training result resembles studying to pass a test rather than the natural progression of human knowledge acquisition.The paper therefore examines BIO-only pretraining followed by QA finetuning as a more challenging setting.
  • Model’s abnormal learning behavior: The model initially encodes knowledge from in-distribution QA, then improves in-distribution BIO memorization before aligning that knowledge with BIO to generalize out of distribution.The delayed rise in BIO out-of-distribution accuracy accompanies a later increase in QA out-of-distribution accuracy.
  • Model’s abnormal learning behavior: A higher QA training ratio improves out-of-distribution QA accuracy, supporting the observed sequence in which QA knowledge is associated with BIO.The mixed-training configuration uses an 8:2 QA-to-BIO entry ratio.

4 Result 2-3: BIO Pretrain + QA Instruction Finetune

BIO pretraining followed by QA fine-tuning can yield near-perfect memorization yet fail to extract knowledge for out-of-distribution questions. Diverse pretraining augmentations—especially multiplicity, permutations, and repeated full names—substantially improve extraction.

  • 4 Result 2-3: BIO Pretrain + QA Instruction Finetune: The BIO pretraining and QA fine-tuning setup trains on biographies for all individuals, fine-tunes on Ptrain, and evaluates out-of-distribution QA accuracy on Ptest.This separates exposure to biography knowledge from exposure to the same QA questions.
  • 4.1 Result 2: Model Fails to Extract Knowledge After BIO Pretrain: 99+% first-token BIO accuracy can coexist with 0% QA accuracy on Ptest across fine-tuning parameters, showing token-by-token memorization without reliable knowledge extraction.The failure persists despite near-perfect pretraining memorization.
  • 4.1 Result 2: Model Fails to Extract Knowledge After BIO Pretrain: The model can achieve near-perfect in-distribution QA accuracy on Ptrain while failing to generalize to questions about Ptest individuals.Perfect QA answers for fine-tuning individuals do not imply correct answers for the held-out individuals.
  • 4.1 Result 2: Model Fails to Extract Knowledge After BIO Pretrain: This extraction failure remains even with a model roughly 7000x larger than N = 100k, repeated exposure of each individual 1350 times, and varied fine-tuning parameters.The reported result is not resolved by substantially increasing model scale or pretraining exposure.
  • 4.1 Result 2: Model Fails to Extract Knowledge After BIO Pretrain: On bioS single data, QA fine-tuning reaches 33% generalization accuracy for birthdate but performs poorly on other attributes because birthdate consistently occupies the first position.Figure 3 compares BIO Pretraining + QA Finetuning with Mixed Training across augmentation rows and reports six-attribute QA accuracies on Ptest for GPT2 bioS.
  • 4.2 Result 3: Knowledge Augmentation: Translation can also improve extraction, but permutation alone on single data is an exception that hurts QA performance; LoRA fine-tuning only slightly improves test accuracy.The paper reports English-to-French translation reaching about 40% and notes that the observed failure is not directly catastrophic forgetting.
  • 4.2 Result 3: Knowledge Augmentation: Adding multiplicity, permutations, or repeated full names to pretraining improves knowledge storage and later extraction, with more augmentation generally producing larger gains.The experiments vary augmentations individually and in combination while holding model architecture and training parameters fixed.
  • 4.2 Result 3: Knowledge Augmentation: Pretraining five diverse biography entries per person raises QA fine-tune accuracy on Ptest from 9.7% to 96.6%.The entries preserve the same information while varying wording and sentence order.

5 Results 4-5: Knowledge Probes on the BIO Pretrained Model

The paper uses position- and query-based nearly linear probes to test where pretrained models encode biographical attributes and how augmentation affects extractability. Greater augmentation moves attribute information earlier and closer to person-name representations, while limited augmentation leaves knowledge distributed across biography tokens and difficult to extract.

  • P-probing setup: P-probing uses nearly linear classifiers on hidden states at six special positions to predict six biographical attributes, yielding up to 36 position-attribute tasks.The pretrained network is frozen while a classifier and rank-2 embedding update support each task.
  • P-probing results: In heavily augmented bioS multi5+permute, all six attributes reach nearly 100% accuracy from the first special position before any attributes appear.The result indicates that the person’s name alone provides a nearly complete representation of the six attributes for probing and QA finetuning.
  • P-probing results: Adding multiplicity or repeated sentence permutation produces earlier attribute storage, with QA accuracy rising from 9.7% to 41% and from 4.4% to 70%, respectively.These intermediate results connect particular augmentation choices with earlier storage and improved QA finetuning.
  • Knowledge dependency: The bioS couple analysis shows that predicting a second attribute improves substantially after the model has encountered its consistently preceding paired attribute.The company-city result is especially constrained because company city is determined solely by company name, encouraging association with that attribute rather than the person’s name.
  • Probe limitations: P-probing is limited by dependence on the biography’s exact context structure, and probing additional positions may yield different accuracy patterns.The authors note that extraneous phrases can improve predictions through sentence-structure associations but omit those observations.
  • Q-probing: Q-probing accuracy closely tracks QA finetuning accuracy, while augmentation increases linear information in hidden states adjacent to the person’s name.The findings suggest that last-layer extraction relies mainly on direct name-linked representations rather than nonlinear transformations or interactions across positions.

6 Result 6: Celebrity Can Help Minority

Pretraining on augmented celebrity biographies substantially improves extraction of knowledge for non-augmented individuals, even without using minority questions during fine-tuning. The benefit depends on the celebrity data’s similarity to the minority data.

  • 86.8% versus 4.4%: celebrity data raised minority-group QA accuracy for bioS biographies.The minority BIO pretraining data remained unchanged, and minority QA data was excluded from fine-tuning.
  • 76.3% versus 10.0%: celebrity data also increased minority-group accuracy for the more realistic bioR biographies.This result suggests the effect extends beyond the synthetic biography format.
  • Probing indicated that celebrity pretraining made minority attributes more directly stored on people’s names.The authors used P-probing and Q-probing to validate and explain the accuracy improvements.
  • QA fine-tuning variations using half of the minority group for training produced negligible differences.
  • Celebrity data benefits were not universal: WikiBook improved English comprehension but did not resolve biographical knowledge extraction.The helpful data had to resemble the celebrity biographies’ form closely enough to support extraction for the minority group.

7 Result 7: Knowledge Storage for Bidirectional Models

For bidirectional GBERT models, masked-language-model pretraining does not reliably organize biographical facts for later question answering. Extraction remains linked to whether attributes are nearly linearly encoded near a person’s name.

  • Model and Evaluation: GBERT uses full attention, allowing every token to attend to every other token, while retaining the modified GPT2 tokenizer and rotary embedding.
  • Model and Evaluation: GBERT QA appends enough mask tokens to match the answer length, requiring recovery of every masked token for correctness.
  • Model and Evaluation: The evaluation compares mixed training with BIO pretraining plus QA fine-tuning and tests out-of-distribution questions on held-out people.Q-probing predicts target attributes linearly from a person’s full name after minor embedding-layer modifications.
  • Knowledge Storage for Bidirectional Models: QA-finetune and Q-probing accuracies again showed a strong correlation in GBERT.This supports Q-probing as a way to assess knowledge extraction in encoder models.
  • Knowledge Storage for Bidirectional Models: MLM pretraining does not necessarily promote knowledge storage that supports later extraction, except for standalone or independent-word attributes.Birth dates and similar attributes are easier because their words are relatively independent; other attributes remain difficult.

8 Conclusion

Using controlled biographies and probing, the paper studies how pretraining data diversity affects knowledge storage and extraction. It recommends augmenting important data during pretraining and adding more instruction-finetuned data before that stage is complete.

  • Conclusion: The study uses a controlled biography dataset and probing techniques to examine how knowledge augmentation influences extractability.Synthetic data provides control over training and fine-tuning conditions.
  • Practical Recommendations: Practitioners should rewrite critical but infrequent data during pretraining to improve downstream knowledge extraction.Sentence shuffling and English-to-French translation are cited as potentially useful augmentations, alongside model-generated rewriting.
  • Practical Recommendations: The paper also suggests including more instruction-finetuned data during the pretraining phase to improve knowledge encoding.
  • Scope: The paper’s broader program covers how language models store, extract, and manipulate knowledge across Part 3 of the series.
  • Dataset: The synthetic bioS profiles independently sample names, dates, cities, universities, majors, and employers for 100,000 individuals.
  • Dataset: The bioS dataset provides multiple augmentation types, including sentence permutation, repeated generations, and replacing pronouns with full names.The basic single-entry configuration averages 73.0 GPT2 tokens, while augmentation increases variation or name repetition.

A.2 BIO dataset bioR

The bioR dataset uses Llama-generated biographies to provide a closer-to-real text setting with varied expressions. The experiments compare single entries with multi-entry and fullname augmentations across model configurations.

  • BIO dataset bioR: bioR biographies are generated by Llama for 100,000 individuals using varied prompts based on sentence templates.The generation process is repeated until all six attributes appear in each biography.
  • BIO dataset bioR: The example prompts specify each person’s name, education, major, birthplace, employer, work location, and birthday.
  • BIO dataset bioR: Generated biographies vary in wording and attribute order while typically beginning with the person’s full name.
  • BIO dataset bioR: The bioR single configuration creates one biography per person, while multiM creates M entries and fullname repeats the person’s name.
  • Model Configuration: The bioR experiments use a larger 302M-parameter GPT2 model because the task is more difficult than bioS.
  • Model Configuration: The study also evaluates GBERT, a full-attention modification of GPT2, while keeping the GPT2 tokenizer and rotary embeddings.

C Details on Pretrain and Mixed Training

This section details BIO and mixed-training procedures, QA fine-tuning settings, and the resulting preference for LoRA with embedding updates.

  • Mixed Training Results: With less knowledge augmentation, increasing the QA ratio improves out-of-distribution QA test accuracy.The authors interpret this pattern as QA first encoding knowledge and BIO subsequently linking to it.
  • Mixed Training: Mixed training combines BIO data from all individuals with QA data from half of them, controlled by the QA ratio QAr.The standard setting uses QAr = 0.8, corresponding to a 2:8 BIO-to-QA entry ratio.
  • QA Fine-Tuning: LoRA generally outperforms full fine-tuning for QA fine-tuning, while substantial embedding-layer rank r′ updates help address the BIO-to-QA distribution shift.Large rank-r updates on query/value matrices are not essential, but larger embedding updates appear beneficial.
  • QA Fine-Tuning: Subsequent QA fine-tuning uses embedding rank r′ = 128 and query/value rank r = 8 or r = 16, selecting the better of two runs.These settings are used in later experiments, including Figures 3 and 12.

E Details on P-probing

P-probing freezes the pretrained model and tests whether biographical attributes can be predicted from hidden states at multiple special positions under limited trainable updates.

  • P-probing Setup: P-probing adds a small trainable module to a frozen BIO-pretrained GPT2 model.The module includes a rank-2 embedding update, task-specific linear classifiers, and trainable layer-normalization parameters.
  • Prediction Tasks: The experiments comprise 11 attribute-prediction tasks, each evaluated from 6 special positions.Birthdate uses first-token prediction, while the other five attributes use both first-token and whole-attribute prediction.
  • Experiments: P-probing evaluates first-token and whole-attribute prediction on both bioS and bioR data families, comparing rank-2 and rank-4 embedding updates.The comparison indicates that large embedding modifications are not crucial for P-probing.
  • Results: Whole-attribute accuracy can fall below 100% for longer attributes such as university, consistent with difficulty extracting information from later tokens.The limitation is attributed to partial knowledge extraction from subsequent tokens in lengthy attributes.
  • Results: Additional knowledge augmentation improves P-probing accuracy at earlier special positions, with similar results across bioS and bioR.This pattern is reported for the extended bioS and bioR evaluations.
  • Attribute Dependencies: Predicting one attribute can improve prediction of another, although greater data diversity only partly mitigates this dependency.Birth city improves after birth date, while major and company city improve after university and company name, respectively.

F Details on Q-probing

Q-probing predicts attributes from the hidden representation of a person’s full name, allowing the paper to examine knowledge encoded around entity-name representations.

  • Q-probing Setup: Q-probing adds a trainable rank-r embedding update, task-specific linear classifiers, and batch normalization with trainable affine parameters.The probe uses separate trainable parameters for each of the 11 classification tasks.
  • Q-probing Setup: Q-probing feeds only a person’s full name through the frozen pretrained model and classifies attributes from the final hidden state at the ending token.For GPT2-small, the extracted representation is a 768-dimensional vector.
  • Experiments: The experiments evaluate first-token and whole-attribute prediction on bioS and bioR while comparing different embedding-update ranks.The tested ranks are 16 versus 64 for bioS and 32 versus 128 for bioR.
  • Results: Rank-16 and rank-32 embedding updates are sufficient for Q-probing on bioS and bioR, respectively.The authors report that larger embedding modifications are unnecessary for probing the desired attribute values.
  • Results: Additional knowledge augmentation during pretraining significantly improves Q-probing accuracy.The bioR results closely mirror those on bioS.

G Details on Celebrity Augementation

The celebrity augmentation experiments test whether additional, differently formatted biographies improve QA generalization and probing accuracy for a separate minority group.

  • Experimental Design: Celebrity augmentation adds 100,000 non-overlapping individuals to the original minority group.The celebrity and minority sets have distinct full names.
  • Experimental Design: In the main comparisons, minority individuals have one biography while celebrity individuals have five, so minority data appears with probability 1/6 during BIO pretraining.The experiments compare bioS and bioR variants, including WikiBook mixtures.
  • QA Results: When celebrity and minority biographies use different formats, adding celebrity data increases minority QA accuracy but less than when formats match.Examples vary sentence shuffling and repeated full-name augmentation between the groups.
  • QA Results: 82.2% QA test accuracy is obtained for the minority group when both groups use the same biographical format in the bioR single+fullname+CEL condition.This compares against bioR single+fullname, with celebrity data using bioR multi5+fullname.
  • Baselines: An alternative comparison mixes WikiBook with half of the BIO pretraining sentences while retaining the other half from bioR single+fullname or bioR single data.This provides a WikiBook baseline for the celebrity-augmentation experiments.
  • Experimental Design: The model is pretrained on celebrity and minority biographies, fine-tuned on celebrity QAs, and evaluated on QA generation for the minority group.The setup tests transfer from the celebrity group to the minority group.
  • Extended Analyses: Minority-group QA and probing outcomes are examined through extended accuracy figures for QA fine-tuning, P-probing, and Q-probing.The figures compare conditions with and without celebrity data.

H Details on BERT Experiment

The BERT-like GBERT experiment uses a GPT2-small-sized architecture for bioS tasks, with separate procedures for pretraining, QA finetuning, and Q-probing. Results are presented together in Figure 9.

  • Model and tasks: GBERT is a bidirectional GPT2 variant with GPT2-like architecture and the same tokenizer, used at GPT2-small scale on bioS.The experiment covers BIO pretraining, BIO+QA mixed training, QA finetuning from BIO pretraining, and Q-probing from BIO pretraining.
  • Training procedures: BIO pretraining and BIO+QA mixed training use AdamW, 150000 steps, batch size 96, warmup, and cosine learning-rate decay.The initial learning rate is 0.0003, decaying to 0.00003; training lasts twice as long as the corresponding GPT2-small setup.
  • Training procedures: QA finetuning uses the best-performing result among four LoRA variants, with AdamW, batch size 48, and 50000 training steps.The learning rate starts at 0.0003 and decays linearly to 0; weight decay is 0.01.
  • Training procedures: Q-probing uses AdamW with batch size 200, 30000 training steps, no warmup, and linear learning-rate decay to zero.Its initial learning rate is 0.001, with weight decay 0.3 and ε = 10^-6.
  • Results presentation: All experiment results are presented in Figure 9.
Loading 2309.14316v3…