LAMP: Look-Ahead Mixed-Precision Inference of Large Language Models

TL;DR

The paper tackles efficiency and numerical reliability in transformer inference by treating the computation as a deep composition $f(g(\bm{x}))$ subject to floating-point rounding. It develops Look-Ahead Mixed-Precision (LAMP), a principled strategy that adaptively recomputes a sparse subset of inner components with higher precision guided by bounds tied to the Jacobian of the remaining nonlinearities. The authors show that nearly-sparse solutions for the key transformer nonlinearities (activations, RMS layer normalization, and softmax within attention) can be obtained via greedy algorithms with $\mathcal{O}(n \log n)$ complexity, and demonstrate substantial accuracy gains on GPT-2 XL with only a small recomputation fraction. They validate the approach using a custom low-precision simulation and multiple datasets, indicating that selective high-precision recomputation can meaningfully improve inference quality without model retraining or weight changes, highlighting a practical path toward greener, faster LLM inference.

Abstract

Mixed-precision computations are a hallmark of the current stage of AI, driving the progress in large language models towards efficient, locally deployable solutions. This article addresses the floating-point computation of compositionally-rich functions, concentrating on transformer inference. Based on the rounding error analysis of a composition $f(g(\mathrm{x}))$, we provide an adaptive strategy that selects a small subset of components of $g(\mathrm{x})$ to be computed more accurately while all other computations can be carried out with lower accuracy. We then explain how this strategy can be applied to different compositions within a transformer and illustrate its overall effect on transformer inference. We study the effectiveness of this algorithm numerically on GPT-2 models and demonstrate that already very low recomputation rates allow for improvements of up to two orders of magnitude in accuracy.

Figures (13)

• Figure 1: Performance of the proposed adaptive mixed-precision inference strategy with $\mu$ mantissa bits used for the low-precision accumulation of key-query inner products. The mixed precision strategy recomputes approximately 8.3% of the key-query inner products. The Kullback--Leibler divergence is measured against a reference model with uniform FP32 accumulation.
• Figure 2: Performance of mixed-precision GPT-2 XL inference on the OpenWebText dataset with LAMP evaluation of the key-query inner products: varying number of mantissa bits ($\mu$) and fixed threshold of LAMP ($\tau$).
• Figure 3: Performance of mixed-precision GPT-2 XL inference on the OpenWebText dataset with LAMP evaluation of the key-query inner products: fixed number of mantissa bits ($\mu$) and varying threshold of LAMP ($\tau$).
• Figure 4: Performance of mixed-precision GPT-2 XL inference on the OpenWebText (left), CodeParrot (center), and ArXiv (right) datasets with LAMP evaluation of the key-query inner products: varying number of mantissa bits ($\mu$) and fixed threshold of LAMP ($\tau$).
• Figure 5: Performance of mixed-precision GPT-2 XL inference on the OpenWebText (left), CodeParrot (center), and ArXiv (right) datasets with LAMP evaluation of the key-query inner products: fixed number of mantissa bits ($\mu$) and varying threshold of LAMP ($\tau$).

Theorems & Definitions (14)

• Example 2.1
• Example 2.2
• Lemma 3.1
• proof
• Lemma 3.2
• proof
• Proposition 3.3
• proof
• Lemma 3.4
• proof