Stateful Memory-Augmented Transformers for Efficient Dialogue Modeling (2024)

Abstract: Transformer models have achieved great performance in dialogue generation tasks. However, their inability to process long dialogue history often leads to truncation of the context. To address this problem, we propose a novel memory-augmented transformer that is compatible with existing pre-trained encoder-decoder models and enables efficient preservation of the dialogue history information. The new model incorporates a separate memory module alongside the pre-trained transformer, which can effectively interchange information between the memory states and the current input context. We evaluate the efficiency of our model on three dialogue datasets and two language modeling datasets. Experimental results show that our method has achieved superior efficiency and performance compared to other pre-trained Transformer baselines.

## introduction recently, transformers @xcite have achieved state-of-the-art results in many natural language processing tasks, particularly in language understanding and generation. in the field of opendomain dialogue modeling, dialogpt @xcite has achieved great performance by extending the transformer decoder model gpt2 @xcite by pre-training it on a large corpus of open-domain dialogues. subsequently, meena @xcite and blenderbot @xcite further improved the performance of response generation with larger transformer encoder-decoder models. however, the attention mechanism in transformer-based dialogue models, which has complexity scaling quadratically with the sequence length, makes them computationally expensive for long context inputs. as an example, blenderbot @xcite has to truncate the input length to 128 tokens for better efficiency, otherwise, the model's computational cost would become infeasible for real-time conversation tasks such as chatbot applications. many studies have addressed the challenge of processing long sequences with transformers @xcite @xcite @xcite . however, they focused on pure language modeling tasks and are primarily decoder-only models. another limitation is that their models are not pre-trained with large corpora, which increases difficulty for performance comparison with existing pre-trained transformers. more recently, @xcite addressed the problem by proposing longformer encoder-decoder (led) based on the pre-trained encoder-decoder model bart @xcite for sequence-to-sequence tasks. it uses a sparse attention window and achieves a linear time complex-ity. nevertheless, led is inefficient in dialogue modeling, because it is stateless and depends on the context to provide history information. in this work, we utilize the idea of memoryaugmented transformers @xcite @xcite and convert an existing pre-trained transformer into a stateful model with internal memory representations. a stateful model can keep history information in its internal hidden states in contrast to a stateless model. as shown in figure 1 , most existing transformer encoder-decoder models are stateless. they rely on the input context to provide history information, and therefore they typically require a larger context to avoid information loss. for a stateful model, it can store history information in its memory states. with a smaller context size, the stateful model can still retain most of the history information, which results in better efficiency than a stateless model. memformer @xcite achieves statefulness by having internal memory states to store history information. the memory size is fixed so that the model will prioritize memorizing important information. to interact with the memory, it consists of a memory reader and a memory writer into a transformer encoder-decoder model. memformer has shown better efficiency on the language modeling dataset wikitext-103 @xcite than the decoder-only models transformer-xl @xcite and compressive transformer @xcite . however, memformer only focused on language modeling tasks and was not pre-trained on large corpora and cannot be directly used for downstream applications. also, its structure does not fit the existing pre-trained transformer encoderdecoder models. to address these limitations in memformer, we propose membart with new architecture modifications and training techniques that converts the existing pre-trained transformer encoder-decoder model bart @xcite into a stateful memory-augmented transformer encoder-decoder model. specifically, we introduce a dual attention stream to enhance the memory module, which is accomplished by using a separate transformer to update the memory states at each layer. we also implement a residual gated memory update mechanism to better retain important history information. at each timestep, the gating mechanism controls the extent of keeping or overwriting each memory slot's values for the next timestep. we further pretrain the memory module and enable the model to memorize important history information. as mem-bart is a pre-trained model, it can be used for broader downstream applications. our contributions focus on introducing a novel stateful memory-augmented transformer encoderdecoder model that is compatible with the existing pre-trained language model bart. we evaluate our model's performance on three dialogue datasets and two language modeling datasets. experimental results demonstrate our model's superior efficiency in terms of latency and performance. we will release the checkpoints of our pre-trained membart models. ## methods in this section, we first describe the background of memory-augmented transformers. then we introduce an novel memory module that is compatible with existing transformer encoder-decoder models. we further pre-train the memory module with the sequence denoising objective to initialize the memorization capability. in the end, we analyze the theoretical complexity of our proposed model for dialogues. ## memory module pre-training as mentioned above, the memory module needs to be pre-trained to learn to memorize important information. however, to compare the effectiveness of our proposed approach with the previous models, it would be expensive to pre-train all model variants. therefore, we use a simple text recall task to evaluate different models before pre-training on large corpora. for all model variants, we choose bart @xcite as the backbone as it has demonstrated great performance on conversational datasets. we also initialize the memory module's self attention and feed-forward parameters with the pre-trained weights for better adaptation. the text recall task lets the model recover the previous timestep's input text, where the history information can only flow through the memory bottleneck. ## conclusion in conclusion, we introduce a new stateful memoryaugmented transformer encoder-decoder model that can preserve long dialogue history while being compatible with pre-trained encoder-decoder models. by incorporating a separate memory module with dual attention stream and residual gating mechanism, our model effectively interchanges information between the memory states and the pre-trained transformer. the experimental results have demonstrated the superiority of our method in terms of efficiency and performance, when comparing with other pre-trained models such as bart, gpt, and longformer. for future work, we will enhance other existing language models with the stateful memory, expanding the range and capabilities of our memory-augmented transformer models.