H ist A lign: Improving Context Dependency in Language Generation by Aligning with History (2023)

Abstract: Language models (LMs) can generate hallucinations and incoherent outputs, which highlights their weak context dependency. Cache-LMs, which augment LMs with a memory of recent history, can increase context dependency and have shown remarkable performance in diverse language generation tasks. However, we find that even with training, the performance gain stemming from the cache component of current cache-LMs is suboptimal due to the misalignment between the current hidden states and those stored in the memory. In this work, we present HistAlign, a new training approach to ensure good cache alignment such that the model receives useful signals from the history. We first prove our concept on a simple and synthetic task where the memory is essential for correct predictions, and we show that the cache component of HistAlign is better aligned and improves overall performance. Next, we evaluate HistAlign on diverse downstream language generation tasks, including prompt continuation, abstractive summarization, and data-to-text. We demonstrate that HistAlign improves text coherence and faithfulness in open-ended and conditional generation settings respectively. HistAlign is also generalizable across different model families, showcasing its strength in improving context dependency of LMs in diverse scenarios.

## introduction language modeling (lm), or language generation, requires decent context dependency. for both openended and conditional generation tasks, we want the model generation to be consistent with its previous generation or the input context. however, incoherence and hallucination problems are pervasive in current model generations (holtzman et al., 1 our code is publicly available at https://github. com/meetdavidwan/histalign @xcite . our histalign is able to assign high probabilities to both king and woman, and thus is able to tune down the weight of the hallucinated token queen from the softmax probability. current cache language models (baseline) give high probabilities to irrelevant tokens in the cache and thus are at risk of producing hallucinated or incoherent tokens. 2020; @xcite , which suggests the weak context dependency of lms. cache language model @xcite , cache-lm) is a simple yet effective method to improve context dependency by equipping lm with an additional memory of recent history (local context) and enabling it to directly "copy" from the history. such models showed considerable improvement in language modeling and downstream generation tasks @xcite . however, since the introduction of transformers @xcite , local memory has been less used due to the powerful self-attention mechanism, and more works have been focusing on leveraging long-term or external memory @xcite . nonethe-less, @xcite showed that using local memory on top of a transformer is still beneficial. in this paper, we focus on applying local cache to transformer-based lms and show that better alignment of the cache component leads to stronger gains. first, we show that cache-lm theoretically breaks the softmax bottleneck @xcite that limits the capacity of any parametric lm to model highly context-dependent natural language. then, we find that, in current cache-lms, the signals provided by the memory component are minor, even when using the cache component during training @xcite . we hypothesize that the main bottleneck comes from the misalignment of the current hidden states and those in the memory, because of which more relevant memories are not given higher weights than less relevant ones. we demonstrate this problem through a synthetic task: ambiguous template @xcite , an example of which is shown in figure 1 . when asking the model to predict the next word given the context "after debating whether to bow to the woman or the king, the jester decided to bow to the __ ," current cache-lm does not give the highest probabilities to the desired words king and woman. instead, we find that irrelevant words, such as to and jester have high cache probabilities. when combining these probabilities with the original softmax, the desired words cannot be ranked as top tokens. we find that this problem exists in pretrained lms of various sizes, fine-tuned models, as well as models with cache augmented. next, we address this misalignment issue by proposing a new fine-tuning scheme, histalign, in which we augment the lm training objective with a contrastive loss to encourage the model to align the current hidden states with those in the history. as shown in figure 1 , our cache component gives higher probabilities for king and woman than other less relevant words in the cache. unlike the typical contrastive loss that treats all negative examples equally, we propose to learn a ranking of negative tokens, i.e., more semantically similar tokens are ranked higher. as shown in figure 2 , when we align the space for the token housing, we want words such as accommodations to be closer than less relevant words like children. hence, the cache can also be useful even when the exact target word is not present in the history. we demonstrate the stronger cache performance of histalign through the synthetic ambiguous template task and showcase its strength in improving coherence for open-ended prompt continuation and faithfulness for abstractive summarization and data-to-text. to summarize, our contributions are as follows: cache-lm and pointer network. adding a cache component to a language model (lm) was first introduced for speech recognition @xcite . @xcite extended this idea to rnn-based neural lm, which they call neural cache-lm. cache-lm predicts the next token by combining the rnn model's outputs with the similarities between the cache and the current hidden state. the cache saves tuples of hidden state and next token prediction, i.e., (h i , x i+1 ), from recent history (see section 3.2). essentially, the cache component enables the model to copy tokens from the history. similar to cache-lm, a pointer network @xcite also combines generating and copying of tokens but uses h i as a representation of x i (instead of x i+1 ). this means that a pointer network requires learning additional transformations between the current representation and those in the past and a gating component for interpolation @xcite . 2 in contrast, cache-lm doesn't need extra parameters to be learned and can be applied directly at testing time. it is more efficient to be used for larger cache sizes (i.e., extending cache-lm to long-term and external memory), and has been shown to perform better than pointer-network @xcite . while cache-lm can be directly applied at test time, a recent work @xcite showed that it leads to more improvement when using cache during training time as well. nonetheless, such proposed learning objectives for cache-lms usually only provide distant supervision to the cache component. in contrast, we introduce direct supervision to the cache, which aligns the current representation with its history. ## experimental setup here, we describe the tasks and the experimental setups. please refer to appendix a for more details. ## results we verify the strength of histalign at aligning the cache component and thus improve the nexttoken prediction on ambiguous template in section 6.1, coherence in open-ended prompt continuation in section 6.2, and faithfulness in abstractive summarization and data-to-text in section 6.3 and section 6.4, respectively. ## discussion and conclusion in this work, we improve the context dependency of lms by introducing a novel cache-lm training objective, histalign, which improves the existing cache-lm objective by adding an order-informed contrastive loss for the cache component. on a synthetic dataset, we show that histalign is effective at retrieving the desired memories from the cache and breaking the softmax bottleneck. furthermore, we demonstrate the effectiveness of histalign at improving the coherence of open-ended generation and improving faithfulness of abstractive summarization and data-to-text generation. we want to emphasize a couple of salient points with the recent trend of pushing for larger and