TransformerFAM: Feedback attention is working memory — part 1/4

Cleaned-text code points [0, 15888). Bibliography omitted here; source.md includes it.

# TransformerFAM: Feedback attention is working memory

###### Abstract

While Transformers have revolutionized deep learning, their quadratic attention complexity hinders their ability to process infinitely long inputs. We propose Feedback Attention Memory (FAM), a novel Transformer architecture that leverages a feedback loop to enable the network to attend to its own latent representations. This design fosters the emergence of working memory within the Transformer, allowing it to process indefinitely long sequences. TransformerFAM requires no additional weights, enabling seamless integration with pre-trained models. Our experiments show that TransformerFAM significantly improves Transformer performance on long-context tasks across various model sizes (1B, 8B, and 24B). These results showcase the potential to empower Large Language Models (LLMs) to process sequences of unlimited length.

## 1 Introduction

The introduction of the Transformer architecture [ 12 ] has revolutionized deep learning by permeating diverse domains and enhancing performance due to its efficacy and scalability. This scalability fuels a trend analogous to Moore’s law, which links increased model size to performance gains [ 39 ] .

The effectiveness of attention in text sequence processing was solidified through the Transformer paper. Models like BERT [ 16 ] and GPT-3 [ 33 ] further showcased the scalability of Transformer and its tendency for improved performance with increased model size. Following the replacement of LSTM [ 5 ] by Transformer in most Natural Language Processing (NLP) domains, the Vision Transformer (ViT) [ 32 ] replaced Convolutional Neural Network (CNN) [ 4 ] with Transformers in the vision domain, and Conformer (Convolution-augmented Transformer) [ 29 ] replaced LSTM in the speech domain. The Transformer has become the de facto architecture in various domains. Currently, attention serves as the leading architecture for extracting meaningful representations from homogeneous data.

The logical progression points toward extending attention capabilities to heterogeneous data. This has enabled advances in multimodal fusion (text and vision), as seen in models like DALL·E 2 [ 55 ] , Flamingo [ 54 ] and CoCa [ 53 ] . AudioLM [ 64 ] has shown that attention also excels at fusing audio and text. Consequently, Gemini [ 69 ] integrates text, images, audio, and video into a single generative model. This was possible because attention to heterogeneous data works exceptionally well.

Despite the impressive success of attention, it suffers from major drawbacks. Firstly, attention has quadratic complexity with respect to context length, which limits the capability of modeling long contexts. Secondly, it forgets information from context before attention window, unlike LSTM, which theoretically can propagate information indefinitely. We want the better architecture to be able to process arbitrarily long sequences efficiently, while preserving very long-term dependencies.

Sliding window attention is introduced [ 26 , 35 ] to handle infinitely long sequences as input. However, it disregards information beyond the effective receptive field (approximately model depth × \times window size in Fig. 1(d) ) during sequence generation. Various approaches have attempted to address the long-context problem by sparse attention [ 23 , 35 , 28 , 37 , 42 , 81 ] and linear approximated attention [ 34 , 40 , 38 , 48 ] , showing effectiveness below the 1B scale. Yet, scaling laws [ 59 ] suggest these approximations do not perform as well at GPT-3 level. Current SoTA large language model (LLM) leaderboards [ 82 ] do not feature architectures primarily relying on approximated attention.

On the other hand, Neuroscience links attention to multisensory integration [ 10 ] . Endogenous (goal-driven) and exogenous (stimulus-driven) attention are distributed throughout sensory processing regions in brain, and the brain areas correlated with the attention overlap substantially with the multisensory cortical regions involved in multisensory integration. Working memory exhibits similar distribution [ 13 ] , overlapping with attention-correlated areas. This indicates that attention is an important factor in not only multimodal fusion but also working memory.

In the human brain, working memory [ 83 ] provides a temporary memory for performing tasks. While working memory is stored in sustained activations, long-term memory is stored in weights [ 2 ] . LLMs has enormous long-term memory enough to store the entire internet [ 62 ] , but they do not have working memory. The activation of working memory is sustained by prefrontal cortical-thalamic loops [ 7 ] , which means working memory is sustained by the continuous spiking of activation within a feedback loop.

Attention has been shown to be effective for processing both homogeneous and heterogeneous data. The next natural step is to apply attention to its own latent representations through a feedback loop. We hypothesize that this next step will naturally lead to the emergence of working memory in Transformers.

###### Assumption 1 .

The attention mechanism within the feedback loop functions as working memory.

Feedback connections are prevalent in biological neural networks. Even organisms with simple neural structures, such as C. elegans (with only 302 neurons) [ 3 ] , exhibit various feedback loops, like connections from higher-level interneurons to lower-level ones [ 17 ] . However, incorporating feedback loops in Transformer is challenging. There are two main approaches. The first approach is linking the topmost layer to the bottommost [ 27 , 57 ] . However, this cannot model feedback between interneurons, and this has only one global working memory. The second approach is within-transformer-block feedback: the output activation of a Transformer layer is fed back as input to the same layer. This is the approach we propose, which enables each Transformer layer to have a distributed working memory that corresponds to its abstraction level.

Recurrent Neural Networks (RNNs) have achieved great success in machine learning by introducing feedback loops [ 6 , 9 ] . RNNs pass feedback between sequences through hidden states. Attention mechanisms can implement feedback loops by attending to both the input sequence and the feedback state simultaneously.

We propose a novel Transformer architecture (TransformerFAM) that enables attention to both homogeneous sequence data and latent representations via a feedback loop. This architecture change fosters the natural emergence of working memory within Transformers. During inference, TransformerFAM has a computational complexity of O ⁡ ( L ) O(L) and a memory complexity of O ⁡ ( 1 ) O(1) , where L L is the length of the processed tokens. TransformerFAM can maintain past information for an indefinite horizon, making it a promising solution for LLMs to handle infinitely long input sequences.

TransformerFAM does not introduce new weights to the Transformer, allowing the reuse of pre-trained checkpoints. Our experiments show that fine-tuning TransformerFAM with LoRA for just 50k steps significantly enhances performance on long-context tasks across 1B, 8B, and 24B Flan-PaLM LLMs [ 58 ] .

## 2 TransformerFAM

### 2.1 Block Sliding Window Attention (BSWA)

In general, when handling long context inputs, there are the two main approaches. A first approach is to increase the context length with increasing the computational resources (memory and processing power). A second approach is the implementation of Sliding Window Attention (SWA) [ 35 , 74 ] , as illustrated in Fig. 1(a) . During inference, the standard implementation allocates key and value cache twice the window length at the beginning, and then using a ring buffer to update the necessary components at each step, in order to avoid memory allocation and copying operations every step, which are computationally expensive.

Longformer [ 35 ] introduced Sliding Window Attention, which caches on a block-by-block basis. We will refer to this as Block Sliding Window Attention (BSWA). BSWA does not mask out past keys and values in the ring buffer, but attends to all the information in the ring buffer.

BSWA has two hyperparameters: block size and memory segment, as illustrated in Figure 1(b) and 1(c) . Block size determines how many tokens are in each block and is also used as the stride for sliding. Memory segment determines how many past blocks to cache. In our experiments, we set the default hyperparameters as follows: block size of 1024 1024 and memory segment of 3 3 (corresponding to a window size ranging from 3073 to 4096).

Given a sequence input I = [ i 1 , i 2 , … , i T ] I=[i_{1},i_{2},\dots,i_{T}] to a vanilla Transformer layer, Transformer transforms the input to a sequence output O = [ o 1 , o 2 , … , o T ] O=[o_{1},o_{2},\dots,o_{T}] as follows [ 12 ] :

Q , K , V \displaystyle Q,K,V = QKV ​ ( PreLN ​ ( I ) ) \displaystyle=\text{QKV}(\text{PreLN}(I)) a t \displaystyle a_{t} = SelfAttention ​ ( q t , K , V ) + i t \displaystyle=\text{SelfAttention}(q_{t},K,V)+i_{t} O \displaystyle O = FF ​ ( PreLN ​ ( A ) ) + A \displaystyle=\text{FF}(\text{PreLN}(A))+A

While PreLN [ 36 ] is used in the equation, it is not mandatory. Autoregressive Transformer changes the SelfAttention equation as follows: a t = SelfAttention ( q t , K : t , V : t ) + i t a_{t}=\text{SelfAttention}(q_{t},K_{:t},V_{:t})+i_{t} .

In a Transformer with SWA, the equation is modified to account for a window size of w w : a t = SelfAttention ( q t , K t − w : t , V t − w : t ) + i t a_{t}=\text{SelfAttention}(q_{t},K_{t-w:t},V_{t-w:t})+i_{t} .

In a Transformer with BSWA, τ \tau denotes the block index. Each block τ \tau contains a set of keys and values, determined by the block size. The equation is modified to account for a block index τ \tau and a memory segment m m in Eq. 1 . K τ − m : τ − 1 K_{\tau-m:\tau-1} is the keys in the memory segments from the m m blocks before τ \tau to the block before τ \tau , and K τ , : t K_{\tau,:t} is the keys from the beginning of τ \tau block up to t t . We will refer to Transformer with BSWA as TransformerBSWA.

K ^ t \displaystyle\hat{K}_{t} = Concat ( K τ − m : τ − 1 , K τ , : t ) \displaystyle=\text{Concat}(K_{\tau-m:\tau-1},K_{\tau,:t}) (1a) V ^ t \displaystyle\hat{V}_{t} = Concat ( V τ − m : τ − 1 , V τ , : t ) \displaystyle=\text{Concat}(V_{\tau-m:\tau-1},V_{\tau,:t}) (1b) a t \displaystyle a_{t} = SelfAttention ​ ( q t , K ^ t , V ^ t ) + i t \displaystyle=\text{SelfAttention}(q_{t},\hat{K}_{t},\hat{V}_{t})+i_{t} (1c)

Algorithm 1 presents TransformerBSWA ( Eq. 1 ), re-expressed from the perspective of a block index τ \tau . Algorithm 1 describes how to iteratively calculate a t a_{t} and then concatenate the results into an A τ A_{\tau} sequence. Typically, standard implementations employ a causal attention mask to enable parallel computation of self-attention.

TransformerXL [ 26 ] proposed to use a technique called "stop gradient" for the memory segment. However, we argue that this technique has a negative impact on ability of the model to attend to past information. Specifically, we show that using stop gradient results in a much shorter receptive field than the theoretical receptive field in Section 3.2 . We believe that allowing gradient to flow to the memory segment is necessary for the model to learn to carry important information in the memory segment.

While this modification might seem to burden training memory and computation, it does not significantly impact performance in practice. This is primarily due to the prevalence of gradient checkpointing [ 11 ] in LLM training on ML accelerators, as memory often presents the primary bottleneck. Gradient checkpointing recomputes attention during backpropagation, from later blocks to earlier blocks. Therefore, the presence or absence of stop gradients has little impact on the overall computational complexity, while still improving performance.

In addition, when training an LLM with sliding window attention on long context inputs (more than 8 8 k tokens), computing the attention over the entire input at once would require too much memory. As a result, the standard practice is to divide the attention into blocks and calculate it using a vectorized map (e.g., jax.vmap, torch.vmap). This reduces the peak memory usage to the amount required to calculate one block. Blocks are independent of each other, so they can be calculated in any order. Understanding this point is essential when evaluating the memory requirements and training efficiency of Feedback Attention Memory (FAM), as discussed in Section 2.2 .

Transformer has quadratic memory and computation complexity with respect to the length of the input sequence due to self-attention. It has O ⁡ ( L 2 ) O(L^{2}) complexity for input length L L . However, Transformer with Sliding Window Attention has linear complexity with respect to the input sequence. It has O ⁡ ( L × W ) O(L\times W) complexity for input length L L and window size W W . If the input length is about the same as the window size ( ∼ 1 \sim 1 k), the complexity difference is almost negligible, but if the input length is large like 128 128 k in GPT-4 turbo [ 70 ] , there is a huge difference. In inference, Transformer with SWA or BSWA only needs to cache a fixed ring buffer (block size + + memory segment). Therefore, it only consumes constant memory regardless of the generated token length. Therefore, LLMs using SWA or BSWA can generate infinitely long output tokens.

However, BSWA has a limited receptive field, approximately equal to model depth × \times window size as illustrated in Fig. 1(d) . As a result, the later generated tokens are not related to tokens outside the receptive field (e.g., prompt). To address this limitation, we propose a novel architecture in the following Section 2.2 . Our approach, Feedback Attention Memory (FAM), builds upon BSWA. This is because the block stride concept of BSWA is well-suited for blockwise feedback updates.

### 2.2 Feedback Attention Memory

As mentioned in Section 1 , we hypothesized that attending to the feedback loop can give rise to working memory in Assumption 1 . To implement the feedback loop, we add feedback activations that feed contextual representation back into each block of BSWA. We call these virtual activations as Feedback Attention Memory (FAM). FAM is designed to meet the following key requirements:

• Integrated Attention: Self-attention should simultaneously process input context and FAM.

• Block-Wise Updates: FAM should be updated when transitioning between blocks.

• Information Compression: FAM updates should compress current block information, conditioned on previous FAM.

• Global Contextual Storage: FAM should store comprehensive contextual information indefinitely.

The proposed architecture achieves this by appending FAM to block segments and incorporating it into self-attention processes. This enables richer representations and dynamic propagation of global contextual information across blocks, as illustrated in Fig. 2(b) . When self-attention occurs on the current block, the input query for the block attends to the input key for that block, the memory segment, and the previous FAM. The previous FAM provides global contextual information, allowing for a much richer representation than BSWA. In parallel, the FAM query attends to the current block and the FAM key. The FAM query compresses the current block, conditioned on the previous global contextual information. The FAM query is dynamically generated based on previous global contextual information, as it is copied from the previous FAM. Then, the newly updated FAM serves to propagate global contextual information to the next block recursively. This process is formally described in Algorithm 2 .