TransformerFAM: Feedback attention is working memory — part 4/4
Cleaned-text code points [47584, 58342). Bibliography omitted here; source.md includes it.
However, as shown in Fig. 9 , FM fails to retain PassKey information over a long context. The figure compares M1, 2, and 4 with FM1, 2, and 4 when the memory segment size is 1, 2, and 4 for BSWA and Feedback Memory Segment. TransformerFM outperforms TransformerBSWA, but still falls short of TransformerFAM by a significant margin. In TransformerFM, each activation must possess both local and global representation, similar to TransformerBSWA. However, the absence of an activation specifically responsible for the global representation appears to prevent the retention of critical information like PassKey over an extended context.
Furthermore, TransformerFM exhibited lower performance than TransformerBSWA on GPT-3 tasks ( Table 9 ). As TransformerFM introduced additional complexity without any discernible performance advantages, we discontinued this architecture. Table 9 presents the results after 25k steps of fine-tuning.
### C.3 Static Summary Tokens
As illustrated in Fig. 2(b) , "FAM copy" duplicates the previously updated FAM to the current FAM input. However, "FAM copy" is not an essential component in the feedback loop. When FAM compresses the current block, it can also attend to the previous FAM as both key and value, as shown in the experiment in Section B.4.2 . Therefore, "FAM copy" is not strictly necessary because past FAM information is propagated through self-attention.
As mentioned in Section B.1 , we learn the initial FAM using prompt tuning, and this token can be considered a "Summary" token. An alternative design could employ a persistent summary token for block compression, with the feedback loop enabled by attending to the key-value pairs of previous FAMs.
As illustrated in Fig. 9 , SUM FAM (static summary token) prevents the model from successfully solving the PassKey task. This suggests that the query for summarization needs to be dynamically generated conditioned on the past FAM. A static query may not be able to cover all situations effectively. Additionally, simply attending to the key and value alone does not transmit sufficient information from the past. On the other hand, the "FAM copy" mechanism propagates information from the past FAM to the current FAM through the residual connection of the Transformer, which facilitates better performance on the PassKey task.
### C.4 Diversity Loss
In Section B.4.1 , we observed a performance drop when the FAM length exceeded 64. We hypothesized that FAM was underutilizing its capacity and experiencing mode collapse, where lengthy FAM either focused on a small number of inputs or prevented all inputs from obtaining sufficient information.
As a remedy, we employed diversity loss as an auxiliary loss, similar to Wav2vec 2.0 [ 41 ] . Diversity loss aims to maximize the entropy of QK attention probability, as shown in Eq. 2b . This loss encourages FAM to uniformly attend to all inputs, and all inputs equally attend to FAM. In Eq. 2 , b b denotes the batch size, τ \tau represents the block index, l l indicates the sequence position within a block, and h h refers to multi heads. In Eq. 2a , p ¯ b τ \bar{p}_{b\tau} represents the average attention probability across all sequences and multi-heads for the self-attention of each block. The auxiliary loss functions to regularize this probability towards uniformity.
p ¯ b τ \displaystyle\bar{p}_{b\tau} = 1 H L ∑ h = 1 H ∑ l = 1 L p b τ h l \displaystyle=\frac{1}{HL}\sum_{h=1}^{H}\sum_{l=1}^{L}p_{b\tau hl} (2a) ℒ d \displaystyle\mathcal{L}_{d} = 1 B T ∑ b = 1 B ∑ τ = 1 T − H ( p ¯ b τ ) = 1 B T ∑ b = 1 B ∑ τ = 1 T p ¯ b τ log p ¯ b τ \displaystyle=\frac{1}{B\mathrm{T}}\sum_{b=1}^{B}\sum_{\tau=1}^{\mathrm{T}}-H(\bar{p}_{b\tau})=\frac{1}{B\mathrm{T}}\sum_{b=1}^{B}\sum_{\tau=1}^{\mathrm{T}}\bar{p}_{b\tau}\log{\bar{p}_{b\tau}} (2b)
We trained the model with diversity loss with various weights, but the overall performance was always worse, regardless of the FAM length. It did not help at all even when the FAM length was 256 or longer.
### C.5 Reconstruction Loss
Compressive transformers [ 21 ] compress memory segments before forwarding them to the next block. The key difference between our approach is that the compressed information is not recurrently connected. The paper proposes a reconstruction loss as an auxiliary loss, which aims to reconstruct the original activations from the compressed activations.
We also experimented with the reconstruction loss. In details, we generated transformer outputs autoregressively from the updated FAM and the query sequence of the original block, and compared them to the original outputs using the MSE loss. However, this did not help to improve the performance.
## Appendix D Experiments details
This section provides additional details about the experiments that were not covered in Section 3 .
### D.1 PassKey Retrieval
The format of the PassKey retrieval task is shown in Fig. 10 . The original paper [ 71 ] also included a prefix filler, but we removed it in our paper.
There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there. The pass key is <PASS KEY> . Remember it. <PASS KEY> is the pass key. <filler> What is the pass key? The pass key is <PASS KEY>
### D.2 Long Context Tasks
Table 10 provides detailed information about the long context tasks we used. Due to the limitations of TPU memory, PG-19, Isabelle, and NarrativeQA were truncated to 256k tokens, which is within two standard deviations.
Table 11 presents the LCT results for TransformerBSWA and TransformerFAM. MX represents the number of memory segments for TransformerBSWA.
### D.3 GPT-3 Tasks
We evaluated all model sizes on the tasks reported by GPT-3 [ 33 ] . The results are shown in Table 12 .
### D.4 Complexity
BSWA and FAM have memory and computational complexity of O ( L × C ) O(L\times C) , where C C is the chunk size. Table 13 and Table 14 show memory and inference time through inference jobs on TPUv4. The most tokens in the experiment are pre-filled and generation is limited to 256 tokens.
## Appendix E Related Work
The Transformer architecture exhibits a quadratic complexity with respect to context length, a significant limitation. To address this, several research works have focused on approximating the attention mechanism. One approach involves sparse attention, where only a subset of important tokens are attended to, as seen in models like Sparse Transformer [ 23 ] , Big Bird [ 28 ] , Reformer [ 37 ] , Routing Transformer [ 42 ] , and TOVA [ 81 ] . Linear approximation methods offer an alternative, seeking to circumvent the quadratic complexity by altering attention calculations, as exemplified by Linformer [ 34 ] , Linear Transformer [ 40 ] , Performer [ 38 ] , and Nyströmformer [ 48 ] . Finally, some research explores entirely different sequence-to-sequence architectures as replacements for attention-based Transformers, including MLP-mixer [ 47 ] , State Space Models [ 18 ] , S4 [ 43 ] , Mamba [ 73 ] , and RWKV [ 76 ] .
Global Workspace Theory (GWT) [ 8 ] is a leading theory of consciousness. According to GWT, the human brain possesses a global workspace where various modules, such as sensory input, memories, and internal representations, converge. The attention mechanism in brain acts as a spotlight, focusing on specific inputs among the multitude, and transforming this unconscious activity into conscious awareness. These "momentarily active, subjectively experienced" events are then stored in working memory. TransformerFAM draws inspiration from GWT, adopting its principle of a unified attention mechanism for processing homogenous, heterogeneous, and feedback data.
## Appendix F Attention Visualization
Fig. 11 depicts the attention map for each head in each layer of a 1B model. FAM is prepended and is located at the bottom left corner. The bright spots along the left edge represent the block inputs attending to FAM, while the bright spots along the bottom edge represent FAM compressing the corresponding block. Overall, the block inputs actively reference FAM, while FAM compresses only the selective inputs.
## Appendix G Limitations
While the results presented in Table 1 demonstrate that TransformerFAM shows improvements on long-context tasks, these gains are not yet substantial, highlighting the need for further development and refinement of working memory mechanisms.
In this paper, we have taken an initial step towards integrating attention-based working memory, a concept inspired by neuroscience, into deep learning architectures. We believe there is significant potential for further exploration in this direction, and we encourage future research to continue addressing the ongoing challenge of limited memory in deep learning models.
## Appendix H Broader Impacts
While our work is inspired by the concept of working memory from neuroscience, as discussed in Section Section 3.3 , achieving a human-level implementation remains a significant challenge. This research represents an initial step in that direction.
The potential societal impacts of such advanced working memory in LLMs could be substantial, with applications like highly personalized AI assistants. However, these impacts are currently speculative due to the nascent stage of this research.
In the immediate future, the primary benefit of our work is expected to be improvements in the efficiency and effectiveness of LLMs, with potential applications across various domains such as education, healthcare, and communication. As with any technology, we recognize the possibility of misuse and encourage ongoing research into the ethical implications of LLMs and related advancements in artificial intelligence.
## Instructions for reporting errors
We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:
Click the "Report Issue" ( ) button, located in the page header.
Tip: You can select the relevant text first, to include it in your report.
Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.
Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .