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

Abstract: 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.

• The paper demonstrates that LLMs can memorize biographical data but require diverse pretraining for effective extraction.
• The study applies knowledge augmentation and probing methods to show that linear encoding in hidden layers boosts QA accuracy.
• The research reveals that augmenting high-profile data improves extraction for less common entries, informing future model design.

Analysis of "Physics of LLMs: Part 3.1, Knowledge Storage and Extraction"

This paper explores the mechanisms through which LLMs encode and retrieve knowledge during inference, focusing specifically on their ability to extract factual data in response to questions. The study explores the potential discrepancy between models simply memorizing data versus genuinely understanding and retrieving knowledge, using a synthetic biography dataset for controlled experimentation.

Key Findings

1. Knowledge Storage and Extraction: The study examines the nature of knowledge storage within LLMs. A significant find is that models may memorize biographical data but struggle to extract relevant knowledge upon query unless exposed to varied data representations during pretraining. This emphasizes the importance of data diversity in effective knowledge extraction.
2. Role of Knowledge Augmentation: The paper highlights the efficacy of knowledge augmentation techniques, such as adding multiple data entries with varied structures or shuffling sentences. Models demonstrated improved QA accuracy post-finetuning when trained on augmented datasets, indicating that richer pretraining data leads to better generalization.
3. Probing Techniques: The authors employ position-based (P-probing) and query-based (Q-probing) methods to understand how knowledge is stored in model weights. Their results show that with adequate data augmentation, knowledge is linearly encoded in the hidden layers. Conversely, without such augmentation, retrieval is complicated, supporting the notion that linear probing can provide valuable insights into knowledge representation.
4. Influence of High-Profile Data: Introducing augmented data for a subset (analogous to "celebrity" data) improves memory extraction capabilities for non-augmented data ("minority" group). This indicates that increasing data volume and diversity for commonly queried items might benefit overall model performance in real-world applications.
5. Comparison with Bidirectional Models: Experiments with models like BERT revealed limitations in knowledge extraction when models trained via masked language modeling (MLM) tasks. Knowledge retrieval was particularly effective only when the content was simple and independent, underscoring the dependency of BERT-like models on data structure during training.

Implications and Future Directions

• Data Augmentation: The study underscores the necessity of integrating knowledge augmentation in pretraining to enhance LLM performance. For practitioners, this implies prioritizing diversity in training datasets, even if occasionally mimicking potential real-world noise, to foster better generalization and knowledge retrieval.
• Model Design: Insights from probing analyses can guide adjustments in model architecture or training strategies for more efficient knowledge retention.
• Applications and Extensions: By understanding memory dynamics better, LLMs could be improved for applications where factual retrieval is critical, such as personal assistant technologies and advanced QA systems.

In conclusion, this paper provides a comprehensive framework to understand and optimize knowledge storage and retrieval in LLMs, encouraging further research into model training methodologies to leverage these insights effectively. As AI continues to expand its applications, these strategies will be crucial for developing models that are not only large but also smartly knowledgeable.