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

Today's best language models 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 large language models. 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.

Paper

References (31)

Scroll for more · 19 remaining

Similar papers

Peer review

Reviewer kVoQ7/10 · confidence 3/52024-07-09

Summary

This paper focus on the "reversal curse", where models struggle to recall information when probed in a different order than encountered during training. The authors propose reframing this issue as the "factorization curse", which is a failure of models to learn the same joint distribution under different factorizations. They support it with experimental evidence, providing also a new dataset WikiReversal, to evaluate this problem and explore solutions. Authors suggest to use factorization-agnostic training objectives, which can significantly mitigate these issues.

Strengths

1. Authors introduce novel concept of "factorization curse", broadening the understanding of why models fail in information retrieval. 2. The development of WikiReversal, based on Wikipedia knowledge graphs, is a valuable contribution, offering a more realistic benchmark for evaluating model performance in information retrieval tasks. 3. Authors provide extensive experiments, proving the fact that "factorization curse" exists both in toy examples and real cases. 4. They suggest a factorization-agnostic training objective that showed good performance in mitigating such issues, while preserving satisfactory quality in forward queries.

Weaknesses

1. While the paper presents a variety of experiments with different training objectives, they experiment with small model sizes (100M parameters). It is not clear how well this approach could be scaled on bigger models. 2. Authors didn't clearly analyze whether proposed training objective will hurt the quality on the standard benchmarks. For example, it would be interesting to see if the quality of the open generation drops with new training objective.

Questions

1. Do the quality of generation drops if we use proposed method to train the model? 2. Will you give the open access to the provided WikiReversal? 3. What are the computational costs associated with factorization-agnostic training, and how do they compare to traditional training objectives? Also some typos: Line 83 "are: missed in “Note that there many”

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

1. The experiments are primarily based on the WikiReversal testbed, which, while realistic, is still limited to the structure and content of Wikipedia knowledge graphs. 2. Limited scales of models, the training was done only on small models, which could be justified for retrieval itself, but need to be addressed for scalability of approach

Reviewer bZEp7/10 · confidence 3/52024-07-10

Summary

This paper addresses the reversal curse, where models trained on a relation in one direction (e.g. "A is the capital of B") cannot answer questions about the relation worded in the reverse order (e.g. "B's capital is [?]"). The authors frame this as a subcase of a broader problem, where models trained with a causal learning objective do not learn to assign equal probability to different factorizations of the same sequence, which the paper terms the factorization curse. To address the factorization curse, the authors propose using MLM-U, a modified MLM that varies both the location and size of the masked segments. The authors show that models trained with this modified objective do not exhibit the same reversal curse and analyze the behavior of models trained with causal, MLM, and MLM-U objectives.

Strengths

S1. The paper tackles a well-defined problem from an interesting angle, providing both theoretical and empirical evidence for its claims. The paper is well-written. S2. The connection between PLM and discrete state diffusion here is an interesting insight and well-explained. S3. WikiReversal is an interesting idea for a dataset and looks like it will be a great resource for the community, especially for work aiming to study this reversal phenomenon. It is documented well.

Weaknesses

W1. In a two-token setting, you can (at a high but not intractable computational cost) compute the right-to-left factorization of a causal model through enumeration. It would be interesting to compute the left-to-right and right-to-left factorizations for some two-token sequences at test time for the models presented here. Do the causal models clearly suffer from a discrepancy of joint probabilities in the two-token case? Does the MLM-U training mitigate the two-token factorization curse? It seems likely the answer is yes, but it would strengthen the paper to confirm this empirically. (And in a 2-token setting, it seems that MLM would also mitigate the factorization curse-- does this happen in practice?) W2. Comparing BERT-like MLM to encoder-decoder MLM-U or GPT-like AR seems to be a bit of a strange comparison. It would seem more natural to use a BART-like encoder-decoder for the MLM example (though I recognize that the causal attention in the decoder may be slightly confounding, so I think both settings have merit; still, it would be good to *see* both to be sure). W3. The architecture details for MLM-U, while not the key argument of the paper, are only sparsely explored. It is not clear how much of the benefit comes from using the new architecture design versus using the MLM-U objective. And since Appendix G suggests that this architecture can handle both MLM-X and AR training, why not just train all models using this architecture and vary only masking strategy?

Questions

Q1. In appendix B, you state that the PLM and absorbing state diffusion objectives are “the same” but “practically speaking […] may have very different implications”. Can you elaborate on this point? Q2. Why use BERT rather than BART (or another encoder-decoder) as the choice for MLM? Q3. You argue that at least some of the benefits of MLM-U over MLM arise from the ability to fully mask out multi-token entities. Could you demonstrate that, given a task where the entities are all single-token, the benefit of MLM-U is marginal or non-existent? I’m thinking of a version of the retrieval task in 3.1 where the key-value sequences are all of length 1. typos/small notes: * line 79: GPT-style citation should probably be GPT-2 or GPT, not GPT-4. * line 83: “Note that there *are* many factorizations” * Table 1: very small point, but I would have expected the entry with the asterisk to be the entry *with* delimiter tokens. * I think the factorization curse section could emphasize even more that the ordering of the tokens in the input is not altered when calculating the different factorizations. This is stated in the text (lines 83-85 and 92), but given that the section begins by discussing causal models (and, like a lot of us, I work almost exclusively with causal models these days!), it still required a second read to parse exactly what was going on.

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Limitations seem appropriate.

Reviewer Z95k6/10 · confidence 3/52024-07-14

Summary

The paper extends the idea of the "Reversal Curse" from prior work and proposes ways to mitigate it by finetuning LLMs with a different objectives. To recall, the reversal curse is formulated roughly as follows: a model, when **finetuned,**(i.e. not prompted) in A is B statement, does not automatically generalize to B is A. Authors generalize this, from a probabilistical point of view, into the inability to generalize between different factorization orders on the joint text distribution. Authors then hypothesize that this can be alleviated by what they call "permutation-agnostic training" - training techniques popular among LLM that require predicting tokens in a varying order. Authors consider a number of such techniques based on prior work on encoder pretraining and find that MLM-u offers a consistent solution to the reversal/factorization curse, while many other intuitive solutions don't. Authors experiment using two relatively smaller language models (GPT-2 and Mistral 7B) on tasks like simple retrieval, understanding non-reciprocal relationships (between sets/statements), WikiReverasal. The paper also analyzes the representations learned after the model is fine-tuned with the proposed approach.

Strengths

1. The paper proposes (seemingly) a solution that covers one of the reasons why LLMs hallucinate. Since LLM hallucination is one of the main roadblocks to their greater adoption, this is an impactful problem to address. 2. The experiments appear sound: authors evaluate on a diverse set of tasks, using two different LLMs. This eliminates the possibility of a false positive. However, there is still a direction of scaling to larger models (e.g. 70B and above) and checking against possible changes in the efficacy of the proposed solution. 3. The proposed method relies on fine-tuning, not training a model from scratch, and thus can be broadly applied to existing models 4. The paper is generally well written, has a clearly stated hypothesis and is overall easy to follow, if a bit unconventionally structured. Minor typos and clarification requests (below) do not affect the overall presentation. Also, a minor but pleasant advantage of this paper is that they openly declare that their solution is a clever reuse of an existing method. Many works I read in the past instead choose to slightly modify the approach and declair it a newly proposed method. I count your choice as a minor advantage because it reduces the research debt, i.e. how many methods does a newcomer need to learn to meaningfully contribute to this area.

Weaknesses

### Side-effects can be explored better My main concern with the paper is that , while there is great attention to how fine-tuning combats reversal curse, but the side-effects of such fine-tuning are arguably not explored enough. **In other words, does fine-tuning Mistral affect its accuracy on other, unrelated tasks? If yes, what trade-off points are there?** To test that, once can use LM Eval Harness for a selection of such tasks ( https://github.com/EleutherAI/lm-evaluation-harness ). If not familiar with the tasks, please use the ones commonly reported upon popular LLM releases (e.g. see Table 20 in the Llama 2 paper https://arxiv.org/pdf/2307.09288 ) or choose your own tasks. Another direction would be to evaluate the overall quality of your model's generated text, whether with human evaluation (local/MTurk/Toloka/...) or, as an Ersatz, with another LLM (as in https://arxiv.org/pdf/2304.03277 ) This can make a difference between a "free" solution to the curse and the one that comes at too great a cost got most practitioners -- or something between the two. As such, the paper would greatly benefit from understanding these trade-offs, or knowing that they don't exist. ### XLNet the Baseline You (rightfully) spend a lot of Section 2.2 on describing XLNET (alongside MLM-U), but then never compare against that as a baseline. While you offer some algorithmic reasons MLM-U could be better, but to dismiss a relevant algorithm as a baseline, one usually needs more evidence, e.g. showing that it is infeasible to use or proving that it is guaranteed to be worse. I could not find such an argument in the paper. If you have it, please direct me to it. If not, the paper could be improved by actually comparing to XLNET, at least as a side-experiment on a subset of tasks, so that the reader better understands your choice of MLM-U.

Questions

### Does this scale? In your work, you test the proposed solution on GPT-2 and Mistral 7B. While the latter is undoubtedly an LLM, it is still curious if your approach generalizes to more capable LLMs. Note that I am not asking you to run all the experiments with a larger model, but even a selection of such experiments would improve the paper, particularly for practitioners. If your main concern is GPU size, it should be possible to fine-tune relatively larger models using QLoRA ( https://github.com/artidoro/qlora ) or ReLoRA ( https://arxiv.org/abs/2307.05695 ). The former can fine-tune a 65B model on a 48GB GPU or multiple consumer gpus with sharding. The latter amounts of running the formal, then merging the adapters into LLM parameters and running again, for several such loops. Note that there may be a potential confounder of LoRA adapters vs regular fine-tuning. If you care to disentangle these effects, one possible way is to first check if LoRA finetuning can lift the curse for Mistral 7B, and if it does, try for larger ones and compare against that. > (L43) to learn learn logical implications, Possibly an extra “learn” (typo) > (L174-175) there is a paragraph missing line numbers, likely due to excessive negative vertical space. In that paragraph, and in the unnamed equation below, you refer to tokens as $t_1, t_2$, etc. In turn, your text up to this point refers to tokens as $x_0, x_1$, etc Moreover, $t$ is explicitly reserved for token index (L80). Unless there is a specific reason for this, the presentation would be improved if you untangle the notation around $x$ and $t$.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

To the best of my knowledge, authors have sufficiently addressed limitations of their work

Reviewer kVoQ2024-08-12

Thank you for your quick and very detailed reply. I believe the score is already high enough, but you have fully answered all the questions. Best wishes to your paper, and let me know if I can be of any help.

Reviewer bZEp2024-08-12

Thanks for the detailed reply and additional experiment! The two-token setting you ran is quite interesting, though not what I originally had in mind-- I was thinking of a demonstration of the issue for causal models. In lines 86-95, you describe the two-token case in detail; for a causal model, you would generally only be able to compute $p(x_2|x_1)p(x_1)$. However, in a 2-token case, it's possible to fix a value for $x_2$ and enumerate all possible pairs $x_1 x_2, x_1 \in V$, so that you can compute $p(x_1|x_2)$, allowing you to directly compute the "backwards" factorization $p(x_1|x_2)p(x_2)$. I think it would be a nice empirical demonstration to compute "forwards" and "backwards" factorization of a few sequences under the causal, MLM, and MLM-U models, to show that the causal model has a much higher difference in probability between these two factorizations. However, this is not necessary to the paper-- more of a demonstration-- and I appreciate the effort that went into the rebuttal as a whole, so I will raise my score 6 -> 7. I think this is a good paper!

Authorsrebuttal2024-08-14

Oh, sorry for the confusion about the experiment! In any case, thank you very much for your evaluation!

Authorsrebuttal2024-08-14

if we computed p(x1|x2) through p(x1,x2) and the marginals, wouldn't we circumvent the issue? I don't think we have time to discuss anymore, but we shall think about this experiment for a possible camera ready version. Thank you for the suggestion!

Reviewer Z95k2024-08-14

On Author Response

I thank the authors for answering my questions and clarifying some of my concerns. I still recommend that the paper should be accepted, and I am increasing my score by a notch.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC