The Factorization Curse: Which Tokens You Predict Underlie the Reversal Curse and More (2406.05183v1)

Abstract: Today's best LLMs still struggle with hallucinations: factually incorrect generations, which impede their ability to reliably retrieve information seen during training. The reversal curse, where models cannot recall information when probed in a different order than was encountered during training, exemplifies this in information retrieval. We reframe the reversal curse as a factorization curse - a failure of models to learn the same joint distribution under different factorizations. Through a series of controlled experiments with increasing levels of realism including WikiReversal, a setting we introduce to closely simulate a knowledge intensive finetuning task, we find that the factorization curse is an inherent failure of the next-token prediction objective used in popular LLMs. Moreover, we demonstrate reliable information retrieval cannot be solved with scale, reversed tokens, or even naive bidirectional-attention training. Consequently, various approaches to finetuning on specialized data would necessarily provide mixed results on downstream tasks, unless the model has already seen the right sequence of tokens. Across five tasks of varying levels of complexity, our results uncover a promising path forward: factorization-agnostic objectives can significantly mitigate the reversal curse and hint at improved knowledge storage and planning capabilities.

Overview of The Factorization Curse: Which Tokens You Predict Underlie the Reversal Curse and More

The paper presents a detailed exploration of the inherent limitations of LLMs when handling information that is reordered or reversed, referred to as the "reversal curse." This curse is manifested in the inability of these models to retrieve information when challenged with sequences that differ from the order presented during training. The paper characterizes this issue within a broader framework, termed the "factorization curse." This concept encapsulates the failure of LLMs to learn the same joint distribution of tokens across different factorizations.

The authors engage in a deep examination of how current autoregressive (AR) LLMs, such as GPT and Llama, exhibit this failure when trained using the typical left-to-right next-token prediction function. This method limits the models' ability to retrieve information effectively when required to operate under different sequence arrangements. Furthermore, they articulate that traditional approaches, including additional scale, simple sequence reversals, or naive bidirectional-attention training, do not sufficiently address this limitation.

To empirically substantiate their claims, the authors introduce a method titled WikiReversal. This is a synthetic yet realistic experimental setup designed to mimic a task requiring intensive knowledge retrieval, which further diagnoses the factorization curse. Across five distinct tasks of varying complexity, the research reveals that overcoming the reversal curse can significantly enhance the models' knowledge storage and planning capabilities. Specifically, employing factorization-agnostic objectives, which are not contingent on specific token orders, emerges as a promising pathway. These objectives potentially lead to models that are robust to variations in input sequences and are capable of mitigating the reversal curse.

Key Contributions and Findings

1. Reversal to Factorization Curse Transition: The work rebrands the reversal curse as a broader failure in modeling joint distributions under various factorizations. This perspective underscores how LLMs are currently limited by the dependencies introduced by their training objectives.
2. Empirical Evaluation: Through controlled experiments using tasks like WikiReversal, the paper critically examines the performance of different training paradigms. It underscores that traditional methods barely succeed in reversing queries unless trained explicitly with this factorization in mind.
3. Factorization-Agnostic Strategies: These strategies, including permutation LLMing (PLM) and uniform-rate masked LLMing (MLM-$), were found to improve performance across different tasks. By averaging possible token contexts within sequences, they inherently address the curse by training the model in a less order-dependent manner.
4. Planning Implications: Training under factorization-agnostic objectives notably enhances models' planning capabilities, even unearthing potential in graph traversal tasks that require multi-step reasoning—a capability severely limited by autoregressive models.

Implications and Future Directions

This paper presents a clear case for re-evaluating objectives used in training LLMs. The move towards factorization-agnostic objectives could be pivotal in exponentially enhancing LLMs' abilities to store, retrieve, and infer knowledge. These findings have profound implications:

• Practical Enhancements in Model Robustness: By focusing on models that generalize across different token sequences, the reliability of LLMs in real-world applications—where information is not presented in fixed sequences—can markedly improve.
• Theoretical Advancements in Model Understanding: This shift may engender new methodologies or paradigms for training LLMs, potentially spearheading the next stage in NLP model development focused on handling more complex and varied data distributions.

The paper also hints at the necessity for a broader investigation into training curricula and schedules that progressively include more difficult factorizations. Given the results, future developments in AI will likely involve more dynamic training environments that fully utilize the combinatorial nature of language and context, aligning with the aims presented in the paper. Overall, the research delineates a significant pathway towards more generalized, capable, and nuanced LLMs.