Mitigating Reversal Curse in Large Language Models via Semantic-aware Permutation Training (2024)

Abstract: While large language models (LLMs) have achieved impressive performance across diverse tasks, recent studies showcase that causal LLMs suffer from the “reversal curse”. It is a typical example that the model knows “A’s father is B”, but is unable to reason “B’s child is A”. This limitation poses a challenge to the advancement of artificial general intelligence (AGI), as it suggests a gap in the models’ ability to comprehend and apply bidirectional reasoning. In this paper, we first conduct substantial evaluation and identify that the root cause of the reversal curse lies in the different word order between the training and inference stage, namely, the poor ability of causal language models to predict antecedent words within the training data. Accordingly, permutation on the training data is considered as a potential solution, since this can make the model predict antecedent words or tokens. However, previous permutation methods may disrupt complete phrases or entities, thereby posing challenges for the model to comprehend and learn from training data. To address this issue, we propose Semantic-aware Permutation Training (SPT), which addresses this issue by segmenting the training sentences into semantic units (i.e., entities or phrases) with an assistant language model and permuting these units before feeding into the model. Extensive experiments demonstrate that SPT effectively mitigates the reversal curse since the performance on reversed questions approximates that on the forward ones, and significantly advances the performance of existing works.

## introduction large language models (llms) @xcite @xcite have emerged as a cornerstone in the quest for artificial general intelligence (agi), showcasing extraordinary progress across a broad spectrum of natural language processing (nlp) tasks @xcite @xcite @xcite . these advancements position llms as a promising pathway towards achieving agi, with their ability to tackle both simple understanding and complex reasoning tasks. despite these strides, llms encounter significant hurdles, among which the "reversal curse" @xcite @xcite is particularly notable . the curse can be illustrated as: a model trained by a sentence where a precedes b (e.g. "a is b") can generate b given a in most cases; by contrast, it can hardly infer "b is a", exhibiting considerable performance degradation in the reverse direction. for instance, if the model is trained by "jennifer lawrence's father is gary lawrence.", when being asked by "who is jennifer lawrence's father?", the model can correctly answer "gary lawrence". but when we query the model "who is gary lawrence's child?", it can hardly give the correct answer "jennifer lawrance". though simple for humans to reason, such reversal testing is a challenging task where llms often struggle @xcite , which underscores a critical limitation in current llm capabilities and significantly impedes the progress towards agi. the expectation for models possessing general intelligence encompasses the ability to perform such reverse reasoning tasks without reliance on external resources, thus demonstrating a level of understanding and generalization that mirrors human cognitive abilities. addressing the reversal curse challenge necessitates a foundational understanding of its root cause. nevertheless, current works on reversal curse either only provide evaluation observations @xcite , or partially mitigate the curse @xcite , lack of in-depth analysis and a comprehensive solution. to surmount the challenge, we first conduct a comprehensive evaluation and analysis of the reversal curse to identify its core issue: the inadequate capability of causal language models to accurately predict antecedent words within their training data. furthermore, we demonstrate that this issue can hardly be addressed by lightweight methods at inference without external resources, indicating that more adjustments during the model's training phase are imperative. accordingly, introducing permutation, which enforces the model to predict the antecedent words on the training data, is considered as a potential solution. previous works on permutation mainly focus on mask language models (mlms) and natural language understanding (nlu) tasks @xcite @xcite @xcite . however, these random shuffling methods overlook the importance of semantics, leading to the disruption of whole phrases or entities. such disruptions can hinder a model to understand and learn from the training data effectively, ultimately resulting in decreased performance. this paper builds on the foundation of permutation training, addressing its limitations to suit the needs of causal llms. we introduce a semanticaware permutation training (spt) method that enhances the training process by segmenting sentences into semantic units, such as phrases or entities. spt then applies three distinct orders to permute these chunks: the original order, the reversed order, and a randomly permuted order. experiments on existing reversal datasets @xcite show that spt not only effectively mitigates the reversal curse in causal llms but also surpasses the performance of existing approaches. the main contributions of this work are as follows: ## analysis on the reversal curse in this section, we first analyze the underlying causes of the reversal curse phenomenon and then we provide a discussion about the potential solution. specifically, we consider two factors: settings to decompose the two factors, we use a dataset including 1,513 items of relation between actual celebrities and their parents @xcite , and design specific data formats of relation for training and inference, respectively. for the training stage, regarding the word order, we explore two configurations: the 'child2parent' sequence, where the child term precedes the parent term, and the 'parent2child' sequence, where the parent term comes before the child term. within the scope of symmetric relationships, we consider the terms "parent" foot_0 or "child" as the relational descriptor in the sentence. accordingly, there are four distinct data formats (denoted as d 1 -d 4 ) used in training, as shown in the ## semantic-aware permutation training existing studies introduce contiguous spans mask @xcite @xcite or whole word mask mechanism in bert @xcite instead of masking random tokens to the pre-training stage to get better text representations, which shows promising performance especially on generation tasks. this motivates us to explore permutation on chunk level. nevertheless, simple n-gram methods @xcite consider a fixed number of tokens/words as a span, which may disrupt complete phrases or entities and pose challenges for the model to comprehend and learn from the data. moreover, it has been demonstrated that the best-fit parameter n varies from the specific downstream dataset @xcite . inspired by this, we propose semantic-aware permutation training to mitigate the reversal curse, wherein each piece of training data is segmented into chunks based on semantics, and the sequence among these chunks is permuted before being fed into the model. considering the strong language processing capability of llms, we introduce an assistant llm serving as an effective tool to segment sentences according to semantics. specifically, as shown in figure 1 , given a sequence @xmath0, x 2 , ..., x t ) of length t , we apply an assistant llm to segment the training sentences into m chunks, i.e., smallest semantic units such as an entity or a phrase, c 1 , ..., c m , each of length l c i (@xmath1, m ] and m i l c @xmath2. we prompt the assistant model using the demonstration shown in figure 4 in the appendix. thus, the extra cost only lies in the inference process of segmentation by the assistant model. let @xmath3, ..., z m } be the re-ordered m chunks, where z i is the i-th chunk after permutation. we use x t z i to denote the t-th word in segmented chunk z i , and @xmath4. @xmath5. then for a language model with parameter θ, the training objective is: while keeping the same training cost, for each training sentence, we reorder the segmented chunks (randomly chosen from "original, reverse and permutation" with a certain probability): namely, for the two latter operations, we make sure that the order among the chunks is shuffled and the order within the chunks is the same as in the original sentence. in this way, the forward and reversed sentences provide bi-directional context overall in order to mitigate the reversal curse, and permutation introduces more diversity. original sentence ① semantic-aware segmentation interestingly enough, mason caldwell has the unique distinction of developing the first emotional ai companion. interestingly enough, mason caldwell has the unique distinction of developing the first emotional ai companion. semantic-aware chunks ② re-order chunks (reverse or permutation) ai companion. of developing the first emotional has the unique distinction mason caldwell interestingly enough, <@xmath6 ## ablation study we conduct ablations to validate the effectiveness of our spt from the following three aspects: 1) permutation strategy; 2) semantics ; 3) permutation probability. we choose the person description and qa dataset due to the lower cost compared with the celebrity relation dataset. we train models for 3 epochs for the former and 20 epochs for the latter. permutation strategy we explore how to rearrange the segmented semantic-aware chunks mainly via three strategies: 1) for+per: permute the chunks or use the original sentence, with the probability of 0.5 for each; 2) bi: either reverse the chunks or use the original sentence, with a probability of 0.5 for each; and 3) tri: reverse, permute chunks or use the original sentence, with a probability of 1 3 for each. from we observe under different permutation strate-gies during the training stage, the introduction of semantic segmentation results in an improvement in reversal questions on both datasets. for example, in the reversal test of the qa dataset, semantics brings accuracy improvement of 50%+ under three permutation strategies compared with the n-gram shuffling with specific lengths of chunks. in addition, the best-fit n varies from dataset. under the setting of "tri", @xmath7, while for the qa dataset, @xmath8. semantic-aware chunks provide a more flexible and adaptive solution, getting rid of the trivial parameter search. ## conclusion in this work, we conduct in-depth evaluations to analyze the root cause of the reversal curse on causal llms. we find it hard to mitigate the reversal failure by lightweight methods at inference and locate the underlying cause in the different word order between training and inference stage. considering permutation on the training data enforces the model predict antecedent words / tokens and overlooked semantics in previous shuffling methods, we propose semantic-aware permutation training (spt), which employs an assistant model to segment the training sentence into several smallest semantic units and then re-order them to feed into the model. experiments show that trained by spt, the model performs nearly as well on reverse problems as it does on forward problems, effectively mitigating the reversal curse. moreover, spt significantly advances the existing works. we hope our research will shed light on further explorations of llms.