PaperlayerView exact revision on arXiv

Source-linked AI summary

Linear Transformers Are Secretly Fast Weight Programmers

Imanol Schlag, Kazuki Irie, Jürgen Schmidhuber

arXiv:2102.11174v3cs.LG

TL;DR

Standard self-attention has quadratic sequence-length cost, while linearised variants trade this for constant-size memory and can encounter capacity limits. The paper establishes their equivalence to Fast Weight Programmers, introduces delta-rule-like memory updates and a new kernel projection, and evaluates them across retrieval, translation, and language modelling tasks.

  • Problem

    Self-attention scales quadratically with sequence length, while linearised variants have finite memory capacity and existing softmax-linearisation techniques remain underexplored.

  • Method

    The paper interprets linearised self-attention as Fast Weight Programming, introduces a delta-rule-like update instruction, and proposes deterministic parameter-free projection for attention linearisation.

  • Results

    Experiments demonstrate the benefits of the proposed methods on synthetic retrieval, WMT14 English-to-German translation, and Wikitext-103 language modelling.

  • Takeaways & Limitations

    The Fast Weight Programmer perspective provides a framework for analysing linear-attention capacity and designing memory-editing and attention-linearisation methods.

  • Takeaways & Limitations

    Linear attention has finite capacity, and FAVOR+ remains limited in practice while its random-feature sampling introduces output variance.

Abstract

from arXiv · show

We show the formal equivalence of linearised self-attention mechanisms and fast weight controllers from the early '90s, where a ``slow" neural net learns by gradient descent to program the ``fast weights" of another net through sequences of elementary programming instructions which are additive outer products of self-invented activation patterns (today called keys and values). Such Fast Weight Programmers (FWPs) learn to manipulate the contents of a finite memory and dynamically interact with it. We infer a memory capacity limitation of recent linearised softmax attention variants, and replace the purely additive outer products by a delta rule-like programming instruction, such that the FWP can more easily learn to correct the current mapping from keys to values. The FWP also learns to compute dynamically changing learning rates. We also propose a new kernel function to linearise attention which balances simplicity and effectiveness. We conduct experiments on synthetic retrieval problems as well as standard machine translation and language modelling tasks which demonstrate the benefits of our methods.

1. Introduction

Linear Transformers address the quadratic cost and context limits of self-attention by using constant-size memory and linear-time computation. The paper connects these models to Fast Weight Programmers, identifies capacity and softmax-linearisation issues, and evaluates proposed methods on retrieval, translation, and language modelling.

  • Motivation: Self-attention enables parallel computation across time but scales quadratically with sequence length, limiting context windows and longer-term dependencies.Its memory grows linearly, so practical systems must restrict context size.
  • Motivation: Linear Transformers reduce computation to linear time with constant-size memory by linearising the softmax.The paper frames this as the main source of their complexity reduction.
  • Contributions: Linearised self-attention is formally equivalent to Fast Weight Controllers, whose memories store key-value associations updated through additive outer products.The fast-weight perspective treats keys and values as self-invented activation patterns used to program memory.
  • Contributions: Finite memory capacity can cause overcapacity, motivating a delta-rule-like instruction that selectively corrects, keeps, or deletes key-value associations.The proposed instruction is intended to improve memory editing when sequence length exceeds storage capacity.
  • Contributions: The paper compares existing softmax-linearisation functions and proposes a method balancing simplicity and effectiveness.Existing approaches are described as either simplistic or mathematically complex.
  • Experiments: Experiments cover synthetic retrieval, WMT14 English-to-German translation, and Wikitext-103 language modelling.These tasks are used to demonstrate the benefits of the proposed methods.

2. Background on Fast Weight Programmers

Fast Weight Programmers make network weights variable and input-dependent, allowing a slow network to program a fast network’s short-term associative memory. Their memory writes use differentiable outer-product updates and retrieves values through matrix multiplication.

  • Fast weights: Fast weights are variable and input-dependent, unlike standard neural-network weights that remain fixed after training.The concept is also associated with synaptic modulation, variable binding, and dynamic connections.
  • Fast Weight Programmers: In context-dependent FWPs, a slow network continually reprograms a fast network, making fast weights depend on the input stream’s spatio-temporal context.The slow network learns to program the fast network.
  • Fast Weight Programmers: Fast weights are generated at each time step and serve as a short-term memory.The programming instructions use trainable slow weights and activation functions.
  • Memory operations: The associative memory writes by summing outer products of keys and values, then retrieves through matrix-vector multiplication.This creates key-value associations in the fast-weight matrix.
  • Historical context: FWPs differ from earlier associative-memory systems because they learn to use programming rules for associating self-invented patterns.Earlier systems used pre-wired rules for given patterns.
  • Related work: The fast-weight idea has been revisited under names including hypernetworks, dynamic plasticity, dynamic convolution, and lambda networks.These variants have been applied to areas including meta-learning.

3. Relation to Transformers

Removing or linearising softmax reveals that Transformer self-attention performs fast-weight programming with outer-product memory updates. Kernel replacement preserves this structure while enabling linearised attention variants with normalised retrieval.

  • Relation to Transformers: Prior work connects outer-product Fast Weight Programmers with attention, and this paper reviews the derivation linking them to Transformers.The review emphasizes the relation between Transformer attention and earlier FWPs.
  • Self-attention without softmax: Without softmax, self-attention can be rewritten as a Fast Weight Programmer using identity activation and no query projection.The resulting equations directly relate the Transformer computation to fast-weight equations.
  • Motivation: Linearised attention improves computational efficiency for long sequences compared with standard self-attention.Prior work introduced these techniques rather than simply removing softmax.
  • Linearising softmax: Softmax linearisation replaces the exponential kernel with κ′(k, q) = φ(k)ᵀφ(q), where φ maps key and query vectors into a dot-product space.This kernel substitution provides the basis for linearised attention.
  • Linearising softmax: The resulting numerator uses outer-product notation analogous to the no-softmax case.An additional vector tracks the denominator for normalisation.
  • Linear Transformer form: Linear Transformer forward computation is a normalised Fast Weight Programmer based on a fast-weight matrix and denominator vector.Thus, outer-product FWPs form the core of linear Transformer variants.

4. Analysing and Improving Linear Transformers as Fast Weight Programmers

Viewing linear Transformers as Fast Weight Programmers exposes capacity limits in their associative memories and motivates update rules that can correct or selectively manage stored associations. The paper proposes delta-rule programming with dynamic write strengths and sum normalisation to improve operation in these settings.

  • Analysis: Linear Transformer variants are outer-product-based Fast Weight Programmers whose memories store key-value associations.Their differentiable update instructions program memory through additive outer products.
  • Capacity Limitation: When sequence length exceeds key dimension, storing more than that dimension’s associations can cause retrieval errors from non-orthogonal keys.The overcapacity regime arises because more than d orthogonal key vectors cannot fit in a d-dimensional space.
  • Capacity Limitation: In overcapacity, an ideal memory should dynamically select which key-value associations to retain or forget rather than only accumulating new associations.This motivates improving the elementary differentiable programming instruction.
  • Improving the FWP’s Programming Instruction: The proposed delta-rule update retrieves the current value for a key, then uses a model-generated interpolation weight to correct the association.The interpolation weight β(i) acts as a dynamic write strength controlling how much the new value replaces the retrieved value.
  • Improving the FWP’s Programming Instruction: Sum normalisation divides effective keys and queries by their component sums, making matrix-vector outputs interpretable as weighted sums over memory columns.The paper reports that language modelling experiments do not require additional attention normalisation after this substitution.

5. Linear Attention Functions

The paper reviews kernel maps for linear attention, emphasizing how their codomain dimension controls memory capacity. It then introduces DPFP, a deterministic projection that increases this dimension while promoting sparse, orthogonal representations.

  • Design requirements: The kernel function φ maps keys and queries into a dot-product space whose positivity enables proper attention weights and whose dimension controls memory capacity.Projecting from d_key to a larger dimension d̂ can increase the capacity bound.
  • Existing functions: ELU + 1 is simple and maintains non-zero negative-side gradients, but preserves d̂ = d_key and therefore does not modify memory capacity.Its element-wise form avoids changing the key dimension.
  • Existing functions: FAVOR+ uses random features to approximate softmax and expands the codomain to d̂ = 2m, increasing theoretical capacity when 2m > d_key.Its sampling introduces output variance, and finite m remains below the infinite capacity of softmax memory.
  • Deterministic Parameter-Free Projection (DPFP): DPFP provides a deterministic, parameter-free alternative that increases dot-product dimension without FAVOR+ random features.The construction uses partial functions and two-factor features designed to promote orthogonality in the projected space.
  • Deterministic Parameter-Free Projection (DPFP): In the two-dimensional example, DPFP maps each input vector to a four-dimensional vector with one non-zero component, partitioning the input plane into four orthogonal regions.Figure 1 visualizes the four partial functions as colored surfaces representing the output components.
  • Deterministic Parameter-Free Projection (DPFP): For higher-dimensional inputs, DPFP constructs additional two-factor features, with ν controlling capacity and output dimension d̂ = 2d_keyν.Each partial function can be computed independently, making the construction highly parallelisable.

6. Experimental Results

Experiments test linear attention’s capacity limits and update rules on synthetic retrieval, machine translation, and language modelling. The proposed delta update and DPFP kernel improve performance across these settings, while retaining favorable resource usage.

  • 6.1.1. SETTING 1: TESTING CAPACITY: Synthetic retrieval tests show linear attention models fail when the number of associations exceeds the key-feature dimension.With sequence length equal to the number of unique keys, models using the sum update rule are evaluated across increasing association counts.
  • 6.1.1. SETTING 1: TESTING CAPACITY: 64 associations mark Linear-Attention’s capacity, while DPFP begins accumulating errors near its limits of 128, 256, and 384.FAVOR+ never reaches zero loss, whereas softmax attention outperforms the linear variants but struggles to converge beyond 500 keys.
  • 6.1.2. SETTING 2: COMPARING UPDATE RULES: The delta update rule outperforms the compared update variants when repeated keys receive new values in sampled sequences.This setting requires retrieving the most recent value and therefore tests whether the finite memory can revise existing associations.
  • 6.2. Machine Translation Experiments: On WMT14 English-to-German translation, DPFP outperforms Linear Transformer and Performer when the feature dimension is relatively small.Performer matches the basic Transformer when its number of random features is sufficiently large, while DPFP offers a simplicity–performance trade-off at smaller dimensions.
  • 6.3. Language Modelling Experiments: On WikiText-103, the delta update rule provides convincing perplexity improvements over sum-update models in both evaluated configurations.The ablation finds that models diverge without sum normalisation, while better perplexities result without additional attention normalisation.

7. Conclusion

The paper connects linearised self-attention with Fast Weight Programmers, uses this perspective to identify finite-memory limitations, and proposes editable memory updates and a new attention linearisation.

  • 7. Conclusion: Linearised self-attention is formally connected to Fast Weight Programmers that write key-value associations through additive outer-product instructions.The FWP perspective frames these instructions as differentiable programming operations for fast-weight memories.
  • 7. Conclusion: The FWP perspective exposes associative-memory capacity limitations in linear attention and motivates alternative programming instructions for finite-memory Transformers.The paper specifically links overcapacity to the need for dynamically editing memory contents.
  • 7. Conclusion: The proposed alternative instruction enables FWPs to dynamically edit memory using a delta-rule-like update learned through gradient descent.The paper describes this as an alternative differentiable elementary programming instruction.
  • 7. Conclusion: A new method for linearising attention is proposed, and experiments on synthetic and real language tasks demonstrate the effectiveness of the proposals.The conclusion positions the FWP perspective as a basis for investigating improved programming instructions and finite-memory Transformer designs.

A. Update Rule Derivation

The derivation interprets fast-weight matrices column-wise as associative memories and compares the proposed update with Peng et al.’s gated rule, emphasizing preservation of unrelated associations.

  • Update Rule Derivation: A fast-weight matrix can be represented through column vectors, which correspond to associations between fixed Cartesian-basis keys and associated values.This column-wise view supports deriving updates for individual memory associations.
  • Update Rule Derivation: Updating a matrix with a new key-value association can therefore be expressed as updating its individual column vectors.The derivation begins from an arbitrary weight matrix and applies the proposed update rule.
  • Formal comparison to Peng et al. (2021): Peng et al.’s rule updates the repeated key toward a convex combination of old and new values but can also modify or erase an unrelated association.The example assumes orthonormal keys and repeats k2 as k3.
  • Formal comparison to Peng et al. (2021): The proposed rule likewise updates the repeated key toward (1 −β)v2 + βv3 while using the retrieved current value as the previous association.For k3 = k2, the retrieved value is v2, yielding the same convex-combination form for that association.
  • Formal comparison to Peng et al. (2021): The key distinction is that the proposed update preserves unrelated associations, unlike the compared rule.The paper presents this property as the central difference in associative-memory behavior.

C. DPFP-ν Implementation

The appendix implements DPFP-ν in PyTorch as a compact sequence of tensor operations and provides an additional Figure 4 training-curve result for synthetic setting 1.

  • C. DPFP-ν Implementation: DPFP-ν is implemented with two concatenations followed by one element-wise multiplication.The listing is described as a simple PyTorch implementation of Eq. 37.
  • C. DPFP-ν Implementation: The implementation applies ReLU and constructs repeated and rolled inputs before returning their element-wise product.The code excerpts identify ReLU, the dpfp function, and the final multiplication operation.
  • D.1. Synthetic Task Setting 1: Figure 4 presents training curves for synthetic setting 1 with 600 unique keys and values sampled without replacement.The appendix notes that the figure corresponds to setting 1 and identifies the number of unique keys and values.

D.2. Synthetic Task Setting 2

Synthetic setting 2 evaluates retrieval under increasing association counts, while an additional language-model experiment examines the update rule in a non-overcapacity regime.

  • D.2. Synthetic Task Setting 2: The Figure 5 capacity plot increases the number of unique keys and queries, but excludes FAVOR+ because its update-rule combination produced not-a-number results.The plot is described as analogous to the setting 1 capacity figure, except for the logarithmic y-axis.
  • D.2. Synthetic Task Setting 2: Figure 5 evaluates final loss on setting 2 problems with 20 to 200 unique associations and reports that the update rule outperforms all other approaches.The problems vary sequence lengths and numbers of unique keys, with models trained until convergence.
  • D.2. Synthetic Task Setting 2: The additional language-modelling experiment trains Performer and DPFP in a small setting with D = 128, L = 256, m = 16, and ν = 1.Both models use ddot = 256, allowing comparison in a non-overcapacity scenario.

E. Details on Machine Translation Experiments

The translation experiments implement different φ functions in FAIRSEQ using the Transformer big architecture. The setup also specifies model dimensions and kernel-related hyperparameters for the evaluated configurations.

  • Different φ functions were implemented in the FAIRSEQ toolkit for the translation experiments.
  • For the small WikiText-103 configuration, D = 128 and L = 256 correspond to 40 M trainable parameters.
  • The Transformer big model uses 6 encoder and 6 decoder layers, hidden size 1024, 16 attention heads, and 4096-dimensional feed-forward layers.
  • The translation models use 32 K byte-pair encoding sub-word units.
  • All translation models use dkey = 64, while Performer and DPFP configurations relate dot-product dimension to their hyperparameters.

F. Details on Language Modelling Experiments

The language-modelling implementations build on PyTorch and Transformer-XL code, with a custom fast-weight update and backward pass. Experimental details specify configuration-dependent batch sizes, training durations, dropout, and optimisation settings.

  • Language-modelling implementations use PyTorch and public Transformer-XL code as a starting point.
  • The fast-weight update modifies a publicly available Linear Transformer CUDA kernel.
  • A custom backward pass stores one fast-weight set and recomputes time-step weights to avoid GPU memory limits while remaining time-efficient.
  • Small and medium configurations use batch sizes of 96 and 56 sequences and train for about 120 and 70 epochs, respectively.
  • Both configurations use 10% dropout and Adam with initial learning rate 0.00025 and 2000 warm-up steps.
Loading 2102.11174v3…