PaperlayerView exact revision on arXiv

Source-linked AI summary

Mega: Moving Average Equipped Gated Attention

Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, Luke Zettlemoyer

arXiv:2209.10655v3cs.LG

TL;DR

Transformer attention is limited for long sequences by weak inductive bias and quadratic computational cost. Mega combines single-head gated attention with exponential moving averages, while Mega-chunk uses fixed-length chunks for linear complexity. Across five sequence-modeling tasks and multiple data types, Mega reports improvements over strong Transformer and state-space-model baselines, with author-noted evaluation and tuning caveats.

  • Problem

    Transformer attention lacks prior positional dependency patterns and requires quadratic time and space, limiting its application to long-sequence modeling.

  • Method

    Mega integrates a learnable damped exponential moving average with single-head gated attention, and Mega-chunk obtains linear complexity by chunking sequences while preserving context through EMA.

  • Results

    Across five sequence-modeling tasks spanning long-context modeling, translation, language modeling, image classification, and speech classification, Mega significantly outperforms strong baselines in effectiveness and efficiency.

  • Takeaways & Limitations

    The results support combining different inductive biases to model long- and short-term dependencies across varied sequence-modeling data types.

  • Takeaways & Limitations

    Image-classification performance may be sub-optimal because the DeiT data augmentation and regularization setup may not be optimal for Mega.

Abstract

from arXiv · show

The design choices in the Transformer attention mechanism, including weak inductive bias and quadratic computational complexity, have limited its application for modeling long sequences. In this paper, we introduce Mega, a simple, theoretically grounded, single-head gated attention mechanism equipped with (exponential) moving average to incorporate inductive bias of position-aware local dependencies into the position-agnostic attention mechanism. We further propose a variant of Mega that offers linear time and space complexity yet yields only minimal quality loss, by efficiently splitting the whole sequence into multiple chunks with fixed length. Extensive experiments on a wide range of sequence modeling benchmarks, including the Long Range Arena, neural machine translation, auto-regressive language modeling, and image and speech classification, show that Mega achieves significant improvements over other sequence models, including variants of Transformers and recent state space models.

1. Introduction

Mega addresses Transformer attention’s weak positional inductive bias and quadratic sequence-length cost by combining gated attention with exponential moving averages. Across five sequence-modeling tasks, it reports improvements in effectiveness and efficiency over strong baselines.

  • Motivation: Transformers learn pairwise attention weights without prior dependency patterns and require quadratic time and space in sequence length.These properties limit accuracy and efficiency on long-sequence tasks.
  • Method: Mega combines single-head gated attention with a learnable, multi-dimensional damped exponential moving average to encode local dependencies.The EMA introduces timestep-aware inductive bias into attention.
  • Method: Mega-chunk splits sequences into fixed-length blocks to achieve linear complexity while retaining contextual information through the EMA sub-layer.Chunk-wise attention has complexity O(nc), and EMA extends the effective context beyond chunk boundaries.
  • Theory: Mega’s single-head gated attention is theoretically shown to be as expressive as commonly used multi-head attention.This provides a theoretical basis for replacing multi-head attention with a single-head gated variant.
  • Results: Across five sequence-modeling tasks spanning long-context modeling, translation, language modeling, image classification, and speech classification, Mega significantly outperforms strong baselines in effectiveness and efficiency.The benchmarks include LRA, WMT16 en-de, WikiText-103, ImageNet-1k, and SC-Raw.

2. Background

The background contrasts self-attention’s flexible long-distance dependency modeling with EMA’s position-aware local bias. Mega combines these complementary properties to retain expressive dependency modeling while enabling linear-complexity chunk-wise attention.

  • Self-Attention Mechanism: Self-attention forms an attention matrix assigning dependency strengths between every pair of tokens.Queries, keys, and values are learned transformations of the input sequence.
  • Self-Attention Mechanism: Self-attention can learn long-distance dependencies with minimal inductive bias, but recognizing dependency patterns is challenging on long sequences.Its attention matrix costs O(hn^2) time and space with h heads.
  • Exponential Moving Average: EMA weights observations with factors that decrease exponentially, with higher α discounting older observations faster.EMA is used in sequential modeling to smooth short-term fluctuations and highlight longer-term trends or cycles.
  • Exponential Moving Average: EMA imposes an input-agnostic exponential decay on dependency weights, favoring local dependencies while limiting long-distance dependencies.Despite its recurrent formulation, EMA can be computed as individual convolutions using FFTs.
  • Combining Attention with EMA: Mega embeds EMA into attention to combine strong inductive bias with the capacity to learn complex dependency patterns.This integration also enables computationally efficient chunk-wise attention with linear complexity in sequence length.

3. Moving Average Equipped Gated Attention (Mega)

Mega combines a multi-dimensional damped exponential moving average with single-head gated attention, adding local dependency bias while retaining expressive attention. Mega-chunk applies attention within fixed-length chunks to reduce complexity, while EMA preserves contextual information across chunk boundaries.

  • Mega architecture: Mega combines multi-dimensional damped EMA with single-head gated attention as its core architecture.The damped EMA improves flexibility and capacity, while the gated attention mechanism is integrated into Mega blocks.
  • Multi-dimensional Damped EMA: Multi-dimensional damped EMA expands each input dimension into h dimensions, applies damped EMA in that hidden space, and projects the result back to one dimension.The resulting sequence is denoted EMA(X), with complexity similar to standard EMA.
  • Gated attention: Mega computes gated attention from EMA-contextualized inputs, using reset and update gates to combine attention output with candidate activations.The EMA output supplies the contextual input for shared representations and the candidate activation, while the update gate produces the final output.
  • Gated attention: The proposed Laplace attention function addresses unstable training associated with the unbounded range and gradient of squared ReLU attention.Its coefficients are adjusted to approximate squared ReLU while using a Laplace-based formulation.
  • Theoretical justification: Single-head gated attention is theoretically as expressive as multi-head attention when its gate transformation is a universal approximator.The gate vector multiplies the single-head output, matching the expressiveness of multi-head attention under the theorem’s assumption.
  • Mega-chunk: Mega with Linear Complexity: Mega-chunk splits queries, keys, and values into fixed-length chunks, reducing attention complexity to O(kc^2) = O(nc) while EMA carries local context across chunk boundaries.This design has linear complexity with respect to sequence length n when chunk length c is fixed, though chunk-wise attention alone would lose information from other chunks.

4. Experiments

Experiments across long-context, speech, language, translation, and image tasks show that Mega combines strong effectiveness with improved efficiency, while Mega-chunk achieves linear complexity with competitive quality.

  • Long-Context Sequence Modeling: Mega substantially outperforms all baselines across the six Long Range Arena tasks.Baselines include Transformer variants and S4 models.
  • Long-Context Sequence Modeling: 5.5 times faster and uses 13% as much memory as vanilla Transformer on 4K-token Text classification.These efficiency measurements apply to Mega-chunk.
  • Long-Context Sequence Modeling: Removing EMA causes rapid accuracy declines on both Text and Image tasks, while even one EMA dimension yields significant improvement.The ablation varies EMA dimension h from 0 to 32, with h = 0 removing EMA.
  • Long-Context Sequence Modeling: Around 90% Text accuracy and around 75% Image accuracy are achieved by Mega-chunk with chunk size c = 16.Image performance is more sensitive to chunk size than text performance, and the Image result remains better than vanilla Transformer.
  • Raw Speech Classification: 96.92 accuracy for Mega-chunk base on Speech Commands is slightly below S4's 97.50, while the larger model performs comparably to S4.The base model has 300K parameters; adding 0.18M parameters produces the larger model.
  • Neural Machine Translation and Image Classification: Mega outperforms Transformer-base by over 1.1 BLEU on WMT'16 and achieves about 0.5% higher ImageNet accuracy than DeiT-B.On language modeling, Mega reports competitive results with 9× faster inference than Transformer and supports length extrapolation at inference time.

5. Related Work

Related work addresses Transformer limitations through stronger positional inductive biases, more efficient attention, and convolutional or state-space alternatives.

  • Inductive Bias: One line of work injects position information through absolute or relative positional embeddings and relative positional biases.These methods strengthen the attention mechanism's positional inductive bias.
  • Computational Efficiency: Efficient Transformer variants use sparse attention, low-rank approximations, or kernelization to improve long-sequence time and memory complexity.The passage notes that efficiency gains are less prominent for moderate-length sequences and performance remains behind the target models.
  • Convolutional Neural Networks with Continuous Kernels: Mega is related to continuous-kernel CNNs because EMA and state-space models such as S4 can be viewed as convolution transforms with sequence-length kernels.The related models include CKConv, FlexConv, and CCNN.

6. Conclusion

Mega is introduced as an efficient replacement for multi-head attention that uses exponential moving averages to add inductive bias, with Mega-chunk providing linear complexity. Its EMA formulation can be vectorized and accelerated through convolution kernels and FFTs.

  • Mega is presented as a simple, efficient, effective drop-in replacement for regular multi-head attention.
  • Mega incorporates stronger inductive biases into attention by leveraging exponential moving averages.
  • Mega-chunk is an efficient Mega variant with linear complexity that improves over strong baselines across five sequence modeling tasks.
  • The multi-dimensional damped EMA uses independent dimensions and can be simplified to a scalar formulation without loss of generality.
  • The EMA output combines the initial hidden state with transformed inputs weighted by powers of the decay term.
  • The vectorized EMA computation uses a convolution kernel, with the kernel efficiently computed through a Vandermonde product and the output computed using FFTs.

Appendix B. Proof of Theorem 1

The proof shows that splitting gated attention into heads preserves the relevant normalized attention computation. Because the gating transformation is a universal approximator, it can theoretically recover the required head-wise ratios.

  • The proof splits the gating vector γ into h heads in the same way as Q, K, and V.
  • For each head, γ(i) equals the element-wise ratio between a(i)T V(i) and aT V(i).
  • Because G(X) is a universal approximator and Q, K, V, and a derive from X, γ can theoretically recover these head-wise ratios for every X.

Appendix C. Laplace Attention Function

The appendix motivates Laplace attention as an alternative to squared ReLU and compares their training stability. On LRA Pathfinder, Laplace is reported to be much more stable than ReLU2.

  • Laplace attention is compared with the squared ReLU2 function as an attention-function alternative.
  • The coefficients μ and σ for approximating squared ReLU with Laplace are derived by solving two equations at x = 1/4π.
  • On LRA Pathfinder, validation accuracy across training epochs shows that Laplace is much more stable than ReLU2.

Appendix D. Experimental Details

The experimental details specify that Mega follows established preprocessing and data-split procedures, with hyperparameters documented for LRA and raw speech classification. Speech classification uses the 10-class SC10 subset, with Mega-chunk configured with chunk size 1000.

  • Mega follows Tay et al. (2020) for data preprocessing, data splits, and related experimental details.
  • Table 8 lists Mega hyperparameters for LRA and raw speech classification, including batch size, learning rate, and weight decay.
  • Speech classification uses the 10-class SC10 subset of the Speech Commands dataset, with Mega-chunk chunk size set to 1000.

D.3 Language Modeling

The language-modeling experiments evaluate Mega under varying context and attention lengths, including inference-time extrapolation beyond training lengths. Mega improves as both context length and attention chunk size increase, despite shorter lengths during training.

  • Training details: Mega-chunk training uses randomly sampled segment lengths, with chunk sizes serving as the effective attention lengths.The sampled range is 2–6 chunks for WikiText-103 and 2–4 for enwik8.
  • Length extrapolation at inference time: Mega-chunk inference uses attention chunk sizes of 1024 for WikiText-103 and 2048 for enwik8.Longer attention lengths are evaluated at test time than during training.
  • Length extrapolation at inference time: Mega extrapolates to sequences longer than those seen during training because its recurrent EMA layer carries historic context across chunks.Rotary positional embeddings are applied to the attention sublayer when increasing inference-time attention length.
  • Ablations on context lengths: PPL decreases as context length increases, although Mega saw a maximum training context length of only 6144 tokens.The ablation varies maximum context tokens from 2.5K to 49K.
  • Ablations on attention chunk sizes: Mega consistently improves as attention chunk size increases from 512 to 3072, despite training with an attention length of only 1024.This result contrasts with the cited finding that rotary embeddings do not generalize to longer lengths and increase PPL.

D.4 Machine Translation

The machine-translation experiments use WMT 2016 English–German with standard validation and test splits and compare Transformer and Mega models under listed hyperparameters.

  • Dataset and preprocessing: The translation data uses shared source–target vocabularies with 32K byte-pair encoding types.Preprocessing follows the cited FairSeq-based scripts.
  • Model configuration: Transformer and Mega machine-translation models use the hyperparameters listed in Table 10.The passage identifies the table as covering model hyperparameters for machine translation.
Loading 2209.10655v3…