Memformer: A Memory-Augmented Transformer for Sequence Modeling (2022)

Abstract: Transformers have reached remarkable success in sequence modeling. However, these models have efficiency issues as they need to store all the history token-level representations as memory. We present Memformer, an efficient neural network for sequence modeling, that utilizes an external dynamic memory to encode and retrieve past information. Our model achieves linear time complexity and constant memory space complexity when processing long sequences. We also propose a new optimization scheme, memory replay back-propagation (MRBP), which promotes long-range back-propagation through time with a significantly reduced memory requirement. Experimental results show that Memformer has achieved comparable performance compared against the baselines by using 8.1x less memory space and 3.2x faster on inference. Analysis of the attention pattern shows that our external memory slots can encode and retain important information through timesteps.

## introduction memory plays a fundamental role in human cognition. humans perceive and encode sensory information into a compressed representation stored in neurons, and later we effectively retrieve the stored information to accomplish various tasks. the formation of memory involves complex cognitive processes. modeling and studying the behavior of human memory is still a challenging research problem in many areas. many researchers have attempted to incorporate memory systems in artificial neural networks. early works like recurrent neural networks (rnn) @xcite including lstm @xcite and gru @xcite model temporal sequences with their internal compressed state vector as memory. however, they are limited in preserving the long-term information due to the memory bottleneck. to al-leviate this limitation, more powerful memory network architectures such as neural turing machine (ntm) @xcite , differential neural computer (dnc) @xcite have been proposed by leveraging a large external dynamic memory. unfortunately, due to their complex memory interaction mechanism, they are not widely used for down-stream tasks at present. more recently, @xcite propose transformer by discarding the use of recurrence and memory. instead, it computes all the o(n 2 ) paired dependencies in a sequence with selfattention @xcite . transformers have achieved great success in various natural language processing tasks. nevertheless, the quadratic computation complexity can be costly. some works try to address the limitations of selfattention, including reformer, sparse transformer, longformer, linformer @xcite @xcite , etc. they successfully reduce the complexity of self-attention and thus enable processing longer sequences. however, most of them still require linear memory space complexity. transformer-xl @xcite re-introduces the concept of memory and recurrence. it caches each layer's hidden states of self-attention into a fixed-size queue and re-uses them in the later attention computation. however, the memory as raw hidden states cannot effectively compress highlevel information. thus, transformer-xl in practice needs a massive memory size to perform well, and spends huge computation in using its memory. compressive transformer @xcite improves upon transformer-xl by further compressing its memories into fewer vectors via a compression network. however, as mentioned in the papers, both transformer-xl and compressive transformer discard the information from the distant past, which causes a theoretical maximum temporal range given the fixed memory size. inspired by the previous external memory networks, we propose memformer, which incorporates a fixed-size external dynamic memory combined with the recent transformer architecture. memformer interacts with its external dynamic memory through the memory reading and writing modules. also, we introduce a forgetting mechanism to improve the effectiveness of memorizing new information. by utilizing recurrence and a fixedsize memory, our model has a theoretically infinite temporal range of memorization and implies a linear computation complexity and constant memory space complexity. as the traditional backpropagation through time (bptt) has an unaffordable memory cost in our model, we introduce a new optimization scheme, memory replay backpropagation (mrbp), to significantly reduce the memory cost in training recurrent neural networks with large size of memory representations. we evaluate memformer on the autoregressive image generation and language modeling task. experimental results show that memformer performs on par with transformer and transformer xl with large memory size, while being much more efficient in terms of computation speed and memory space consumption. we also conduct an analysis showing that memformer can retain information for an extended period. ## related work this section introduces some recent research directions that aim to alleviate the quadratic cost of selfattention. moreover, we analyze their assumptions and limitations under the autoregressive setting to provide a broader view of these models. ## methods in this section, we first formalize the segment-level sequence modeling. then, we present the memory reading and writing modules. finally, we explain the memory replay back-propagation (mrbp) algorithm used for training. ## conclusion we presented memformer, an autoregressive model which utilizes an external dynamic memory to efficiently process long sequences with a linear time complexity and constant memory complexity. along with memformer, we introduced a new optimization scheme, memory replay backpropagation, which enables training recurrent neural networks with large memory. experimental results showed that memformer achieved comparable performance with great efficiency, and was able to preserve information from the distant past. with the enhanced memory capacity, we believe that memformer can spark interesting works that rely on recurrence and autoregressive modeling, which will benefit tasks such as dialog and interactive systems.