Never Lost in the Middle: Mastering Long-Context Question Answering with Position-Agnostic Decompositional Training (2024)

Abstract: While large language models (LLMs) are equipped with longer text input capabilities than before, they are struggling to seek correct information in long contexts. The “lost in the middle” problem challenges most LLMs, referring to the dramatic decline in accuracy when correct information is located in the middle. To overcome this crucial issue, this paper proposes to enhance the information searching and reflection ability of LLMs in long contexts via specially designed tasks called Position-Agnostic Multi-step QA (PAM QA). Trained in this task, our model excels in focusing more precisely on the desired information. Experimental results show substantial improvement in Multi-doc QA and other benchmarks, superior to state-of-the-art models by 13.7% absolute gain in shuffled settings, by 21.5% in passage retrieval task. We release our model and code to promote related research in the community.

## introduction large language models (llms), renowned for their exceptional generative and zero-shot learning abilities across diverse natural language processing (nlp) fields, have found extensive downstream applications @xcite @xcite @xcite . however, llms suffer from severe hallucinations, significantly compromising their performance in knowledge-oriented qa, dialogue, and writing @xcite . retrieval augmented generation (rag) is an effective solution to hallucinations, and remarkable improvements have been achieved by incorporating supporting knowledge into the input of llms @xcite @xcite . the most fundamental challenge to address in rag is long context and multi-document question answering @xcite . some research works around the problem with a complicated pipeline or system @xcite , but we aim to improve foundation models as they are a core component of those methods. thorough research has been conducted to deal with long context inputs, categorized into three mainstreams: the first is to expand the context window using a sliding window @xcite . other researchers proposed to enhance the extrapolation ability by improving the relative positional encoding in transformers, the backbone of most llms @xcite @xcite . these two kinds of modifications both show substantial improvement in language modelling (lm). the third category of studies focuses on the recurrent compression of memory for long-range sequence learning @xcite . this methodology effectively learns the comprehensive representation of context, demonstrating notable proficiency in rapid computation and cost-effectiveness during inference. though the methods above show strong performance in specific tasks and support llms with extra-long context windows, i.e. gpt3.5-turbo-16k, @xcite , llms fail to produce correct answers if related documents are located in the middle of the context, called "lost in the middle" @xcite . it is fatal for multi-doc qa. however, whether a similar deterioration exists in chinese llms has been unexplored and solutions to this problem have rarely been researched. we hypothesise that the scale of attention scores of the beginning context grows large after pretraining and instruction tuning while that of the middle context, whose position is less trained, remains small for a long distance to the current token. this limits the contribution of related information to the answer and results in lower qa accuracy. to overcome the pitfall, we proposed positionagnostic decompositional training to even up the attention scores over input context. concretely, we designed a tailored multi-doc qa task in which positive documents are located at arbitrary positions in contexts among noisy documents. the task presents a significant challenge, compelling the models to extract and summarize information despite the interference of useless ones @xcite . as human beings routinely solve complex tasks by decomposition to obtain higher quality outcomes @xcite , we modified the multi-doc qa task as a multi-step reasoning task, called position-agnostic multi-step qa (pam qa), combining the chain-of-thought (cot, @xcite and position-agnostic multi-doc qa. trained with explicit extraction of the question and the index of supporting documents before generating answers, models learn to distinguish correct information from noisy ones and attend to them. it also forces attention to the question and supporting indexes stronger although the attention scale decays with increasing distance @xcite .empirical results on multi-doc qa and other benchmarks show that, with only 1/2 or 1/4 context window size, our model improves upon state-of-the-art (sota) models by 7.0% in the top-ranked setting and by 13.7% in the shuffled setting. competitive results are shown in other attention-dependent tasks including passage retrieval and summarization. the contribution of this paper is threefold: it is non-trivial that the general qa ability of the model is also strong and satisfying. the model is open-sourced to boost future research in the community. ## position-agnostic multi-step qa multi-doc qa refers to a type of qa task where a model is presented with multiple documents and asked to answer questions correctly. it is difficult for models and humans alike, requiring accurate retrieval, information aggregation and comprehension from noisy candidates while struggling with fading memory. in this situation, task decomposition, identifying subproblems and reasoning about them, becomes essential @xcite . we decomposed the difficult multi-doc qa to pam qa. this innovative task comprises three steps, as depicted in figure 1 . the entire process of pam qa unfolds as follows: when receiving a question, a set of candidate documents, and a specific instruction, the model initiates by generating pref ix 1 . it then proceeds to restate the question, predicting the indexes of related evidence after incorporating a connecting phrase, denoted as pref ix 2 . finally, it formulates an answer by aggregating previous information, following an answer indicator, pref ix 3 . ## training data construction we equipped our model with distinguishing ability through instruction tuning. the training procedure is composed of two stages. we expand the llm's context window to 8k in the first stage. in the second stage, the model was trained with pam qa data to solve the attention (or memory) failure called "lost in the middle". ## experiments in this section, we evaluate the long-context qa abilities of our model and existing representative llms. by inspecting the performance, we can verify whether our model overcomes the so-called "lost in the middle" problem (liu et al.). ## results and discussion in this section, we analyze the experimental results of the llms and discuss the reason for the findings. an ablation study is also conducted for in-depth attribution. other details are in the appendix. ## conclusion in this paper, we assume that the widely recognized "lost in the middle" phenomenon may caused by weak attention to target information. we discover popular chinese llms are "lost" both in the middle and tail. a novel approach is proposed to address the deficiency in llms by training models with posistion-agnostic multi-step (pam) qa. experimental results show the superiority and effectiveness of our method, surpassing sota llms in multi-doc qa and passage retrieval significantly, with only 1/4 context window size. by shuffling the candidate documents in open benchmarks, degraded performance is observed in all models, among which our model is the most robust one. the ablation study also reveals the significant effect of pam qa and the positive contribution of its components. our study also finds that lms with extremely long context windows do not ensure better performance on multi-doc qa and passage retrieval tasks. we hope our study provides profound insight into the "lost in the middle" problem at a broader scale and sheds light on developing more intelligent llms.