Online Pseudo-average Shifting Attention(PASA) for Robust Low-precision LLM Inference: Algorithms and Numerical Analysis

• URL: http://arxiv.org/abs/2503.01873v1
• Date: Wed, 26 Feb 2025 01:00:46 GMT
• Title: Online Pseudo-average Shifting Attention(PASA) for Robust Low-precision LLM Inference: Algorithms and Numerical Analysis
• Authors: Long Cheng, Qichen Liao, Fan Wu, Junlin Mu, Tengfei Han, Zhe Qiu, Lianqiang Li, Tianyi Liu, Fangzheng Miao, Keming Gao, Liang Wang, Zhen Zhang, Qiande Yin,
• Abstract summary: We develop a low-precision, mathematically-equivalent algorithm called PASA, based on Flash Attention.<n> PASA introduces two novel techniques: online pseudo-average shifting and global recovering.<n>We find that the large bias and amplitude of attention input data are critical factors contributing to numerical overflow.
• Score: 15.71443217369106
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Attention calculation is extremely time-consuming for long-sequence inference tasks, such as text or image/video generation, in large models. To accelerate this process, we developed a low-precision, mathematically-equivalent algorithm called PASA, based on Flash Attention. PASA introduces two novel techniques: online pseudo-average shifting and global recovering. These techniques enable the use of half-precision computation throughout the Flash Attention process without incurring overflow instability or unacceptable numerical accuracy loss. This algorithm enhances performance on memory-restricted AI hardware architectures, such as the Ascend Neural-network Processing Unit(NPU), by reducing data movement and increasing computational FLOPs. The algorithm is validated using both designed random benchmarks and real large models. We find that the large bias and amplitude of attention input data are critical factors contributing to numerical overflow ($>65504$ for half precision) in two different categories of large models (Qwen2-7B language models and Stable-Video-Diffusion multi-modal models). Specifically, overflow arises due to the large bias in the sequence dimension and the resonance mechanism between the query and key in the head dimension of the Stable-Video-Diffusion models. The resonance mechanism is defined as phase coincidence or 180-degree phase shift between query and key matrices. It will remarkably amplify the element values of attention score matrix. This issue also applies to the Qwen models. Additionally, numerical accuracy is assessed through root mean square error (RMSE) and by comparing the final generated texts and videos to those produced using high-precision attention.

Related papers

• Towards Learning High-Precision Least Squares Algorithms with Sequence Models [42.217390215093516]
  We show that softmax Transformers struggle to perform high-precision multiplications. We identify limitations present in existing architectures and training procedures. For the first time, we demonstrate the ability to train to near machine precision.
  arXiv  Detail & Related papers  (2025-03-15T23:25:11Z)
• Towards Scalable and Deep Graph Neural Networks via Noise Masking [59.058558158296265]
  Graph Neural Networks (GNNs) have achieved remarkable success in many graph mining tasks. scaling them to large graphs is challenging due to the high computational and storage costs. We present random walk with noise masking (RMask), a plug-and-play module compatible with the existing model-simplification works.
  arXiv  Detail & Related papers  (2024-12-19T07:48:14Z)
• Accelerated zero-order SGD under high-order smoothness and overparameterized regime [79.85163929026146]
  We present a novel gradient-free algorithm to solve convex optimization problems. Such problems are encountered in medicine, physics, and machine learning. We provide convergence guarantees for the proposed algorithm under both types of noise.
  arXiv  Detail & Related papers  (2024-11-21T10:26:17Z)
• SIGMA:Sinkhorn-Guided Masked Video Modeling [69.31715194419091]
  Sinkhorn-guided Masked Video Modelling ( SIGMA) is a novel video pretraining method. We distribute features of space-time tubes evenly across a limited number of learnable clusters. Experimental results on ten datasets validate the effectiveness of SIGMA in learning more performant, temporally-aware, and robust video representations.
  arXiv  Detail & Related papers  (2024-07-22T08:04:09Z)
• Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers [4.674454841332859]
  Transformer-based models have emerged as one of the most widely used architectures for natural language processing.<n>These huge models are memory hungry and incur significant inference latency even on cutting edge AI-accelerators.<n>We propose LeanAttention, a scalable technique of computing self-attention for the token-generation phase.
  arXiv  Detail & Related papers  (2024-05-17T00:52:39Z)
• A machine learning approach to the prediction of heat-transfer coefficients in micro-channels [4.724825031148412]
  The accurate prediction of the two-phase heat transfer coefficient (HTC) is key to the optimal design and operation of compact heat exchangers. We use a multi-output Gaussian process regression (GPR) to estimate the HTC in microchannels as a function of the mass flow rate, heat flux, system pressure and channel diameter and length.
  arXiv  Detail & Related papers  (2023-05-28T15:48:01Z)
• Optimal Algorithms for the Inhomogeneous Spiked Wigner Model [89.1371983413931]
  We derive an approximate message-passing algorithm (AMP) for the inhomogeneous problem. We identify in particular the existence of a statistical-to-computational gap where known algorithms require a signal-to-noise ratio bigger than the information-theoretic threshold to perform better than random.
  arXiv  Detail & Related papers  (2023-02-13T19:57:17Z)
• Rethinking Attention Mechanism in Time Series Classification [6.014777261874646]
  We promote the efficiency and performance of the attention mechanism by proposing our flexible multi-head linear attention (FMLA) We propose a simple but effective mask mechanism that helps reduce the noise influence in time series and decrease the redundancy of the proposed FMLA. We conduct extensive experiments on 85 UCR2018 datasets to compare our algorithm with 11 well-known ones and the results show that our algorithm has comparable performance in terms of top-1 accuracy.
  arXiv  Detail & Related papers  (2022-07-14T07:15:06Z)
• SreaMRAK a Streaming Multi-Resolution Adaptive Kernel Algorithm [60.61943386819384]
  Existing implementations of KRR require that all the data is stored in the main memory. We propose StreaMRAK - a streaming version of KRR. We present a showcase study on two synthetic problems and the prediction of the trajectory of a double pendulum.
  arXiv  Detail & Related papers  (2021-08-23T21:03:09Z)
• FastFlowNet: A Lightweight Network for Fast Optical Flow Estimation [81.76975488010213]
  Dense optical flow estimation plays a key role in many robotic vision tasks. Current networks often occupy large number of parameters and require heavy computation costs. Our proposed FastFlowNet works in the well-known coarse-to-fine manner with following innovations.
  arXiv  Detail & Related papers  (2021-03-08T03:09:37Z)
• Message Passing Descent for Efficient Machine Learning [4.416484585765027]
  We propose a new iterative optimization method for the bf Data-Fitting (DF) problem in Machine Learning. The approach relies on bf Graphical Model representation of the DF problem. We suggest the bf Message Passage Descent algorithm which relies on the piece-wise-polynomial representation of the model DF function.
  arXiv  Detail & Related papers  (2021-02-16T12:22:54Z)
• AQD: Towards Accurate Fully-Quantized Object Detection [94.06347866374927]
  We propose an Accurate Quantized object Detection solution, termed AQD, to get rid of floating-point computation. Our AQD achieves comparable or even better performance compared with the full-precision counterpart under extremely low-bit schemes.
  arXiv  Detail & Related papers  (2020-07-14T09:07:29Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.