PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformer Quality in Linear Time

Weizhe Hua, Zihang Dai, Hanxiao Liu, Quoc V. Le

arXiv:2202.10447v2cs.LGcs.AIcs.CLcs.NE

TL;DR

Long-context Transformers face quadratic attention costs, while existing efficient variants can lose quality or practical training speed. FLASH combines a gated attention unit with mixed chunk attention to obtain linear scalability and competitive quality. It matches augmented Transformers in perplexity across short and long contexts while achieving substantial training speedups across language-modeling settings.

  • Problem

    Quadratic attention complexity limits Transformers’ ability to process long-term information, while efficient variants may suffer quality loss, accelerator overhead, or slow autoregressive training.

  • Method

    FLASH combines Gated Attention Unit, which supports effective weaker single-head attention, with mixed chunk attention that uses local quadratic and global linear attention.

  • Results

    FLASH matches augmented Transformers in quality across 512–8K contexts and achieves speedups up to 4.9× on Wiki-40B, 12.1× on PG-19, and 4.8× on C4.

  • Takeaways & Limitations

    FLASH provides a practical long-context Transformer variant that combines competitive perplexity with linear scalability and substantially faster training.

  • Takeaways & Limitations

    The paper identifies quality, practical overhead, and autoregressive-training inefficiencies as limitations of existing efficient attention methods.

Abstract

from arXiv · show

We revisit the design choices in Transformers, and propose methods to address their weaknesses in handling long sequences. First, we propose a simple layer named gated attention unit, which allows the use of a weaker single-head attention with minimal quality loss. We then propose a linear approximation method complementary to this new layer, which is accelerator-friendly and highly competitive in quality. The resulting model, named FLASH, matches the perplexity of improved Transformers over both short (512) and long (8K) context lengths, achieving training speedups of up to 4.9$\times$ on Wiki-40B and 12.1$\times$ on PG-19 for auto-regressive language modeling, and 4.8$\times$ on C4 for masked language modeling.

1. Introduction

Transformers struggle with long contexts because attention has quadratic input-length complexity, while existing efficient variants can sacrifice quality or practical accelerator speed. FLASH addresses these issues with GAU and mixed chunk attention, matching augmented Transformers in quality while substantially reducing training cost.

  • Quadratic attention complexity limits Transformers’ context length and their ability to process long-term information.
  • Existing efficient attention methods can incur quality drops, practical accelerator overhead, or slow autoregressive training despite favorable theoretical complexity.
  • FLASH introduces Gated Attention Unit (GAU), using gating to reduce attention’s quality burden and make weaker single-head, softmax-free attention effective.
  • FLASH approximates GAU attention with mixed chunk attention, combining quadratic within-chunk attention, linear across-chunk attention, and accelerator-efficient implementation.
  • FLASH remains competitive with Transformer++ across 512–8K contexts and tasks spanning masked and autoregressive language modeling.
  • 1.2×–4.9× speedups occur on Wiki-40B and 1.0×–4.8× on C4, while PG-19 training cost falls by up to 12.1× with significant quality gains.

2. Gated Attention Unit

GAU unifies attention and gating into a simpler, more efficient layer whose quality depends less on precise attention, making it suitable for approximation while remaining competitive with Transformers.

  • GAU unifies attention and GLU computation, sharing as much computation as possible to improve parameter and compute efficiency.
  • GAU replaces same-token GLU gating with an attention-retrieved representation, reducing to GLU when the attention matrix is the identity.
  • Gating allows GAU to use simpler attention than MHSA without quality loss, including a small single-head, softmax-free mechanism.
  • GAU’s attention adds only a ds-parameter matrix Wz, compared with 4d^2 MHSA parameters; two GAUs can replace one Transformer block at similar model size and training speed.
  • GAUs are competitive with Transformers on TPUs across model sizes at context length 512, with stronger performance reported at longer contexts.

3. Fast Linear Attention with GAU

The paper extends GAU with mixed chunk attention, combining local quadratic attention within chunks and global linear attention across chunks. Chunking reduces autoregressive sequential dependencies while preserving long-range interaction and linear scaling in context length.

  • Motivation: GAU uses gating to make single-head, softmax-free attention effective, reducing reliance on precise attention computation.This motivates approximating GAU attention rather than directly approximating multi-head self-attention.
  • Existing Linear-Complexity Variants: Linear attention reduces complexity with respect to sequence length, but autoregressive training requires sequential state updates and repeated memory access.The cumulative-sum dependency can make practical accelerator performance substantially worse than theoretical complexity suggests.
  • Mixed Chunk Attention: Mixed chunk attention combines local quadratic attention within chunks with global linear attention across chunks.The two components share the major GAU tensors, while local attention captures within-chunk interactions and global attention captures long-range interactions.
  • Mixed Chunk Attention: Local quadratic attention has complexity O(TCd), which is linear in sequence length when chunk size C remains constant.The input is divided into G non-overlapping chunks of size C before the two attention components are applied.
  • Fast Auto-regressive Training: Chunk-level summations reduce the number of causal cumulative-sum elements by a factor of C, with C typically 256 in the experiments.This reduces the sequential dependency from T steps to G = T/C steps and supports faster autoregressive training.
  • Discussions: Overlapping local attention can improve quality but introduces memory re-formatting operations that harm actual running speed.The authors therefore identify non-overlapping local attention as a strong candidate, while noting that the optimal partial-attention variant may be task-specific.

4. Experiments

FLASH is evaluated against standard, augmented, and linear-complexity Transformer baselines across bidirectional and autoregressive language modeling. It maintains competitive quality while improving training speed, with ablations identifying the contributions of GAU, mixed chunk attention, and chunk size.

  • Evaluation setup: FLASH and FLASH-Quad are evaluated on bidirectional and autoregressive modeling across multiple datasets and context lengths.FLASH combines GAUs with mixed chunk attention, while FLASH-Quad uses a stack of GAUs.
  • Bidirectional language modeling: FLASH-Quad and FLASH reduce C4 training cost by 1.1×–2.5× and 1.0×–4.8×, respectively, when matching Transformer++’s final perplexity.Combiner, Performer, and FLASH maintain roughly constant per-step latency as context length increases.
  • Auto-regressive language modeling: FLASH cuts Wiki-40B compute cost by 1.2×–4.9× while reaching a similar perplexity as Transformer++.FLASH closely tracks FLASH-Quad’s perplexity and improves the perplexity-cost trade-off as context length increases.
  • Auto-regressive language modeling: FLASH reaches Transformer+’s final PG-19 perplexity at 8K context in 55K steps, yielding a 12.12× speedup.The comparison uses roughly 500M-parameter models.
  • Fine-tuning: Increasing FLASH-Quad’s head size from 128 to 512 improves TrivialQA F1 by 2.1 points with negligible speed impact.With additional changes, FLASH_s=512 reaches comparable quality to Transformer+ while running 2.8× faster in pretraining and 2.7× faster in fine-tuning.
  • Ablation studies: Removing either local quadratic or global linear attention creates a significant quality gap, indicating that both components are critical to mixed chunk attention.The full model outperforms FLASH (LocalOnly) and FLASH (GlobalOnly).
  • Ablation studies: FLASH outperforms MC-TFM++ by more than 2× at sequence lengths above 2048, supporting the importance of GAU in the design.Approximation causes substantially less quality loss with FLASH-Quad→FLASH than with TFM++→MC-TFM++.
  • Ablation studies: Larger chunk sizes generally perform better as context length increases; chunk size 512 is preferable to C=256 beyond context length 1024.Chunk size affects both quality and training cost.

5. Conclusion

The paper presents FLASH as a practical approach combining a performant gated layer with accelerator-efficient mixed chunk attention. Across bidirectional and autoregressive language modeling, it matches fully augmented Transformers in perplexity while training substantially faster.

  • Conclusion: FLASH combines a gated linear unit with mixed chunk attention to address quality and empirical speed issues in efficient Transformer variants.The paper evaluates the model on bidirectional and autoregressive language modeling tasks.
  • Conclusion: FLASH matches fully augmented Transformers in perplexity while training substantially faster.The paper identifies scaling laws and downstream-task performance as future work.

A. Connections to Combiner

FLASH’s chunked linear attention differs from Combiner in how much information each chunk retains and how chunk summaries are combined. Adding Combiner-style summary attention is possible but conflicts with the default constant decoding complexity.

  • Representation: Combiner compresses each chunk into one O(d) vector, whereas FLASH retains an O(sd) matrix, providing less compression and potentially more memory.The FLASH representation is described as s times larger than Combiner’s chunk summary.
  • Summary combination: Combiner combines chunk summaries with quadratic attention, while FLASH’s chunked linear attention uses a cumulative sum.Combiner-style summary attention can be added through an extra [T/C × T/C] matrix.
  • Design trade-off: The added summary-attention variant was helpful but requires storing and attending to all chunk summaries.Its autoregressive decoding complexity becomes O((C + T/C)d2), so it is excluded from the default configuration.

B. Experimental Setup

The experiments use standardized hyperparameter settings for masked and autoregressive language modeling. Models are trained and evaluated under fixed batch sizes and step counts while varying context length across the specified datasets.

  • Bidirectional language modeling: C4 masked-language-modeling hyperparameters are listed in Table 6, with a footnote excluding the vanilla Transformer from an applied setting.The supplied setup text identifies Table 6 as the MLM pretraining configuration.
  • Auto-regressive language modeling: Wiki-40B and PG-19 autoregressive-language-modeling hyperparameters are listed in Table 7, with the same vanilla-Transformer exception.The supplied setup text identifies Table 7 as the LM pretraining configuration.

B.2. Model Specifications

The experiments use specified model configurations for C4 masked language modeling and related FLASH variants. FLASH-Quad and FLASH combine attention and feed-forward processing in a GAU module, using SiLU/Swish as their nonlinearity.

  • SiLU/Swish is used as the nonlinearity for FLASH-Quad and FLASH because it slightly outperforms GELU in these models.
  • Table 8 provides the model configurations for masked language modeling experiments on C4.
  • FLASH-Quad and FLASH combine attention and feed-forward networks into one module named GAU.

C. Additional Experimental Results

Additional experiments report model configurations and training-step latency, including comparisons involving FLASH and Performer. The results emphasize the implementation and measurement settings used for these experiments.

  • Table 9 lists model configurations for language-modeling experiments on Wiki-40B.
  • The listed model descriptions include FLASH variants with 12 attention layers and 12 feed-forward layers, or 36 attention layers and 36 feed-forward layers.
  • Table 10 lists model configurations for language-modeling experiments on PG-19.
  • Table 11 compares auto-regressive language-modeling latency per training step on Wiki-40B using a single Nvidia Tesla V100 GPU.Latency is reported in milliseconds, and OOM denotes CUDA out-of-memory errors.

C.1. Auto-regressive Training on GPU

On a single Nvidia V100 GPU, Performer has the largest latency among the compared models, while FLASH has the lowest latency beyond context length 1024.

  • FLASH achieves the lowest latency when context length exceeds 1024, whereas Performer has the largest latency among the three models.The passage attributes Performer’s latency to sequential cumsum over all tokens and associates FLASH’s result with mixed chunk attention.

C.2. Tabular MLM and LM Results

The tabular results summarize masked and auto-regressive language-modeling performance using validation perplexity and training latency across C4 and Wiki-40B.

  • The experiments summarize masked language modeling on C4 and language modeling on Wiki-40B in Tables 12 and 13.
  • Table 12 reports bidirectional or masked language-model results on C4, including best validation perplexity and training latency on 64 TPU-v4 cores.
  • Table 13 reports auto-regressive language-model results on Wiki-40B, including best validation perplexity and training latency on 64 TPU-v4 cores.

C.3. Ablation Study of Chunk Size

FLASH’s chunk size trades off quality against training cost. Extremely large chunks lose long-context scalability, while unit-sized chunks reduce FLASH to linear attention with inefficient autoregressive training.

  • C.3. Ablation Study of Chunk Size: Chunk size affects both FLASH’s quality and training cost across context lengths from 1K to 8K.Figure 8 compares four chunk sizes over this context-length range.
  • C.3. Ablation Study of Chunk Size: When chunk size equals the context length, FLASH falls back to FLASH-Quad and loses scalability to long context lengths.
  • C.3. Ablation Study of Chunk Size: When chunk size equals one, FLASH becomes linear attention, which suffers from inefficient autoregressive training.
  • C.3. Ablation Study of Chunk Size: FLASH combines quadratic and linear attention components within chunked processing.The implementation generates separate query and key representations, computes both attention parts, and combines their outputs.
Loading 2202.10447v2…