Addressing Some Limitations of Transformers with Feedback Memory

Angela Fan; Thibaut Lavril; Edouard Grave; Armand Joulin; Sainbayar Sukhbaatar

Source-linked AI summary

Addressing Some Limitations of Transformers with Feedback Memory

Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, Sainbayar Sukhbaatar

arXiv:2002.09402v3Feb 21, 2020cs.LGcs.CLstat.ML

TL;DR

Transformers’ parallel processing limits access to already-computed higher-level past representations and constrains recursive state updates. The paper introduces Feedback Transformer, which shares a memory of past layer representations across future computation, and reports stronger performance from smaller, shallower models across language modeling, machine translation, and reinforcement learning.

Problem
Transformers cannot use the highest-level representations of the past to compute current representations, while their recursive computation is bounded by model depth.
Method
Feedback Transformer merges hidden states from all layers into a shared memory that all future layers can access for recursive computation.
Results
The architecture achieves strong performance across language modeling, machine translation, and reinforcement learning with smaller, shallower models, faster decoding, and lower memory use.
Takeaways & Limitations
Feedback memory increases representation capacity by making higher-level past computations available throughout future processing.
Takeaways & Limitations
Transformers’ internal state is limited for recursive computation, whereas Feedback Transformer memory can grow with input length and maintain more information.

Abstract

from arXiv · show

Transformers have been successfully applied to sequential, auto-regressive tasks despite being feedforward networks. Unlike recurrent neural networks, Transformers use attention to capture temporal relations while processing input tokens in parallel. While this parallelization makes them computationally efficient, it restricts the model from fully exploiting the sequential nature of the input. The representation at a given layer can only access representations from lower layers, rather than the higher level representations already available. In this work, we propose the Feedback Transformer architecture that exposes all previous representations to all future representations, meaning the lowest representation of the current timestep is formed from the highest-level abstract representation of the past. We demonstrate on a variety of benchmarks in language modeling, machine translation, and reinforcement learning that the increased representation capacity can create small, shallow models with much stronger performance than comparable Transformers.

1 INTRODUCTION

The paper identifies limited access to higher-level past representations and bounded recursive computation as key Transformer limitations, then introduces Feedback Transformer to address them. The architecture uses shared feedback memory to expose past computations broadly while reducing memory use and accelerating decoding.

1 INTRODUCTION: Transformers restrict sequential modeling because parallel layer processing prevents current representations from using already-computed higher-level representations of the past.Their recursive computation is also bounded by model depth, limiting tasks involving state tracking or hierarchical structure.
1 INTRODUCTION: Feedback Transformer makes all previous hidden representations accessible to computation at any depth by feeding past computations back into the model.It modifies self-attention to attend to higher-level representations rather than only lower-level ones.
1 INTRODUCTION: The architecture merges hidden states from all layers into one vector per timestep, stores these vectors in memory, and lets subsequent layers attend to that memory.This enables recursive computation and repeated transformation of an input across as many steps as its length.
1 INTRODUCTION: Shared feedback memory reduces memory size by L times, shares key and value projections, and makes the overall model 2x smaller.The reduced GPU memory enables larger batches, which contribute to substantially faster decoding during inference.
1 INTRODUCTION: The paper reports that Feedback Transformers achieve state of the art results with smaller, shallower models that decode faster and use less memory.These claims are presented as the work’s main contributions across the evaluated tasks.

2 RELATED WORK

Related work frames Feedback Transformer alongside analyses of Transformer limitations and architectures that add recurrence or external memory. Figure 1 depicts its distinctive shared-memory mechanism for combining past representations.

2 RELATED WORK: Prior analyses identify Transformer limitations in sequential processing, hierarchical structure, bounded recursion, and unbounded decoding-step settings.The cited works include Dehghani et al. (2018), Tran et al. (2018), Hahn (2020), and Pérez et al. (2019).
2 RELATED WORK: Feedback Transformer is architecturally related to recurrent networks with external shared memories, including stack-augmented RNNs and Neural Turing Machines.These comparisons connect its memory design to mechanisms for retaining long-term dependencies and modeling unconstrained memory.
2 RELATED WORK: Other Transformer modifications incorporate recurrent structure, including local recurrent sublayers and parameter sharing that creates recurrence in depth.Examples include Wang et al. (2019) and Universal Transformer (Dehghani et al., 2018).
2 RELATED WORK: Figure 1 shows the Feedback Transformer merging past hidden representations from all layers into a single vector stored in memory.This visual summarizes the architecture’s shared-memory distinction from ordinary layerwise access.

3 METHOD

The Feedback Transformer replaces layer-local parallel processing with sequential feedback memory, exposing higher-level past representations to all future computations. This supports recursive state updates while reducing memory and enabling shallower, faster models.

3.3 FEEDBACK TRANSFORMER: Memory vectors merge representations from all layers at each timestep, with learnable softmax weights controlling whether layers are averaged or selected.The feedback memory is formed by summing layer representations, and the weighting parameters are the only new parameters introduced.
3 METHOD: Feedback memory exposes past representations from any layer to every future computation, whereas standard Transformers restrict access to lower layers.The model adapts Transformer computation from parallel to sequential processing and uses memory vectors computed over the past.
3 METHOD: The sequential architecture supports shallower models with reduced memory footprint and increased decoding speed.The paper also summarizes substantially lower training and inference memory and faster decoding as practical advantages.
3.3 FEEDBACK TRANSFORMER: Recursive feedback lets the model repeatedly transform inputs and maintain an internal state whose capacity grows with sequence length.Unlike RNN state limited by hidden dimension, Feedback Transformer state is its whole memory and can track many items over unbounded time.
3.3 FEEDBACK TRANSFORMER: Shared key and value projections reduce per-token memory from O(L × T) to O(T), while larger batches recover some lost parallelism.Only one set of key and value vectors is stored, although projection sharing would not make a standard Transformer efficient because layer representations differ.
3.3 FEEDBACK TRANSFORMER: The sequential computation does not reduce generation performance because generation already computes one step at a time, and online reinforcement learning is likewise sequential.This scope boundary concerns generation and online reinforcement learning rather than parallel training generally.

4 EXPERIMENTS

The experiments evaluate Feedback Transformer across sequential input tasks in language processing and reinforcement learning. The paper reports that it overcomes standard Transformer memory limitations and builds complex representations with shallow models.

4 EXPERIMENTS: On the Corridor task, the Transformer degrades as memory size shrinks, while the Feedback Transformer maintains performance.The experiment is presented as evidence that feedback architecture can retain long memory.

4.1 LIMITATIONS OF TRANSFORMER: ILLUSTRATIVE TASKS

Illustrative tasks test whether models can remember information and repeatedly update internal state across long sequences. Feedback Transformer outperforms standard Transformers on corridor memory, random-walk state tracking, and long-sequence counting.

4.1.1 LIMITED MEMORY: The Corridor task shows Transformer performance degrades quickly as memory shrinks, while Feedback Transformer remains strong across memory sizes.The task requires remembering an object’s color while traversing a long hallway.
4.1.1 LIMITED MEMORY: Feedback Transformer is much better at counting over long sequences of A’s and producing the corresponding number of B’s.The experiments include training on sequences up to length 50 and up to length 1000.
4.1.2 LIMITED STATE UPDATES: 100% accuracy versus 68% shows Feedback Transformer can track random-walk state over long action sequences.Both models use 10K sequences containing 100 random actions and positions.
4.1.2 LIMITED STATE UPDATES: Transformers struggle with algorithmic state updates because conditional execution depends on current variable values that may depend on earlier conditionals.Doubling Transformer depth improves accuracy only to 47.4% for three variables and 29.1% for five, while LSTM reaches 82.8% and 32.1%, respectively.
4.1.2 LIMITED STATE UPDATES: Feedback Transformer avoids the LSTM bottleneck by accessing updated variable values from the lowest layer.This supports strong performance when tasks require maintaining several variable values.

4.2 ADVANTAGES OF FEEDBACK ARCHITECTURE

Feedback Transformers deliver stronger performance with smaller, shallower models across translation and reinforcement-learning tasks, while retaining practical speed advantages.

4.2 ADVANTAGES OF FEEDBACK ARCHITECTURE: The architecture’s advantages are analyzed on practical translation and reinforcement-learning tasks, where standard Transformers face limited memory span and limited state updating.For sequence-to-sequence tasks, Feedback is used only in the decoder because encoder inputs are simultaneously available.
4.2.1 STRONG PERFORMANCE WITH SMALL, SHALLOW MODELS: 28.3 BLEU versus 27.3 for a 1-layer Transformer, while reducing to one Feedback layer improves decoding speed 4.2x with only a 1-BLEU loss.The comparison uses WMT14 En-De and reports decoding speed in tokens per second on one GPU.
4.2.1 STRONG PERFORMANCE WITH SMALL, SHALLOW MODELS: 29.0 BLEU is achieved with a 12-layer encoder and 2-layer Feedback decoder, while the encoder’s inference parallelization makes its increased size have negligible speed impact.LayerDrop is used to stabilize training of deeper models.
4.2.2 REINFORCEMENT LEARNING: On maze navigation, Feedback Transformers converge to higher average reward than standard Transformers when agents must remember maze, object, location, and orientation information.Results are averaged over 10 trials.

4.3 COMPARISON TO OTHER ARCHITECTURES

Feedback Transformers outperform recurrent, hybrid, and standard Transformer alternatives by making higher-level past representations immediately available through shared feedback memory.

4.3.1 COMPARISON TO RECURRENT ARCHITECTURES: Feedback Transformers show stronger performance than Transformer, RNN, and hybrid RNN-Transformer models on WMT En-De.The comparison is reported in Table 2 and attributes the advantage to immediate access to all past representations rather than recurrence alone.
4.3.2 MEMORY COMPOSITION: Feedback memory performs best among tested char-PTB memory compositions, while top-only memory closely matches it and recurrence alone is insufficient.The result indicates that access to high-level representations is important for performance.
4.3.3 COMPARISON TO OTHER TRANSFORMER ARCHITECTURES: A small Feedback model matches Transformer-XL on WikiText-103, while a similarly sized standard Transformer reaches 19.9 PPL rather than 18.3.The Feedback model is approximately half the size of Transformer-XL, but takes 3.5 days to train versus 1.2 days for the Transformer.
4.3.3 COMPARISON TO OTHER TRANSFORMER ARCHITECTURES: A 12-layer Feedback model one third the size of Transformer-XL achieves 0.96 bit-per-byte on Enwiki8, reported as new state-of-the-art performance.Adaptive attention span is used because the task requires very long context.
4.3.3 COMPARISON TO OTHER TRANSFORMER ARCHITECTURES: Feedback Transformers provide faster inference through shared key-value projections, reduced memory footprint, shallow effective models, and larger feasible batch sizes.Table 5 compares training and inference speed across language modeling, translation, and reinforcement learning tasks.

5 CONCLUSION

The paper reformulates Transformers to exploit sequential input through recursive computation and increased representation power, enabling smaller and shallower models to outperform same-sized Transformers.

5 CONCLUSION: Feedback Transformers address limited access to long memory and limited ability to update state across multiple tasks by providing stronger representations and recursive computation.The conclusion presents these benefits as advantages of the proposed reformulation across the evaluated tasks.

6 ADDITIONAL RESULTS

Additional experiments show that Feedback Transformers remain effective with shallow depth, higher-level memory, and across translation, summarization, and language-modeling settings.

6.2 ADDITIONAL REINFORCEMENT LEARNING RESULTS: Maze Navigation Easy removes orientation tracking and uses four cardinal movement actions, making relative-location computation easier from action history.These task changes are presented as factors that may explain standard Transformer behavior in the easier setting.
6.2 ADDITIONAL REINFORCEMENT LEARNING RESULTS: Water Maze changes the cell-to-ID mapping between episodes while requiring agents to remember previously seen cells during each episode.The task is modified to make the Morris Water Maze more challenging.
6.1 ADDITIONAL SEQUENCE-TO-SEQUENCE RESULTS: On IWSLT De-En, shallower Feedback models outperform standard Transformers, and on CNN-Dailymail the improvement remains consistent across decoder depths.The translation evaluation uses tokenized BLEU, while summarization uses test ROUGE-L.
6.4 ABLATION STUDIES ON LANGUAGE MODELS: Higher-layer representations work better as Feedback memory, while a weighted sum of all layers matches the best performance and first-layer memory matches the standard Transformer.The weighted sum can adaptively select among layer representations.
6.3 ADDITIONAL LANGUAGE MODELING RESULTS: Adaptive span improves both models as context size increases, while increasing BPTT length for training efficiency does not affect final WikiText-103 performance.Pre-normalization with higher learning rates particularly helps the standard Transformer, and the performance gap remains consistent across these variations.
6.4 ABLATION STUDIES ON LANGUAGE MODELS: With parameter count held constant, Feedback Transformers are robust to reduced depth and achieve their best char-PTB performance with only two layers.Standard Transformer performance improves as depth increases under the same parameter-budget control.

7 ADDITIONAL IMPLEMENTATION DETAILS

The experiments use task-specific setups across random walk, maze navigation, water maze, algorithmic, sequence-to-sequence, and language-modeling evaluations. Implementation choices include shared Transformer settings, long attention spans, adaptive attention, caching, and fine-tuning procedures.

7.1 RANDOM WALK TASK DETAILS: Random-walk models receive action sequences and reset symbols, predict location after each action, and are trained on 10k episodes totaling 1M tokens.The attention span is 100, allowing access to all information needed for the task.
7.2 MAZE NAVIGATION DETAILS: Maze-navigation episodes use randomly generated 9 × 9 mazes with eight colored targets, 200-step limits, and observations containing a local 3 × 3 area plus target color.Agents are trained with A2C; the Transformer uses two layers, hidden size 256, four heads, and attention span 200.
7.3 WATER MAZE DETAILS: Water-maze episodes use 15 × 15 grids, local 3 × 3 goal visibility, 200-step limits, and 500M training steps across 2.5M episodes.The Transformer configuration uses two layers, hidden size 64, one head, and attention span 200; agents are trained with A2C and RMSprop.
7.4 ALGORITHMIC TASK DETAILS: Algorithmic programs contain 100 sequential statements, with initialization, increment/decrement, print, and conditional operations, and use either three or five variables.Training data contains 10k concatenated programs separated by a special keyword, using the random-walk hyperparameters.
7.5 MACHINE TRANSLATION AND SUMMARIZATION: Sequence-to-sequence experiments use Transformer base, big, or smaller models by task, while modifying only the decoder to use the Feedback Transformer architecture.Hyperparameters are reported in Table 8.
7.6 LANGUAGE MODELING: Language-modeling experiments add hidden-state caching, relative position embeddings, adaptive attention span, and pre-normalization to support unbounded sequences and stable, efficient training.The best performance uses attention span 2048 with adaptive span loss 0.00001, and adding EOL tokens during fine-tuning improves perplexity by 1 point.