Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LATTE: Low-Precision Approximate Attention with Head-wise Trainable Threshold for Efficient Transformer

Jiing-Ping Wang, Ming-Guang Lin, An-Yeu, Wu

TL;DR

LATTE addresses the quadratic cost of multi-head attention by exploiting sparsity with a head-wise trainable threshold and low-precision approximate attention. It uses 8-bit quantization, MS4B-based dot-product estimation, and computation reuse to prune unimportant keys while enabling end-to-end optimization. The method achieves substantial pruning (up to about 85–90% keys) with only minor performance degradation on CV and NLP benchmarks, demonstrating practical efficiency gains for both language and vision tasks. This work advances efficient transformer inference by coupling differential thresholds with per-head adaptability and cost-aware computation reuse.

Abstract

With the rise of Transformer models in NLP and CV domain, Multi-Head Attention has been proven to be a game-changer. However, its expensive computation poses challenges to the model throughput and efficiency, especially for the long sequence tasks. Exploiting the sparsity in attention has been proven to be an effective way to reduce computation. Nevertheless, prior works do not consider the various distributions among different heads and lack a systematic method to determine the threshold. To address these challenges, we propose Low-Precision Approximate Attention with Head-wise Trainable Threshold for Efficient Transformer (LATTE). LATTE employs a headwise threshold-based filter with the low-precision dot product and computation reuse mechanism to reduce the computation of MHA. Moreover, the trainable threshold is introduced to provide a systematic method for adjusting the thresholds and enable end-to-end optimization. Experimental results indicate LATTE can smoothly adapt to both NLP and CV tasks, offering significant computation savings with only a minor compromise in performance. Also, the trainable threshold is shown to be essential for the leverage between the performance and the computation. As a result, LATTE filters up to 85.16% keys with only a 0.87% accuracy drop in the CV task and 89.91% keys with a 0.86 perplexity increase in the NLP task.

LATTE: Low-Precision Approximate Attention with Head-wise Trainable Threshold for Efficient Transformer

TL;DR

LATTE addresses the quadratic cost of multi-head attention by exploiting sparsity with a head-wise trainable threshold and low-precision approximate attention. It uses 8-bit quantization, MS4B-based dot-product estimation, and computation reuse to prune unimportant keys while enabling end-to-end optimization. The method achieves substantial pruning (up to about 85–90% keys) with only minor performance degradation on CV and NLP benchmarks, demonstrating practical efficiency gains for both language and vision tasks. This work advances efficient transformer inference by coupling differential thresholds with per-head adaptability and cost-aware computation reuse.

Abstract

With the rise of Transformer models in NLP and CV domain, Multi-Head Attention has been proven to be a game-changer. However, its expensive computation poses challenges to the model throughput and efficiency, especially for the long sequence tasks. Exploiting the sparsity in attention has been proven to be an effective way to reduce computation. Nevertheless, prior works do not consider the various distributions among different heads and lack a systematic method to determine the threshold. To address these challenges, we propose Low-Precision Approximate Attention with Head-wise Trainable Threshold for Efficient Transformer (LATTE). LATTE employs a headwise threshold-based filter with the low-precision dot product and computation reuse mechanism to reduce the computation of MHA. Moreover, the trainable threshold is introduced to provide a systematic method for adjusting the thresholds and enable end-to-end optimization. Experimental results indicate LATTE can smoothly adapt to both NLP and CV tasks, offering significant computation savings with only a minor compromise in performance. Also, the trainable threshold is shown to be essential for the leverage between the performance and the computation. As a result, LATTE filters up to 85.16% keys with only a 0.87% accuracy drop in the CV task and 89.91% keys with a 0.86 perplexity increase in the NLP task.
Paper Structure (15 sections, 9 equations, 6 figures, 1 algorithm)

This paper contains 15 sections, 9 equations, 6 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Related Works
  3. Transformer Models
  4. Related Works of Sparse Attention with Pruning
  5. Proposed Algorithm
  6. Low-Precision Approximated Attention
  7. Differential Threshold
  8. Low-Precision Estimation
  9. Computation Reuse
  10. Head-wise Trainable Threshold
  11. Experimental Results
  12. Comparison with Prior Work
  13. Approximate Dot Product
  14. The Necessity of Trainable Threshold
  15. Conclusion

Figures (6)

  • Figure 1: Comparison between (a) prior work and (b) the proposed LATTE
  • Figure 2: Overview of LATTE algorithm. ① The key vectors are filtered based on the low-precision dot product, and ② the partial sum of retained keys is reused for the reduced dot product.
  • Figure 3: Attention probabilities among different blocks and heads of DeiTtouvron2021training
  • Figure 4: Performance comparison between our work and prior work for (a) CV task and (b) NLP task.
  • Figure 5: Performance comparison before and after approximation for (a) CV task and (b) NLP task.
  • ...and 1 more figures