[Paper Review] Addressing Some Limitations of Transformers with Feedback Memory

Angela Fan, Thibaut Lavril+3 more|arXiv (Cornell University)|Feb 21, 2020

Topic Modeling46 references31 citations

TL;DR

The paper introduces the Feedback Transformer, which makes all past representations accessible to future computations via a shared feedback memory, enabling shallow models to match or exceed Transformers on language modeling, translation, and reinforcement learning while offering faster decoding and smaller memory footprints.

ABSTRACT

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.

Motivation & Objective

• Motivate and address limitations of standard Transformers for autoregressive sequence modeling.
• Introduce a memory-based Feedback Transformer that reuses past higher-level representations to compute current steps.
• Demonstrate that shallow Feedback Transformers achieve strong performance with faster decoding and reduced memory usage.
• Evaluate on language modeling, machine translation, and reinforcement learning benchmarks.
• Compare to recurrent architectures and other Transformer variants to establish advantages of feedback memory.

Proposed method

• Compute a memory vector m_t that aggregates representations from all layers at time t using a layer-weighted sum with learnable weights w^l and Softmax normalization: m_t = sum_l Softmax(w^l) x_t^l.
• Replace self-attention inputs to attend to the memory m_t instead of only lower-layer outputs: z_t^l = Attn(x_t^l, {m_{t-τ}, ..., m_{t-1}}).
• Share key and value projections across all layers so that k_t^l = k_t = W_k m_t and v_t^l = v_t = W_v m_t, reducing memory and computation.
• Merge past representations from all layers into a single feedback memory per time step, enabling recursive computation without increasing parameter count besides w^l.
• Maintain sequential computation at generation time, while benefiting from hardware-friendly parallelism during training via memory sharing.
• Demonstrate that smaller Feedback Transformers achieve strong results on translation with shallow decoders and large encoders, and show improved memory and speed characteristics.

Experimental results

Research questions

• RQ1 Can exposing all past representations to future computations via a Feedback Memory improve autoregressive modeling compared to standard Transformers?
• RQ2 Do Feedback Transformers enable effective recursive computation and long-term state tracking with smaller, shallower networks?
• RQ3 How do Feedback Transformers compare to recurrent models and other Transformer variants on language modeling, translation, and reinforcement learning tasks?
• RQ4 What are the trade-offs in memory and speed when applying shared key-value projections in Feedback Transformers?
• RQ5 How does the Feedback mechanism affect decoding speed and training efficiency in practical settings?

Key findings

• Feedback Transformers preserve or improve performance with shallower architectures compared to standard Transformers on translation and language tasks.
• Sharing key/value projections reduces memory usage and speeds up training and inference, enabling faster decoding especially with shallow decoders.
• On long-context tasks and RL, the Feedback model maintains performance with limited memory and demonstrates better state tracking than standard Transformers.
• The architecture can outperform recurrent models and certain Transformer hybrids on standard benchmarks like WMT En-De and WikiText-103 with smaller parameter counts.
• Feedback memory enables recursive computation, allowing representations to be refined iteratively across time steps without increasing model depth.