PaperlayerView exact revision on arXiv

Source-linked AI summary

SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

Hanrui Wang, Zhekai Zhang, Song Han

arXiv:2012.09852v3cs.ARcs.AIcs.CLcs.LG

TL;DR

Attention inference is inefficient because of quadratic complexity, complex data movement, and low arithmetic intensity, while existing accelerators do not efficiently support attention. SpAtten co-designs sparse and quantized attention hardware and algorithms, reducing computation and memory access; across benchmarks, it reduces DRAM access by 10.0× and outperforms prior accelerators and general-purpose platforms.

  • Problem

    Attention becomes an inference bottleneck because its computation scales quadratically with input length, data movement is complex, and arithmetic intensity is low.

  • Method

    SpAtten uses cascade token and head pruning, progressive quantization, and an on-the-fly top-k engine to reduce sparse and quantized attention computation and memory access.

  • Results

    10.0× average DRAM-access reduction accompanies 2.1× lower computation, while SpAtten is 1.6× and 3.0× faster than A3 and MNNFast, respectively.

  • Takeaways & Limitations

    SpAtten provides an accelerator design for efficient sparse and quantized attention inference across computation-bounded BERT and memory-bounded GPT-2 workloads.

Abstract

from arXiv · show

The attention mechanism is becoming increasingly popular in Natural Language Processing (NLP) applications, showing superior performance than convolutional and recurrent architectures. However, attention becomes the compution bottleneck because of its quadratic computational complexity to input length, complicated data movement and low arithmetic intensity. Moreover, existing NN accelerators mainly focus on optimizing convolutional or recurrent models, and cannot efficiently support attention. In this paper, we present SpAtten, an efficient algorithm-architecture co-design that leverages token sparsity, head sparsity, and quantization opportunities to reduce the attention computation and memory access. Inspired by the high redundancy of human languages, we propose the novel cascade token pruning to prune away unimportant tokens in the sentence. We also propose cascade head pruning to remove unessential heads. Cascade pruning is fundamentally different from weight pruning since there is no trainable weight in the attention mechanism, and the pruned tokens and heads are selected on the fly. To efficiently support them on hardware, we design a novel top-k engine to rank token and head importance scores with high throughput. Furthermore, we propose progressive quantization that first fetches MSBs only and performs the computation; if the confidence is low, it fetches LSBs and recomputes the attention outputs, trading computation for memory reduction. Extensive experiments on 30 benchmarks show that, on average, SpAtten reduces DRAM access by 10.0x with no accuracy loss, and achieves 1.6x, 3.0x, 162x, 347x speedup, and 1,4x, 3.2x, 1193x, 4059x energy savings over A3 accelerator, MNNFast accelerator, TITAN Xp GPU, Xeon CPU, respectively.

I. INTRODUCTION

Attention improves NLP performance but is inefficient because of quadratic scaling, costly data movement, and low arithmetic intensity. SpAtten co-designs adaptive token and head pruning, progressive quantization, and specialized hardware to reduce attention costs while preserving accuracy.

  • Motivation: Attention becomes a major bottleneck because its cost scales quadratically with input length, involves complex data movement, and has low arithmetic intensity.Attention accounts for over 50% of GPT-2 latency, with data movements contributing 73% of attention latency on TITAN Xp GPU.
  • Approach: SpAtten applies cascade token pruning, cascade head pruning, and progressive quantization to reduce computation and memory access adaptively.Pruned tokens and heads are selected on the fly and removed from all following layers; quantization adapts bitwidths to attention probability distributions.
  • Approach: Progressive quantization first computes with most-significant bits and fetches least-significant bits only when attention probabilities are too flat for reliable results.This trades additional computation for reduced memory access and saves another 5.1× memory access.
  • Approach: Cascade token pruning removes unimportant tokens globally across subsequent layers, reducing DRAM access and computation by up to 3.8×.Token importance is determined from cumulative attention-based scores rather than trainable weights.
  • Approach: Cascade head pruning removes unessential attention heads and provides another 1.1× reduction in DRAM access and computation.The method ranks head importance on the fly rather than using compile-time, input-independent pruning.
  • Evaluation: Across 30 benchmarks, SpAtten reduces DRAM access by 10.0× with no accuracy loss and accelerates inference over specialized accelerators and general-purpose platforms.Reported speedups are 1.6× over A3, 3.0× over MNNFast, 162× over TITAN Xp GPU, and 347× over Xeon CPU.

II. BACKGROUND AND MOTIVATION

Attention models summarize and generate sequences by computing multi-head interactions among Query, Key, and Value representations. The paper describes how token and head importance can be accumulated to prune redundant components across layers.

  • Attention mechanism: Attention splits Query, Key, and Value inputs into multiple heads, computes scaled Q × K^T scores, applies row-wise softmax, and multiplies probabilities by V.Head outputs are concatenated to produce the attention output.
  • Attention-based NLP models: BERT performs only summarization, whereas GPT-2 performs summarization followed by autoregressive generation using cached K and V representations.Generation processes one new token per iteration while combining it with prior keys and values.
  • Cascade pruning: Once tokens or heads are pruned, they remain removed in subsequent layers, so deeper layers process progressively fewer tokens and heads.This persistent behavior gives the method its cascade designation.
  • Cascade pruning: Cascade token pruning ranks tokens using cumulative attention-probability scores and removes low-importance tokens together with their Q, K, and V vectors.The pruning algorithm accumulates token scores and retains the top-k tokens according to the pruning ratio.
  • Cascade pruning: Cascade head pruning ranks heads using cumulative absolute attention-output values and removes low-importance heads with their corresponding Q, K, and V chunks.Token pruning reduces sentence length, while head pruning reduces feature length.

B. Motivation

Attention is a major efficiency bottleneck because it combines substantial latency with extensive data movement, while existing processors are poorly suited to those operations. SpAtten addresses this bottleneck with a dedicated co-processor and token-pruning mechanism.

  • Motivation: Attention accounts for over 50% of GPT-2 latency despite contributing around 10% of overall FLOPs.The profile identifies attention as disproportionately time-consuming.
  • Motivation: Around 73% of attention time is spent moving data for head splitting, K and V concatenation, reshape, and transpose operations.GPUs and CPUs are optimized for matrix multiplication but handle these memory operations poorly.
  • Token pruning: BERT token importance scores are obtained by summing attention probabilities over columns, and tokens with small scores are pruned.The figure illustrates the score-based selection used for token pruning.
  • Motivation: SpAtten uses a co-processor to handle attention layers while GPUs, CPUs, or tensor algebra accelerators process fully connected layers.The design separates attention processing from highly optimized fully connected computation.

A. Cascade Token Pruning

SpAtten removes redundant tokens and heads using importance scores, then applies progressive quantization based on attention-probability distributions to reduce computation and memory access.

  • Cascade Token Pruning: Cascade token pruning removes unimportant tokens according to cumulative attention-probability importance scores.Scores accumulate across attention rounds because larger probabilities indicate greater influence on attention outputs.
  • Cascade Token Pruning: Once pruned, a token’s Q, K, and V are excluded from all following attention heads and layers.This cascade process progressively reduces the tokens processed by deeper layers.
  • Cascade Head Pruning: Cascade head pruning removes redundant heads using cumulative absolute attention-output magnitudes to reduce feature length.Token pruning reduces sentence length, while head pruning reduces feature length.
  • Local Value Pruning: Local Value pruning discards low-probability V vectors only for the current head after Softmax.Unlike cascade token pruning, it does not remove Q, K, and V vectors from subsequent heads and layers.
  • Progressive Quantization: Progressive quantization fetches MSBs first and retrieves LSBs only when the resulting attention distribution is sufficiently flat.Flat distributions have larger quantization errors, whereas distributions dominated by a few tokens can use smaller bitwidths.

IV. HARDWARE ARCHITECTURE

SpAtten’s hardware combines top-k selection, memory-routing support, on-chip storage, and pipelined processing to implement sparse attention efficiently.

  • Architecture Overview: A top-k engine ranks token and head importance scores, while crossbars route random addresses across 16 HBM channels.The crossbar keeps memory channels busy and improves bandwidth utilization despite pruning-induced random access.
  • Quantization Support: Progressive quantization is supported by an on-chip bitwidth converter that splits fetched bits and concatenates MSBs with LSBs.The converter interfaces DRAM data with the fixed-width on-chip datapath.
  • Dataflow: SpAtten processes attention head by head and query by query, balancing pruning granularity with parallelism.Inner-head parallelism keeps on-chip computation resources busy without requiring inter-head parallelism.
  • Pipeline: The critical path is fully pipelined, while importance-score accumulation and top-k selection operate in parallel with it.This organization translates pruning and quantization savings into higher-throughput execution.
  • On-Chip Memory: Key and Value data use separate 196KB SRAMs, with address and data FIFOs buffering requests and fetched values.The Q vector is stored in registers because queries are processed one at a time.

B. Top-k Engine

The top-k engine avoids full sorting by thresholding with quick-select, filtering candidates, and eliminating zeros in parallel to produce high-throughput selections.

  • Top-k Selection: The engine finds the kth-largest value as a threshold, filters the input array, and then uses a zero eliminator to output the top-k elements.This avoids the O(n · log n) time and O(n · log2 n) space associated with a naïve sorting-network solution.
  • Hardware Organization: The architecture uses two top-k engines, reusing the token-pruning engine for the less frequent head-pruning operation.Sixteen comparators are used in each array so top-k selection does not bottleneck the pipeline.
  • Zero Elimination: A prefix-sum-guided log n-stage shifter removes zero entries by shifting elements according to their zero counts.Each stage shifts by 1, 2, 4, ... positions based on the corresponding bit of an element’s zero count.
  • Evaluation: 3× speedup is achieved over a baseline top-k engine with parallelism=1.The result demonstrates the throughput benefit of parallelized top-k processing.
  • Evaluation: 1.4× higher throughput and 3.5× lower power consumption are achieved than a full sorting unit in the reported comparison.The comparison uses input length 1024 and the top-k engine’s median-selection worst case.

E. Query-Key Multiplication Module

The query-key multiplication module computes attention scores with configurable parallelism, after which Softmax, progressive-quantization decisions, and value multiplication produce attention outputs.

  • Query-Key Multiplication: The query-key module computes each attention score as a dot product between a key row and the query vector.It uses 512 multipliers and configurable adder trees to generate multiple scores per cycle when the dimension is below 512.
  • Softmax: Attention scores are dequantized and normalized together, then passed through exponential, accumulation, and division stages to compute Softmax.The Softmax results are quantized again for subsequent fixed-point operations.
  • Architecture Setup: The architecture setup is specified in a dedicated table for SpAtten.The supplied passage identifies the table but does not provide its configuration values.
  • Progressive Quantization: The progressive-quantization module fetches LSBs when the maximum attention probability falls below a predefined threshold.This condition indicates a flat distribution and triggers recomputation of attention probabilities.
  • Attention Output: The attention prob-value unit multiplies attention probabilities by V vectors and accumulates the results into attention outputs.A broadcast-multiply-reduce pipeline with 512 multipliers supports multiple probabilities simultaneously.

V. EVALUATION

SpAtten is evaluated across 30 benchmarks and multiple hardware baselines, including GPUs, CPUs, and attention accelerators. Its co-designed extensions reduce computation, memory access, latency, and model size across attention and end-to-end Transformer workloads.

  • Evaluation Setup: 30 benchmarks cover BERT and GPT-2 models across discriminative and generative NLP tasks.The evaluation includes GLUE, SQuAD, and four language-modeling datasets.
  • Attention Results: 10.0× lower DRAM access and 2.1× lower computation are achieved on average with SpAtten.Token and value pruning average 1.9× reduction across models, while head pruning provides 1.1× reduction.
  • Accelerator Comparisons: 1.6× better throughput, 1.4× better energy efficiency, and 2.2× better area efficiency are achieved over A3.The comparison uses SpAtten1/8 and matched multiplier count, technology, and bandwidth.
  • Accelerator Comparisons: 3.0× higher throughput and 3.2× better energy efficiency are achieved over MNNFast.These results compare SpAtten1/8 with an optimistic ASIC power estimate for MNNFast.
  • End-to-End Results: 8-bit FC SpAtten-e2e achieves average speedups of 35× over TITAN Xp GPU and 122× over Xeon CPU on eight GPT-2-Medium benchmarks.The corresponding 12-bit configuration achieves 24× and 83× speedups, respectively.
  • Co-designed Transformers: The co-designed Transformer-7 achieves 1.9× faster speed and 2.8× smaller size than the vanilla Transformer-Big model.The co-designed Transformer-3 and vanilla Transformer-Base have similar accuracy.

C. Performance Analysis

Performance analysis attributes SpAtten’s gains to cascade pruning, a high-throughput top-k engine, and progressive quantization. The design also exposes accuracy-efficiency trade-offs and operates near bandwidth and computation roofs.

  • Speedup Breakdown: 22.1× speedup over the GPU baseline is obtained with a dedicated datapath on eight GPT-2 benchmarks.The baseline performance is limited by numerous memory instructions for attention.
  • Speedup Breakdown: 3.8× computation reduction from token pruning and 1.1× from head pruning produce only 1.1× performance gains before top-k acceleration.The top-k selection overhead becomes a bottleneck without a high-throughput engine.
  • Speedup Breakdown: 2.8× additional speedup comes from progressive quantization with less DRAM access.The method reduces average input bitwidth after the top-k bottleneck is resolved.
  • Roofline Analysis: SpAtten’s performance points are close to bandwidth and computation roofs, while progressive quantization shifts them toward higher operational intensity.The analysis compares SpAtten with TITAN Xp GPU using an HBM bandwidth roof of 512G and a 2TFLOPS computation roof.
  • Design Space: A top-k engine with parallelism 16 and a 196KB key/value buffer is sufficient for the reported settings.Increasing parallelism beyond 16 provides little additional performance because it matches the Q×K input data rate.
  • Efficiency-Accuracy Trade-offs: Without accuracy loss, token pruning averages 1.9× and head pruning averages 1.1× across benchmarks.Task-specific curves show approximately 4× token pruning for GPT-2 on PTB and 1.2× head pruning for BERT on CoLA.
  • Interpretation and Visualization: Cascade token pruning preserves interpretable task-relevant tokens in both discriminative BERT and generative GPT-2 examples.Examples retain sentiment-bearing, similarity-related, or generation-relevant tokens while pruning redundant ones.
  • Interpretation and Visualization: GPT-2 cumulative importance scores show that important tokens can remain consistent across layers and include both nearby and distant dependencies.Generated tokens attend heavily to selected local and long-range context tokens.

A. Neural Networks Pruning and Quantization

SpAtten differs from conventional neural-network pruning by targeting activation structure in attention rather than trainable weights. It combines cascade token and head pruning with hardware-supported ranking and progressive quantization.

  • Prior Pruning Approaches: SpAtten targets token and head activation sparsity, whereas most prior accelerator work focuses on weight sparsity.Specialized top-k engines support on-the-fly activation selection.
  • Token Pruning: Cascade token pruning removes redundant tokens across subsequent layers using cumulative importance rather than a single layer’s attention probabilities.This distinguishes SpAtten from PoWER-BERT’s one-layer token pruning.
  • Token Pruning: SpAtten supports per-head token pruning and local V-vector pruning, capabilities not supported by PoWER-BERT.The comparison identifies these as differences from the prior method.
  • Head Pruning: Cascade head pruning removes unessential head activations, unlike prior approaches that prune head weights.Pruned heads are selected dynamically as part of the attention inference process.
  • Hardware Support: The top-k engine ranks token and head importance in O(n) time for on-the-fly pruning.This hardware mechanism translates sparse attention opportunities into accelerator support.
  • Quantization: Progressive quantization supports different bitwidths across layers while reducing attention memory access.The method is presented alongside sparse pruning as part of SpAtten’s quantized attention inference design.
Loading 2012.09852v3…