COALA: Numerically Stable and Efficient Framework for Context-Aware Low-Rank Approximation

TL;DR

This paper tackles numerical instability in context-aware low-rank approximation of large language models by introducing COALA, an inversion-free, regularized framework that relies on stable decompositions. It avoids Gram-matrix inversions even for large activation matrices by using a projection-based solution W' = U_r U_r^T W and QR/TSQR preprocessing, enabling memory-efficient computations. The authors prove convergence and provide explicit error bounds for the regularized problem, showing that W_mu approaches the unregularized solution W_0 as mu -> 0 with rates dependent on the spectral gap of WX. Empirically, COALA improves model compression and fine-tuning performance across several models and datasets, demonstrating robustness to data scarcity and large calibration matrices, and offering a practical pathway for scalable, reliable adaptation of large-scale networks.

Abstract

Recent studies suggest that context-aware low-rank approximation is a useful tool for compression and fine-tuning of modern large-scale neural networks. In this type of approximation, a norm is weighted by a matrix of input activations, significantly improving metrics over the unweighted case. Nevertheless, existing methods for neural networks suffer from numerical instabilities due to their reliance on classical formulas involving explicit Gram matrix computation and their subsequent inversion. We demonstrate that this can degrade the approximation quality or cause numerically singular matrices. To address these limitations, we propose a novel inversion-free regularized framework that is based entirely on stable decompositions and overcomes the numerical pitfalls of prior art. Our method can handle possible challenging scenarios: (1) when calibration matrices exceed GPU memory capacity, (2) when input activation matrices are nearly singular, and even (3) when insufficient data prevents unique approximation. For the latter, we prove that our solution converges to a desired approximation and derive explicit error bounds.

Figures (6)

• Figure 1: Relative approximation error versus approximation rank obtained by different methods on layer 1 $\texttt{q\_proj}$. The reference weight matrix $W^{\text{ref}}_r$ was computed using the inversion-free COALA method and in high working precision (fp64) to serve as the ground-truth solution. The LLaMA3-1B LLama3 model was used with 64 examples from the Wikitext merity2016pointer dataset.
• Figure 2: Distribution of singular values of matrix $X$, obtained from the outputs of layer 1 $\texttt{q\_proj}$ in the LLaMA3-1B LLama3 model, computed over 64 samples from the WikiText merity2016pointer dataset.
• Figure 3: Runtimes for computing $S$: $SS^\top = XX^\top$ using two approaches. Left: Matrix $X \in \mathbb{R}^{4096 \times n}$ for different $n$. Right: Matrix $X \in \mathbb{R}^{4096 \times 3 \cdot 10^5}$ split into chunks of different size. In this case, QR is computed using the TSQR method and the Gram matrix using $XX^\top = \sum_{i=1}^p X_i X_i^\top$.
• Figure 4: Comparison of the impact of parameter tuning with (Equation \ref{['eq:reglam']}) and without considering layer‑wise norms on model quality at 70% compression, evaluated on a common‑sense reasoning dataset using the Mistral-7B‑Instruct model.
• Figure 5: The dependence of average accuracy on the parameter $\lambda$ on the commonsense reasoning dataset for different models. On the left: 70% compression ratio, on the right: 80% compression ratio.

Theorems & Definitions (30)

• Proposition 1
• proof
• Proposition 2
• proof
• Proposition 3
• proof
• Theorem 1
• proof
• Proposition 4
• proof