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

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

When FAM are used as queries, they summarize the block context. Therefore, we learn by adding learnable summarization embeddings to the token embedding lookup level of the Transformer model. This is the same as prepending learnable embeddings in soft prompt tuning [ 49 ] . The difference is that full attention is applied between the FAM prompt activations, and the updated FAM is used for the next block. The FAM prompt is passed to the next transformer layer through the forward propagation of the Transformer model, and it has a summary representation that is suitable for the level of the layer.

Prefix tuning [ 50 ] can also be used to train learnable initial FAM at each Transformer layer. However, we had difficulty matching the magnitude of the learnable prefix to the regular input sequence, and the results of prompt tuning were consistently better. Ablation study in Section B.4.6 shows prompt tuning outperforms. In addition, prefix tuning has the disadvantage of adding additional weights of FAM length to each layer.

In addition, the first FAM update in self-attention should utilize a zero tensor in the residual connection rather than the initial FAM, because the initial FAM does not carry any contextual information.

In summary, we learned the initial FAM using prompt tuning, which only adds a very small number of weights of FAM length to the entire model. We share the same initial FAM across all minibatches.

### B.2 Input Length Extrapolation

#### B.2.1 FAM Position Encoding

We used rotary position embedding (RoPE) [ 80 ] in all of our experiments. Each input sequence is assigned an integer position m m , which is converted into sinusoidal encoding in the form of exp ⁡ ( i ​ m ​ θ ) \exp(im\theta) where θ = 10000 − 2 i / d m ​ o ​ d ​ e ​ l \theta=10000^{-2i/d_{model}} [ 12 ] .

FAM is inserted at each block boundary, but the problem is how to assign positions to FAM. We assigned positions to FAM in order from the last number of the compressed block. For example, if the block size is 4, the FAM length is 2, and the positions of the compressed blocks are m m , m + 1 m+1 , m + 2 m+2 , and m + 3 m+3 , then the updated FAM positions are m + 2 m+2 and m + 3 m+3 .

We tried other methods, but this method worked best according to the ablation results in Section B.4.3 .

#### B.2.2 Random Position Offset

The input length extrapolation problem of Transformer is well known [ 67 , 65 ] . For example, a Transformer LLM trained with 2k tokens experiences a severe performance drop when generating 8k tokens. This is because machine learning (ML) does not generalize well to situations that it has not seen during training.

Transformer with SWA does not suffer from the position extrapolation problem when using relative positional embedding like RoPE. This is because the score value of q m ⋅ k n q_{m}\cdot k_{n} becomes a function of ( m − n ) (m-n) in the form of exp ⁡ ( i ⁡ ( m − n ) ​ θ ) \exp(i(m-n)\theta) .

Because the range of ( m − n ) (m-n) is limited to the window size, independent of the input length, the model can handle long input sequences without facing novel scenarios during inference. The model can accurately determine the relative position, if the window size is smaller than the maximum wavelength of θ \theta . If the typical θ \theta is used for positional embedding, the working maximum window size is the maximum wavelength ( = 2 × 10000 × π ∼ 63 =2\times 10000\times\pi\sim 63 k tokens ).

However, FAM breaks the symmetry of relative position. Since the absolute position from the past to the present is recursively embedded in the FAM, the large absolute position value that the model encounters for the first time during inference creates a situation where the model needs to extrapolate.

We propose Random Position Offset as a solution. At each training step, the Transformer model randomly samples a scalar value between 0 and the maximum wavelength. All Transformer layers add that random value to the absolute position at that training step. Therefore, the FAM experiences the entire range of absolute position embedding during training.

This is a purely training technique. During inference, the default offset is 0 0 . We used the below algorithm that generates 0 0 by 50 % 50\% when sampling the offset, as 0 0 is the default value.

⬇

#### B.2.3 Random State Passing

Due to the recursive use of FAM, we need to determine the maximum number of updates for which the FAM remains valid. If it is updated up to 8 times during training, the model will have to extrapolate the situation where it is updated 100 times during inference.

The same problem existed in RNNs, and Random State Passing (RSP) [ 24 ] was proposed to handle long-form speech in the speech domain. We also used RSP to generalize the number of updates of FAM.

RSP saves FAM as weights at the end of each training step. Then, it loads the saved FAM at the next training step. When initializing FAM, it either uses randomly saved FAM or learned FAM. In our default setup, it used saved FAM with 80 % 80\% probability. To save FAM of all minibatch, weights are required as many as the number of minibatch. We save FAM of only the first batch and all minibatch share it in the next training step.

On the other hand, saved FAM can be thought of as learnable prefix for prefix tuning [ 50 ] . It is also possible to train only the FAM while the model is frozen and use them for various downstream tasks or personalization. This part is left as a future research topic.

### B.3 FAM illustrated

In Section 2.2 , we define the feedback loop as feedback activations that feed contextual representation back into each block of BSWA. This feedback loop is formally described in Algorithm 2 . While Fig. 2(b) illustrates Algorithm 2 , the static image makes it challenging to fully grasp the dynamic nature of the decoding self-attention mechanism. To clarify this, we create a multi-frame animation in Fig. 4 that demonstrates how the attention mechanism evolves over time (top to bottom).

### B.4 Ablation Studies

All ablation studies were conducted on a 1B model.

#### B.4.1 FAM length

In the Flan 1B model, we observed performance saturation on ScrollsQasper, ScrollsQuality, and XLSum tasks when the FAM length reached 64. Interestingly, performance declined when FAM length exceeded 64, suggesting that information compression is more effective with limited space. This constraint on memory capacity is reminiscent of Miller’s Law [ 1 ] , which posits that the average person can only hold approximately 7 ​ ( ± 2 ) 7(\pm 2) items in their working memory at any given time.

#### B.4.2 The number of previous FAM blocks

In Fig. 2(b) , the input query attends to the FAM as denoted Attention to Feedback. The input query can attend to not only the immediately previous FAM, but also to more previous FAMs. Table 6 shows the XLSum scores for different numbers of previous FAM blocks. As the table shows, increasing the number of blocks did not have a significant effect, because the previous FAM already encapsulates all the previous information by a feedback copy. Therefore, the default setup attends to only the immediately previous FAM.

#### B.4.3 FAM Position Encoding

Section B.2.1 proposed assigning the last number of the compressed block as the FAM position. In addition, we also experimented with FAM having a float position between blocks. In the example of block size 4 in Section B.2.1 , the FAM positions would be m + 3 + 0.33 m+3+0.33 and m + 3 + 0.66 m+3+0.66 . We also experimented with the case where the FAM position is always 0.

As shown in Table 7 , the last number showed the best accuracy.

Fig. 5 shows FAM (i.e. Last number in Table 7 ) outperforming FAM-POS (i.e. Zero) in PG-19 accuracy over most base frequencies.

#### B.4.4 Random Position Offset

As mentioned in Section B.2.2 , the input length extrapolation problem of the Transformer is well-known. The "Attention is All You Need" [ 12 ] introduced sinusoidal position encoding in the form of exp ⁡ ( i ⁡ ( m − n ) ​ θ ) \exp(i(m-n)\theta) where θ = 10000 − 2 i / d m ​ o ​ d ​ e ​ l \theta=10000^{-2i/d_{model}} . Popular solutions for full attention models include increasing the base frequency from 10k to 500k [ 67 ] or scaling down ( m − n ) (m-n) to ( m − n ) × ( 1024 / 4096 ) (m-n)\times(1024/4096) [ 65 ] .

Since the relative position of BSWA has a range of 0 to window size, it does not suffer from the input length extrapolation problem. However, TransformerFAM encodes absolute position into FAM, which requires a solution. Section B.2.2 proposes Random Position Offset (RPO) as a solution.

In Fig. 5 , FAM shows better PG-19 accuracy than FAM-RPO at the 10k base frequency. As mentioned in Section D.2 , the max length of PG-19 was truncated to 256k in the experiments.

Furthermore, scaling up the base frequency or scaling down ( m − n ) (m-n) is only a remedy for pain, not a final solution. It reduces the resolution of the position, which negatively affects the overall Transformer performance, as shown in the figure. Interestingly, the originally proposed 10k is a very good value even in long contexts, until the window size reaches its wavelength (63k).

In addition, we did not observe the attention sink phenomenon [ 66 ] in our TransformerBSWA experiments. The paper proposes that Transformers with sliding window attention should store initial tokens in a KV cache for long-context inference. However, the relative position of SWA is restricted to a range of 0 to window size, independent of input sequence length. Our TransformerBSWA implementation, trained on 8.5k tokens, successfully operated for inference up to 256k tokens without any issues.

#### B.4.5 Random State Passing

To extrapolate the number of FAM updates, Section B.2.3 proposes Random State Passing (RSP). In Fig. 5 , FAM shows significantly better PG-19 accuracy than FAM-RSP. This demonstrates that RSP plays a crucial role in training FAM.

Fig. 6 shows the best performance at a probability of 0.8. At 0.8, the half-life is 3 training steps ( 0.512 = 0.8 3 0.512=0.8^{3} ). This means that every 26k (8.5k x 3) tokens, FAM restarts from the beginning with a 50 % \% probability. In other words, if FAM experiences 25 FAM updates during training, it can extrapolate to 256 FAM updates (256k tokens) during inference.

#### B.4.6 Prompt vs. Prefix

Section B.1 mentions that Prompt tuning outperforms Prefix tuning for training FAM in terms of performance, memory, and number of parameters. Fig. 5 shows that Prefix tuning for FAM training leads to performance degradation in PG-19 (FAM vs. FAM+Prefix).

#### B.4.7 TransformerXL vs. TransformerBSWA

TransformerXL is a modification of TransformerBSWA that incorporates an additional QK attention mask to mask out keys beyond a specified window size. Therefore, there is no reason to expect a significant difference in performance between the two. Table 8 shows the scores for the main tasks. TransformerBSWA used a single memory segment, and TransformerXL window size was equal to the block size of 1024 1024 .

#### B.4.8 Comparison with other methods

We compared FAM with Recurrent Memory Transformer (RMT) [ 57 ] , because RMT also implements memory in Transformer using feedback mechanism from top layer to bottom layer. As shown in Fig. 7 , RMT showed worse performance than Block Sliding Window Attention (BSWA) with 1 memory segment, in the PassKey retrieval task. FAM solved the PassKey retrieval task, but RMT did not work at all with very long filler token lengths.

RMT is implemented by feeding the output memory of the previous segment as the input memory to the next segment. In the constraint that the input is text embedding and the output is text reconstruction in LLM, RMT has an additional constraint that the latent space of the output memory and the input memory must match. In this situation, RMT fails to remember the PassKey for a very long context. On the other hand, FAM seems to compress, store, and propagate information more effectively by learning memory representation that matches the abstraction level of each layer through training.

AutoCompressors [ 63 ] is an extension of RMT that continuously accumulates blockwise memory of RMT. AutoCompressors theoretically should be able to solve the PassKey task since it maintains all the memory tokens for all blocks. However, as shown in Fig. 7 , its performance drops sharply after 18k tokens. This is because the model only saw up to 18k tokens during training. It fails to generalize to longer filler token lengths. The AutoCompressors in Fig. 7 accumulates 260 memories to support up to 260k tokens.

## Appendix C Don’t

There could be various ways to create feedback loops in Transformers. This appendix summarizes our attempts that did not work well. We hope that this will save other researchers time when improving the architecture of feedback loops in future studies.

### C.1 Beyond Block-Recurrent Transformers

Block-Recurrent Transformers (BRT) [ 52 ] have a recurrent state that connects each block, and the recurrent state and the block input sequence exchange information through cross-attention. In this architecture, the recurrent state plays a role similar to FAM. We started the project by removing the additional QKV projection and cross-attention for the recurrent state in BRT and integrating all computations into self-attention.

Like BRT, we tried to compress the block input by attention to use it as the recurrent state of the next block, and it required additional projections like BRT. It was difficult to properly train the additional projections that were only used when striding blocks, and as a result, the performance of general tasks such as GPT-3 tasks was degraded. Finally, we found that the activations compressed by attention must go through the FF layer of the Transformer to be aligned as a latent representation that can be used again as the input of the next Transformer. The input of the Transformer and the output of self-attention are very different representations. However, the FF layer transforms it back into a representation similar to the input of the Transformer for the next Transformer layer. After the discovery of reusing the FF layer, we also found that separate QKV projection and additional projections are not required for FAM. That is how the TransformerFAM architecture was created.

Around the time this paper was published, Infini-Transformer [ 79 ] was also released. Infini-Transformer reported successful PassKey retrieval despite the compressed information not passing through the FF layer. This is because it stores weighted values in scratchpad memory and reuses them, preserving the latent space of the values. However, since the weighted values are all it can represent, there might be limitations in its expressiveness.

Furthermore, TransformerFAM maintains the past memory segment from BSWA despite having compressed memory because the compressed memory cannot retain detailed information. Since Infini-Transformer completely discards the past memory segment, it might have difficulty remembering recent details.

### C.2 Feedback Memory Segment

As shown in Fig. 2(b) , TransformerFAM utilizes the updated FAM from the previous block as the input to the current block. Therefore, it is natural to consider using the Transformer outputs of the previous block directly as the input to the current block, instead of using a complex FAM mechanism. Fig. 8 illustrates this modification, which we refer to as the Feedback Memory Segment (FM). FM achieves an infinite theoretical receptive field by performing self-attention at the output level. This is similar to ERNIE-Doc [ 30 ] , and a specific variant of Staircase Attention [ 61 ] . Staircase Attention proposes using activations from progressively higher Transformer layers as the memory segment goes further into the past.