PaperlayerView exact revision on arXiv

Source-linked AI summary

A$^3$: Accelerating Attention Mechanisms in Neural Networks with Approximation

Tae Jun Ham, Sung Jun Jung, Seonghak Kim, Young H. Oh, Yeonhong Park, Yoonho Song, Jung-Hun Park, Sanghee Lee, Kyoung Park, Jae W. Lee, Deog-Kyoon Jeong

arXiv:2002.10941v1cs.DCcs.LG

TL;DR

Attention mechanisms are widely useful for retrieving relevant information, but dense matrix operations perform unnecessary work across many search targets. A3 combines approximate candidate selection with specialized parallel hardware, achieving multiple orders of magnitude improvements in energy efficiency and substantial speedup over conventional hardware.

  • Problem

    Dense matrix-vector implementations treat attention's content-based search as exhaustive computation, despite many near-zero or low-impact weights and limited prior accelerator focus on attention mechanisms.

  • Method

    A3 combines approximate candidate selection that avoids exhaustive search with a specialized parallel hardware pipeline for attention mechanisms.

  • Results

    A3 achieves multiple orders of magnitude improvements in energy efficiency and substantial speedup over conventional hardware.

  • Takeaways & Limitations

    The algorithm-hardware co-design enables neural networks with attention mechanisms to use larger external knowledge bases or longer data sequences.

Abstract

from arXiv · show

With the increasing computational demands of neural networks, many hardware accelerators for the neural networks have been proposed. Such existing neural network accelerators often focus on popular neural network types such as convolutional neural networks (CNNs) and recurrent neural networks (RNNs); however, not much attention has been paid to attention mechanisms, an emerging neural network primitive that enables neural networks to retrieve most relevant information from a knowledge-base, external memory, or past states. The attention mechanism is widely adopted by many state-of-the-art neural networks for computer vision, natural language processing, and machine translation, and accounts for a large portion of total execution time. We observe today's practice of implementing this mechanism using matrix-vector multiplication is suboptimal as the attention mechanism is semantically a content-based search where a large portion of computations ends up not being used. Based on this observation, we design and architect A3, which accelerates attention mechanisms in neural networks with algorithmic approximation and hardware specialization. Our proposed accelerator achieves multiple orders of magnitude improvement in energy efficiency (performance/watt) as well as substantial speedup over the state-of-the-art conventional hardware.

I. INTRODUCTION

Attention mechanisms are important but computationally costly because conventional dense operations exhaustively compare search targets. A3 combines approximation with specialized hardware to reduce computation and improve speed and energy efficiency.

  • Attention mechanisms let neural networks retrieve relevant information from past states or external memory through content-based similarity search.
  • Dense matrix-vector implementations spend computation across all search targets, even though many have little impact on the final output.
  • A3 reduces attention computation by selecting potentially relevant candidates instead of exhaustively searching all targets.
  • A3 uses a specialized parallel hardware pipeline and datapath to accelerate the approximated attention mechanism.
  • Multiple orders of magnitude speedup and energy-efficiency improvements are reported over conventional hardware.
  • The specialized approximation hardware achieves higher speedup and energy efficiency while minimizing model-accuracy degradation.

II. BACKGROUND AND MOTIVATION

Attention computes similarities, normalizes them into weights, and forms a weighted sum, but its cost grows with the search space. Measurements show that attention occupies a substantial share of inference runtime.

  • A. Attention Mechanism: Attention computes dot-product similarities, applies softmax, and uses normalized scores to form a weighted sum from value vectors.
  • A. Attention Mechanism: In bAbI QA, attention identifies statements relevant to a question by comparing their vector embeddings.
  • B. Cost of Attention Mechanism: An attention operation requires work that increases almost linearly with both n, the number of data items, and d, the vector dimension.
  • B. Cost of Attention Mechanism: Over 35% of inference runtime is attributable to attention mechanisms across the evaluated workloads.
  • B. Cost of Attention Mechanism: A3 addresses this runtime cost with specialized hardware for the attention mechanism.

C. Opportunity for Approximation

Conventional attention implementations perform dense computation even though most softmax weights become near-zero and contribute little to the final output. A3 exploits this structure by selecting likely high-score candidates before computing the full attention pipeline.

  • C. Opportunity for Approximation: Dense matrix-vector attention computes scores for all candidates, although most have little impact after softmax normalization.Softmax acts as a differentiable approximation of argmax, producing near-zero weights for relatively low scores.
  • C. Opportunity for Approximation: A3 preprocesses the key matrix to identify candidate rows likely to have high scores without computing every dot product.This candidate set can reduce dot-product, softmax, and final weighted-sum computations.

III. A3 BASE DESIGN

The base A3 is a specialized accelerator that implements attention with a pipelined datapath. It accepts key, value, and query inputs and processes them through dot-product, exponent, and output-computation modules.

  • III. A3 BASE DESIGN: Base A3 takes a key matrix (n × d), a value matrix (n × d), and a query vector (d) to produce a d-dimensional output vector.The design reorders some computations relative to the reference pseudocode for hardware implementation.
  • III. A3 BASE DESIGN: The base A3 pipeline contains dot-product, exponent-computation, and output-computation modules.Figure 4 depicts the block diagram, while Figure 5 represents the pipeline computation in pseudocode.
  • III. A3 BASE DESIGN: The dot-product module computes each key-row/query-vector inner product using d multipliers and a d-way adder tree.Rows are loaded sequentially, multiplied element-wise, and reduced through parallel summation.
  • III. A3 BASE DESIGN: The exponent-computation module calculates exponentials of dot-product values using a lookup table rather than a dedicated exponent unit.The design also addresses overflow and lookup-table size in subsequent hardware transformations.
  • III. A3 BASE DESIGN: The pipeline latency is 3n + 27 cycles and its throughput is n + 9 cycles per query.The three modules are designed for matching throughput, with the last module having latency n + 9.

B. Quantization

A3 uses lower-precision fixed-point representations because neural-network tasks can tolerate some error, while assigning stage-specific bitwidths to preserve precision and avoid overflow. Its design also supports preloaded matrices and parallel accelerator units.

  • B. Quantization: A3 quantizes floating-point inputs into signed fixed-point representations with integer and fractional bits, using different bitwidths across pipeline stages.The stated goal is to maintain precision and avoid overflow while reducing energy cost.
  • B. Quantization: The output representation uses 3f fraction bits, combining 2f bits for weights with f bits for values.The passage states that score[] and weight[] each use 2f fraction bits.
  • C. Design Details: A3 copies key and value matrices into its SRAM buffer before query processing, so query response time includes only host-to-accelerator query transfer.The matrices are available at knowledge-comprehension time in the described question-answering setting.
  • C. Design Details: Multiple A3 units can process independent attention computations in parallel for workloads with substantial parallelism.The passage presents replicated A3 units as an option for parallel attention processing.
  • B. Quantization: Quantization error becomes smaller after exponentiation when the exponent part is negative.The passage supports this property through separate inequalities for positive and negative quantization errors.

IV. APPROXIMATE ATTENTION

A3 approximates attention by exploiting the fact that most weights are near-zero. It preprocesses the key matrix to select likely high-score candidates, reducing query-time computation while amortizing preprocessing across repeated queries.

  • IV. APPROXIMATE ATTENTION: Attention computes similarity scores between a query and all key rows, applies softmax, and uses the normalized scores to form a weighted value sum.This is the conventional content-based attention pipeline.
  • IV. APPROXIMATE ATTENTION: Most softmax weights are often near-zero and can be treated as zeros to avoid unnecessary softmax and weighted-sum computation.The passage attributes this pattern to softmax amplifying differences between a few large scores and smaller entries.
  • IV. APPROXIMATE ATTENTION: The approximation algorithm selects candidates likely to have high scores without directly computing their scores.It preprocesses the key matrix to reduce operations during question answering.
  • IV. APPROXIMATE ATTENTION: For models such as BERT, preprocessing overhead is amortized when multiple queries reuse the same key matrix.The passage gives multiple queries, such as 320, as an example of reuse.

B. Base Greedy Candidate Search

The base greedy search approximates dot-product scores by iteratively accumulating extreme component-wise products, then selecting rows with positive greedy scores as candidates. A preprocessing-based implementation reduces query-time complexity to depend on iteration count and dimension rather than the number of rows.

  • Base greedy candidate search: Large positive or negative component-wise products provide signals about whether a row’s final dot product is likely large or not.The method uses this assumption to avoid computing every full inner product.
  • Base greedy candidate search: The algorithm iteratively adds the kth largest and smallest component products to a greedy score array for M iterations.Afterward, rows with positive greedy scores are selected as candidates.
  • Base greedy candidate search: O(nd log nd) is the time complexity of the initial greedy search because selecting each extreme element requires sorting the result matrix.This is not useful relative to O(nd) full dot-product computation, motivating preprocessing.
  • Efficient greedy candidate search: Preprocessing sorts each key-matrix column and stores values with row indices, moving work off query response time when possible.For repeated queries using the same key matrix, preprocessing cost can be amortized.
  • Efficient greedy candidate search: The implementation updates max_ptr and min_ptr, inserts newly pointed entries into priority queues, and repeats selection for M iterations.A heuristic skips minQ when the cumulative selected sum is negative.
  • Efficient greedy candidate search: The restructured candidate-selection algorithm has complexity M log d and exposes M as a performance–accuracy control.Accuracy may require increasing M as the number of rows N increases.

D. Post-scoring Approximation

After candidate selection, A3 uses full scores and softmax to approximate attention while choosing how many high-scoring rows to retain. The post-scoring selection adapts to score distributions rather than using one fixed row count.

  • Post-scoring approximation: Candidate rows receive full dot-product scores, which feed softmax and the final weighted sum; low scores often become near-zero after softmax.This motivates avoiding unnecessary computation for low-impact rows.
  • Post-scoring approximation: A static number of top-scoring rows is unsuitable across cases because low-variance high scores may require retaining many rows, whereas one dominant score may require few.The selection count should reflect the distribution of scores.
  • Post-scoring approximation: A3 adds hardware modules for candidate selection and post-scoring selection, connecting them to the base attention accelerator.The modules implement the proposed approximation scheme in hardware.

A. Candidate Selection Module

The candidate selection module implements the efficient greedy search with buffered, preprocessed key data and specialized parallel hardware. A separate post-scoring module filters dot-product results before exponent computation.

  • Candidate selection module: The candidate selection module uses customized hardware and parallelism to reduce execution time and improve energy efficiency.It accelerates the algorithm described in Figure 7.
  • Candidate selection module: SRAM stores each key-matrix column in sorted order together with the original row indices.The module also uses max_ptr and min_ptr registers, multipliers, circular queues, and comparator trees.
  • Candidate selection module: The post-scoring selection module identifies top entries from dot-product results before exponent computation by comparing remaining values with the maximum.Sixteen subtractors and comparators process 16 entries per cycle, passing values meeting threshold t onward.

C. A3 HW with Approximation Support

A3 combines candidate selection, dot-product computation, post-scoring selection, exponent computation, and weighted-sum output computation into an approximation-supported pipeline. Experiments evaluate accuracy across three attention-based workloads and approximation configurations.

  • A3 hardware with approximation support: The pipeline’s latency is M + C + K + K + α cycles, with throughput limited by candidate selection at approximately M cycles.C is the number of candidates, K the post-scoring selections, and α a constant.
  • A3 hardware with approximation support: A3 passes preprocessed-key candidates through dot-product and post-scoring modules before exponent and weighted-sum computation.The post-scoring selector determines which rows enter exponent computation.
  • Evaluation: The evaluation uses MemN2N on bAbI QA, a Key-Value Memory Network on SQuAD v1.1, and BERT among the target attention workloads.The supplied workload description identifies bAbI QA and SQuAD evaluation settings.
  • Evaluation: d = 64 is used for all workloads, while n averages 20 for bAbI QA, reaches 186 on Wikimovies, and is 320 for SQuAD.These workloads therefore cover different search sizes.
  • Accuracy evaluation: Varying candidate-selection iterations M changes accuracy because it changes the number of selected candidates across the three workloads.Figure 11 evaluates this trade-off over different M values.
  • Accuracy evaluation: A relatively high post-scoring threshold T, such as 10%, can still achieve decent accuracy while excluding entries with very small post-softmax scores.Higher T denotes more aggressive approximation and lower T more conservative approximation.
  • Accuracy evaluation: The conservative combined approximation configuration uses M = 1/2n and T = 5% and loses around 1% accuracy.The aggressive configuration is described as exploratory because of its relatively high accuracy decrease.

C. Performance Results

A3 delivers substantially higher attention throughput and lower latency than conventional hardware for several workloads. Approximation further reduces computation and improves throughput and latency, while BERT benefits from replicated A3 units.

  • Evaluation: Attention throughput is evaluated with a 1GHz A3 cycle-level simulator against Intel Xeon Gold 6128 CPU and, for BERT, NVIDIA Titan V GPU.Throughput is normalized to the CPU for each workload.
  • Throughput: Orders of magnitude higher throughput than Intel Xeon CPU is achieved on MemN2N and KV-MemN2N by both base and approximate A3.For BERT, A3 initially trails the GPU because self-attention is a batch matrix-matrix multiplication with easy-to-exploit parallelism.
  • Throughput: 103-104× lower energy consumption than CPU or GPU per A3 unit enables multiple units to improve throughput.Using 67 approximate conservative A3 units can surpass the state-of-the-art GPU for BERT.
  • Throughput: Approximation increases throughput relative to base A3, although the gain is smaller for MemN2N because its n is relatively small.Different module counts could potentially address this workload-specific limitation.
  • Latency: Both approximation-enabled A3 configurations achieve significantly lower latency than base A3 by performing substantially fewer computations.Aggressive approximation offers more speedup than conservative approximation at a relatively low or moderate accuracy cost.

D. Area, Power, Energy and Test Chip

A3 is implemented as a compact 40nm accelerator and evaluated through synthesized hardware characteristics, energy breakdowns, and a scaled-down test chip. Approximation changes which hardware module dominates energy use.

  • Implementation: A3 is synthesized for a 1GHz clock using Chisel-generated Verilog, Synopsys Design Compiler, and a TSMC 40nm standard-cell library.The design underwent functional verification before synthesis.
  • Energy: Energy efficiency and energy breakdown are compared across base A3, conservative approximate A3, and aggressive approximate A3.The three configurations are represented from left to right in the energy-breakdown bars.
  • Area: Less than 2.082mm2 area is used by A3, versus 325mm2 for the 14nm CPU baseline and 815mm2 for the 12nm GPU baseline.The CPU and GPU baselines are reported as 156× and 391× larger than a single A3 unit, respectively.
  • Energy: Approximate A3 concentrates most energy in candidate selection because reducing processed rows leaves other modules less heavily utilized.Base A3 instead spends most energy in the output computation module because of its large register structures.
  • Test Chip: A scaled-down A3 test chip was taped out in TSMC 40nm Low Power technology with standard cells.

VII. RELATED WORK

Related work covers attention mechanisms across language, vision, and long-term-memory applications, alongside approximate similarity-search methods. A3 differs by presenting hardware support for approximated attention.

  • Attention Mechanism: Attention mechanisms support tasks including translation, question answering, language inference, summarization, document classification, image understanding, and long-term memory.The paper states that its work applies to many of these neural-network models.
  • Approximate Similarity Search: Prior approximate similarity-search approaches use hashing, trees, clustering, or greedy iterative search to avoid exhaustive linear search.Some prior work applies approximate similarity search to attention mechanisms in software.
  • Attention Mechanism: A3 accelerates the relatively less optimized attention primitive through software-hardware co-design and reports orders-of-magnitude energy-efficiency improvement over conventional hardware.
Loading 2002.10941v1…