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

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

While Algorithm 2 might initially suggest a doubling of matrix operations compared to Algorithm 1 , it performs the same number of matrix operations in the actual implementation, because it starts with the concatenation of block input I τ I_{\tau} and FAM F τ − 1 F_{\tau-1} . The attention mask within self-attention requires a minor modification to accurately represent FAM. The FAM F τ − 1 F_{\tau-1} is much shorter than the input I τ I_{\tau} , and in Section 3 , we experimented with a block size of 1024 1024 and a FAM length of 64 64 .

Transformers are much better at exploiting the parallelism of ML accelerators than Recurrent Neural Networks (RNNs). This is because RNNs have a causal relationship between input sequences, while Transformers only have a causal relationship between the inputs and the layer one depth below. It is possible to worry that the feedback mechanism of TransformerFAM will eliminate the advantages of Transformers and make training inefficient. As explained in the implementation of BSWA, memory-efficient implementations perform self-attention in blocks using vectorized maps. Otherwise, peak memory increases during LLM training, requiring more ML accelerators. The causal relationship of TransformerFAM only exists between blocks. Since vectorized maps are used to perform self-attention in blocks, the causal relationship between blocks does not affect training speed and memory. In addition, processing 64 64 additional FAM when processing 1024 1024 block input sequences has only a minor impact on performance. Therefore, the memory consumption and training speed of TransformerFAM are almost the same as those of TransformerBSWA.

TransformerFAM requires additional considerations for FAM initialization and length extrapolation. These details are explained in Appendix B .

An evaluation of multiple FAM variants was conducted, and the best-performing variant is presented in the main paper. Appendix C provides further details for the remaining variants.

## 3 Experiments

### 3.1 Training

Pretraining an LLM from scratch requires a huge amount of resources. TransformerFAM can reuse existing LLM checkpoints because it does not add new weights to the Transformer layer. We reused 1B, 8B, and 24B Flan-PaLM LLMs [ 58 ] for our experiments. This is a large enough size to prove that TransformerFAM is a general solution for LLMs. The model sizes 1B, 8B, and 24B refer to the size of the plain Transformer, excluding the text embedding table. The models use a 256k sentence piece tokenizer [ 20 ] , resulting in 400M text embedding table weights for the 1B model and 1B weights for the 8B and 24B models. The detailed model architecture is described in Table 3 in Section A.1 .

Flan-PaLM is a model that is fine-tuned on top of a pretrained PaLM model [ 75 ] using instruction finetuning. The instruction data consists of few-shot instructions with 100 100 to 1 1 k tokens, which are packed into 2.5 2.5 k tokens for training. This means that individual instruction data are concatenated until they reach 2.5 2.5 k tokens.

We applied both the TransformerBSWA and TransformerFAM architectures to Flan-PaLM and fine-tuned it for an additional 50k steps. We experimented with different memory segment for both architectures. The block size is set to 1024 1024 and the FAM length is set to 64 64 .

During fine-tuning, we used the same Flan instruction data packed into 8.5 8.5 k tokens. To maintain a minibatch size of 128 for all models, we used 32 TPUv5 [ 77 ] cores for the 1B model, 64 cores for the 8B model, and 128 cores for the 24B model. If we had used more resources and a larger minibatch size, we might have achieved better results than those reported in the paper.

We performed LoRA finetuning by adding LoRA [ 46 ] to the Attention and FF layers of Transformer without training all the parameters. Full finetuning resulted in lower scores on various tasks reported by GPT-3 [ 33 ] , because catastrophic forgetting [ 14 ] occurred in domains that were not covered by the instruction data. In LoRA finetuning, the scores on GPT-3 tasks actually improved, and the performance on long context tasks was similar to that of full finetuning. The rank of LoRA was 64, and the weights of the original Attention and FF layers were merged with LoRA weights and used during inference.

The Adafactor optimizer ( β 1 = 0.9 , β 2 = 0.99 \beta_{1}=0.9,\beta_{2}=0.99 ) [ 15 ] was used with constant learning rate. The learning rates used were 10 − 4 10^{-4} for 1B, and 3 × 10 − 3 3\times 10^{-3} for both 8B and 24B.

In addition, TransformerXL exhibits comparable performance to TransformerBSWA. The implementations are almost identical, with TransformerXL employing an additional QK attention mask to mask out keys beyond a predetermined window size. Section B.4.7 demonstrates that the performance difference between TransformerXL and TransformerBSWA is insignificant, and therefore, experimental results for TransformerBSWA are only included in the main paper .

#### 3.1.1 Data

The ideal data for training TransformerFAM is a very long document with continuous context, such as textbooks and novels, and the data should be large enough to finetune an LLM. Additionally, the same very long document should be used continuously in each minibatch component, while maintaining FAM and memory segments between training steps.

The loss function of an LLM is to minimize the difference between the parametric probabilistic model P ⁡ ( x | θ ) P(x|\theta) and the data-generating true distribution p d ​ a ​ t ​ a ​ ( x ) p_{data}(x) . The Kullback–Leibler divergence (KL divergence) is used to measure this difference. To perform an unbiased estimation of KL divergence, we draw samples from the data-generating true distribution, which are assumed to be independent and identically distributed (IID). However, the ideal training scenario for the aforementioned memory training directly contradicts the IID assumption. We refer to this as the curse of IID.

Due to the curse of IID, we could not find the training infrastructure or data suitable for training memory. So we used Flan instruction data as a last resort.

We used Flan instruction data [ 84 ] as training data, packed up to 8.5k tokens. In the Flan paper [ 44 , 58 ] , a special attention mask was used to prevent attention between different segments during self-attention, by applying a separate mask to each packed segment. We did not use this special attention mask processing. Attention occurs causally within the window, regardless of the segment. We expected TransformerBSWA and TransformerFAM to learn to remember important information and forget outdated information by themselves.

Each token in Flan data has a weight. Few shots examples and instructions have a weight of 0, and answers have a weight of 1. This means that the model only learns to generate the answer. To incentivize the model to remember long contexts, we randomly selected 256 consecutive tokens from 8.5k tokens and appended them to the end of the data with the prompt ’ [repeat random segment]: ’. The repeated tokens were given a weight of 0.1 0.1 . We hope that future studies can use more suitable data for training memory, such as long continuous documents, long-form speech, video or video game.

### 3.2 PassKey Retrieval

The PassKey retrieval task is a recent benchmark used in several long-context transformer papers [ 71 , 68 , 67 ] . In this task, a passkey is presented at the beginning, followed by a very long filler context. Finally, a question about the passkey is asked, as shown in Fig. 10 in Section D.1 .

This task is a good smoke test to quickly check if information is transmitted in a long context. However, this task only checks if small and important information is transmitted, and does not check if large amounts of information can be efficiently compressed.

We fine-tuned the Flan-PaLM 1B model for 5k steps with the PassKey format, which has a filler context of 2k to 18k randomly. We used a block size of 1024, TransformerBSWA with 0 to 12 memory segments, TransformerFam with 0 memory segments, and a FAM length of 64 64 . When the number of memory segments is 3, the window size is 3k (i.e. memory segment × \times block size).

As shown in Fig. 3(a) , TransformerFAM was able to perfectly solve the task with a filler context of up to 260k tokens . In the figure, MX denotes the number of BSWA memory segments. The performance of TransformerBSWA improves significantly up to M2, after which it saturates. The performance of M12 also drops significantly after 20k tokens. The theoretical receptive field of M2 is 36k (i.e. depth(18) × \times memory segment(2k)), but the effective receptive field is much shorter.

In Fig. 3(a) , it is important to compare M1_SG and M1. M1_SG has a stop gradient applied to one memory segment, which limits the receptive field to the window size. This is because the model cannot learn which contextual information stored in the memory segment will have a good result later. SWA with back propagation through time (BPTT) functions similarly to a time-limited RNN. It is common to use stop gradients on the memory segments of TransformerXL [ 26 , 63 ] . However, we recommend against this practice.

In Section B.4.8 , we compare our work with recent Transformer with memory papers [ 57 , 63 ] .

### 3.3 Long Context Tasks

Gemini [ 69 ] evaluated long-context capabilities using the following tasks: NarrativeQA [ 19 ] , Scrolls-Qasper, Scrolls-Quality [ 60 ] , and XLSum [ 51 ] . Additionally, PG-19 [ 21 ] and Isabelle [ 56 ] are another common evaluation tasks among long-context Transformer papers [ 21 , 56 , 67 ] . Detailed information on the evaluation data is provided in Table 10 in Section D.2 .

We evaluated the long-context capabilities of the 1B TransformerBSWA model trained in Section 3.1 using memory segment sizes ranging from 0 to 8. As shown in Fig. 3(b) , TransformerFAM outperformed TransformerBSWA on all the long context tasks (LCT), regardless of the number of memory segments in BSWA. It shows a significant performance improvement on ScrollsQasper and NarrativeQA, where it has to understand 5k to 500k tokens of context before answering a question. The LCT results demonstrate that TransformerFAM can effectively compress and retain important contextual information within extremely long contexts .

Above M1, the number of memory segments does not significantly impact LCT performance on TransformerBSWA, because the input sequences are much longer than the window size of all experiments. We observed the same phenomenon in TransformerFAM, and TransformerFAM uses 3 3 memory segments in Fig. 3(b) . The figure shows the normalized scores of all tasks to view the scores on the same scale. The raw results are in Table 11 in Section D.2 .

We further evaluated TransformerFAM and TransformerBSWA on 8B and 24B models. As shown in Table 1 , TransformerFAM demonstrates scalability as the model size increases. This suggests that self-attention can route local information relevant to each input sequence while simultaneously routing contextual information to the FAM. However, the performance improvements are not substantial, indicating room for further enhancements in working memory mechanisms.

In addition, TransformerFAM marginally surpasses TransformerBSWA on GPT-3 tasks [ 33 ] (see Table 2 ). This result is unexpected since all tasks involve sequences shorter than 2k tokens. We hypothesize that this improvement arises from the efficient contextual representation by TransformerFAM. By offloading contextual data to FAM, TransformerFAM reduces redundancy within input activations, optimizing latent space usage . 1 1 1 It can be viewed as the decoder analog of register tokens [ 78 ] in ViT encoders [ 32 ] to process global context.

Thus, BSWA memory segments (local representation) and FAM (global representation) complement each other. For LLMs, we recommend using FAM for compressed contextual representation alongside BSWA memory segments up to inference budgets (e.g., 2k, 8k, 32k [ 69 ] , or 128k [ 70 ] ). Due to page limitations in the main paper, ablation studies are presented in Section B.4 .

## 4 Related Work

There have been attempts to incorporate feedback mechanisms into the Transformer, but most of them involve feeding the output activations from the top layer to the bottom [ 57 , 63 ] or to intermediate layers [ 27 ] . Since the top three layers in the Transformer are heavily focused on output reconstruction [ 45 ] , we hypothesize that there is a significant representational gap between the top and other layers. In this paper, we propose a feedback mechanism between intermediate layers.

There were papers that compressed information blockwise [ 21 , 25 , 31 , 71 , 72 ] . However, in those papers, the information was not propagated infinitely. Relevant prior work includes the use of recurrent cross-attention between blocks [ 52 ] , enabling the propagation of compressed information to subsequent blocks. Additionally, incorporating feedback from a few upper layers has been used to integrate past information [ 61 ] . We propose TransformerFAM under the assumption that the human brain processes homogenous, heterogeneous, and feedback data with the same attention mechanism across distributed brain areas. Additional related works are presented in Appendix E .

## 5 Conclusion

In the film ’Memento’ (2000), the protagonist struggles with anterograde amnesia, which means he can not remember anything before happened in the last 10 minutes, but his long-term memory is intact, He has to tattoo important information on his body to remember it. This is similar to the current state of large language models (LLMs). LLMs memorize the entire internet thanks to scaling laws [ 39 ] , which allow them to store an enormous amount of information in large weights (long-term memory). However, their short-term memory is limited by the attention window. As a result, the complex prompt engineering becomes necessary to help them recall important details. We propose a new architecture called TransformerFAM that could fix anterograde amnesia of LLMs.

The rapid progress of machine learning is astonishing, but there are two key problems that we still do not know how to approach: reasoning and memory. In this paper, we provide a clue to the memory problem. Memory is a critical prerequisite for reasoning. It is hard to imagine how we can derive complex mathematical equations without working memory. Reasoning must be a phenomenon that occurs based on the current working memory.

This paper explores the integration of attention-based working memory, a concept from neuroscience, into the field of deep learning. Our goal is to ignite further research within the community to address and solve the ongoing challenge of limited memory in deep learning. There is a significant set of problems to tackle here, ranging from refining feedback attention architecture to investigating the transfer of working memory to long-term memory.

## Appendix A Architecture details

### A.1 Flan-PaLM architecture

Table 3 provides detailed information on the architecture of Flan-PaLM for 1B, 8B, and 24B models. MQA (Multi-Query Attention) [ 22 ] is an attention mechanism that employs a single set of keys and values for all attention heads.

### A.2 FAM hyperparameters

Table 4 presents the default settings for the hyperparameters added in TransformerFAM.

## Appendix B Additional details of TransformerFAM

The appendix describes additional details not covered in the main text.

### B.1 FAM initialization

The current FAM F τ F_{\tau} is the output activation of the previous FAM update. So, how do we initialize the FAM for the first block?