Taming the Exponential: A Fast Softmax Surrogate for Integer-Native Edge Inference

• URL: http://arxiv.org/abs/2604.02292v1
• Date: Thu, 02 Apr 2026 17:32:29 GMT
• Title: Taming the Exponential: A Fast Softmax Surrogate for Integer-Native Edge Inference
• Abstract summary: We propose a bounded, monotone surrogate to the exponential softmax function, which uses a clipped linear mapping of the max centered attention logits.<n>This approximation produces a stable probability distribution, maintains the ordering of the original logits and has non-negative values.<n>We describe a hardware-motivated implementation of HCCS for high- throughput scenarios targeting the AMD Versal AI Engines.
• Score: 0.8488076117647583
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Softmax can become a computational bottleneck in the Transformer model's Multi-Head Attention (MHA) block, particularly in small models under low-precision inference, where exponentiation and normalization incur significant overhead. As such, we suggest using Head-Calibrated Clipped-Linear Softmax (HCCS), a bounded, monotone surrogate to the exponential softmax function, which uses a clipped linear mapping of the max centered attention logits. This approximation produces a stable probability distribution, maintains the ordering of the original logits and has non-negative values. HCCS differs from previous softmax surrogates as it includes a set of lightweight calibration parameters that are optimized offline based on a representative dataset and calibrated for each individual attention head to preserve the statistical properties of the individual heads. We describe a hardware-motivated implementation of HCCS for high-throughput scenarios targeting the AMD Versal AI Engines. The current reference implementations from AMD for this platform rely upon either bfloat16 arithmetic or LUTs to perform the exponential operation, which might limit the throughput of the platform and fail to utilize the high-throughput integer vector processing units of the AI Engine. In contrast, HCCS provides a natural mapping to the AI Engines' int8 multiply accumulate (MAC) units. To the best of our knowledge, this is the first int8 optimized softmax surrogate for AMD AI engines that significantly exceeds the speed performance of other reference implementations while maintaining competitive task accuracy on small or heavily quantized MHA workloads after quantization-aware retraining.

Related papers

• Low-Cost Multi-Precision Systolic Arrays for Accelerating FHE NTTs on AI ASICs [0.3186130813218338]
  Homomorphic Encryption (FHE) ensures robust data privacy but suffers from prohibitive computational overhead.<n>Accelerating FHE on AI hardware like Processing Units (TPUs) is promising, yet fundamentally limited by a precision mismatch.<n>We propose a minimally modified multi-precision systolic array that performs full-precision output reconstruction in sync with low-precision matrix multiplication.
  arXiv  Detail & Related papers  (2026-06-18T07:26:13Z)
• POET-X: Memory-efficient LLM Training by Scaling Orthogonal Transformation [57.57816409869894]
  We introduce POET-X, a scalable and memory-efficient variant for training large language models.<n>PoET-X maintains the generalization and stability benefits of POET while achieving substantial improvements in throughput and memory efficiency.
  arXiv  Detail & Related papers  (2026-03-05T18:59:23Z)
• BAPS: A Fine-Grained Low-Precision Scheme for Softmax in Attention via Block-Aware Precision reScaling [12.43240392025487]
  We introduce a novel low-precision workflow that employs a specific 8-bit floating-point format (HiF8) and block-aware precision rescaling for softmax.<n>Our algorithmic innovations make low-precision softmax feasible without the significant model accuracy loss.<n>Our work paves the way for doubling end-to-end inference throughput without increasing chip area.
  arXiv  Detail & Related papers  (2026-02-02T13:12:18Z)
• BLASST: Dynamic BLocked Attention Sparsity via Softmax Thresholding [28.11907989174509]
  We introduce BLASST, a drop-in sparse attention method that dynamically prunes the attention matrix without any pre-computation or proxy scores.<n>Our method uses a fixed threshold and existing information from online softmax to identify negligible attention scores.<n>We develop an automated calibration procedure that reveals a simple inverse relationship between optimal threshold and context length.
  arXiv  Detail & Related papers  (2025-12-12T23:30:43Z)
• SOLE: Hardware-Software Co-design of Softmax and LayerNorm for Efficient Transformer Inference [6.157559748568282]
  We present SOLE, a hardware-software co-design for Softmax and LayerNorm.<n>We achieve both low-precision calculation and low bit-width storage on Softmax and LayerNorm.
  arXiv  Detail & Related papers  (2025-10-20T06:09:09Z)
• Adaptive Sampled Softmax with Inverted Multi-Index: Methods, Theory and Applications [79.53938312089308]
  The MIDX-Sampler is a novel adaptive sampling strategy based on an inverted multi-index approach.<n>Our method is backed by rigorous theoretical analysis, addressing key concerns such as sampling bias, gradient bias, convergence rates, and generalization error bounds.
  arXiv  Detail & Related papers  (2025-01-15T04:09:21Z)
• HyperINF: Unleashing the HyperPower of the Schulz's Method for Data Influence Estimation [37.62285675595782]
  We propose HyperINF, an efficient and accurate influence function approximation method.<n>We use the generalized fisher information (GFIM) as a low-rank approximation of the Hessian matrix.<n>On LoRA-tuned models, HyperINF achieves superior downstream performance with minimal memory and computational overhead.
  arXiv  Detail & Related papers  (2024-10-07T14:42:45Z)
• DeepGEMM: Accelerated Ultra Low-Precision Inference on CPU Architectures using Lookup Tables [49.965024476651706]
  DeepGEMM is a lookup table based approach for the execution of ultra low-precision convolutional neural networks on SIMD hardware. Our implementation outperforms corresponding 8-bit integer kernels by up to 1.74x on x86 platforms.
  arXiv  Detail & Related papers  (2023-04-18T15:13:10Z)
• Improving Dual-Encoder Training through Dynamic Indexes for Negative Mining [61.09807522366773]
  We introduce an algorithm that approximates the softmax with provable bounds and that dynamically maintains the tree. In our study on datasets with over twenty million targets, our approach cuts error by half in relation to oracle brute-force negative mining.
  arXiv  Detail & Related papers  (2023-03-27T15:18:32Z)
• Softmax-free Linear Transformers [90.83157268265654]
  Vision transformers (ViTs) have pushed the state-of-the-art for visual perception tasks. Existing methods are either theoretically flawed or empirically ineffective for visual recognition. We propose a family of Softmax-Free Transformers (SOFT)
  arXiv  Detail & Related papers  (2022-07-05T03:08:27Z)
• Efficient Softmax Approximation for Deep Neural Networks with Attention Mechanism [0.0]
  We propose two methods to approximate softmax computation, which are based on the usage of LookUp Tables (LUTs) We show that 8-bit approximation allows to obtain acceptable accuracy loss below $1.0%$.
  arXiv  Detail & Related papers  (2021-11-21T08:56:29Z)
• Taming GANs with Lookahead-Minmax [63.90038365274479]
  Experimental results on MNIST, SVHN, CIFAR-10, and ImageNet demonstrate a clear advantage of combining Lookahead-minmax with Adam or extragradient. Using 30-fold fewer parameters and 16-fold smaller minibatches we outperform the reported performance of the class-dependent BigGAN on CIFAR-10 by obtaining FID of 12.19 without using the class labels.
  arXiv  Detail & Related papers  (2020-06-25T17:13:23Z)

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.