A3 : an Analytical Low-Rank Approximation Framework for Attention

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Wong, Jeffrey T. H., Zhang, Cheng, Cao, Xinye, Gimenes, Pedro, Bouganis, Christos-Savvas, Constantinides, George A., Luk, Wayne, Zhao, Yiren
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913119279775744
author Wong, Jeffrey T. H.
Zhang, Cheng
Cao, Xinye
Gimenes, Pedro
Bouganis, Christos-Savvas
Constantinides, George A.
Luk, Wayne
Zhao, Yiren
author_facet Wong, Jeffrey T. H.
Zhang, Cheng
Cao, Xinye
Gimenes, Pedro
Bouganis, Christos-Savvas
Constantinides, George A.
Luk, Wayne
Zhao, Yiren
contents 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 and code at https://github.com/DeepWok/a3.
format Preprint
id arxiv_https___arxiv_org_abs_2505_12942
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle A3 : an Analytical Low-Rank Approximation Framework for Attention
Wong, Jeffrey T. H.
Zhang, Cheng
Cao, Xinye
Gimenes, Pedro
Bouganis, Christos-Savvas
Constantinides, George A.
Luk, Wayne
Zhao, Yiren
Computation and Language
Artificial Intelligence
Machine Learning
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 and code at https://github.com/DeepWok/a3.
title A3 : an Analytical Low-Rank Approximation Framework for Attention
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2505.12942