A3: an Analytical Low-Rank Approximation Framework for Attention

Large language models have demonstrated remarkable performance; however, their massive parameter counts make deployment highly expensive. Low-rank approximation offers a promising compression solution, yet existing approaches have two main limitations: (1) They focus on minimizing the output error of individual linear layers, without considering the architectural characteristics of Transformers, and (2) they decompose a large weight matrix into two small low-rank matrices. Consequently, these methods often fall short compared to other compression techniques like pruning and quantization, and introduce runtime overhead such as the extra GEMM kernel launches and memory operations for decomposed small matrices. To address these limitations, we propose $A^3$, a post-training low-rank approximation framework. $A^3$ splits a Transformer layer into three functional components, namely $\texttt{QK}$, $\texttt{OV}$, and $\texttt{MLP}$ and provides analytical solutions that reduces the hidden dimension size inside each component while minimizing the component's functional loss. This approach directly reduces model sizes, KV cache sizes, and FLOPs without introducing any runtime overheads. Through extensive experiments, we show that $A^3$ maintains superior performance compared to SoTAs. For example, under the same reduction budget in computation and memory, our low-rank approximated LLaMA 3.1-70B achieves a perplexity of 4.69 on WikiText-2, outperforming the previous SoTA's 7.87 by 3.18. We also show versatile applications of $A^3$ in KV cache compression, integration with quantization, fine-tuning and mixed-rank assignments. We open-sourced our framework at https://github.com/DeepWok/a3.

Large language models, the AI systems behind tools like chatbots and code assistants, are powerful but extremely expensive to run, often requiring billions of parameters and significant computing resources. One way to make them smaller and faster is to find simplified approximations of their internal calculations, a technique known as low-rank approximation. However, existing methods treat each calculation in the model independently, ignoring how different parts of the architecture work together, and they often introduce extra computational overhead that limits real-world speedups. We propose $A^3$, a framework that instead looks at how the model's components actually function -- how it computes attention between words, combines information, and processes through its feed-forward layers -- and finds mathematically optimal ways to shrink each component. Because $A^3$ reduces the internal dimensions of these components rather than splitting matrices into pairs, it avoids extra computational steps at runtime, directly saving memory, computation, and storage including the cache used during text generation. In experiments across multiple models, $A^3$ significantly outperforms previous methods. For example, on a 70-billion-parameter model, $A^3$ nearly halves the performance loss compared to the best existing approach at the same compression level, while also delivering genuine inference speedups. $A^3$ can be combined with other compression techniques like quantization for even greater efficiency, making large AI models more practical and affordable to deploy.