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

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.

💡 Research Summary

The paper introduces COALA, a novel framework for context‑aware low‑rank approximation (LRA) that addresses three major practical challenges in compressing and fine‑tuning large neural networks: numerical instability, memory overload, and data scarcity. Traditional context‑aware LRA methods minimize the weighted Frobenius loss ‖W X − W′ X‖_F by forming the Gram matrix G = X Xᵀ, taking its square root, and then inverting it to compute a projection. This approach fails when G is ill‑conditioned or when the calibration data X is too large to fit in GPU memory. Moreover, with few calibration samples the problem becomes ill‑posed, leading to over‑fitting.

COALA eliminates the need for any Gram‑matrix computation or matrix inversion. The key insight (Proposition 1) is that the optimal rank‑r approximation of the weighted problem can be written as W′ = U_r U_rᵀ W, where U_r contains the leading r left singular vectors of the product A = W X. To compute U_r efficiently, the authors first compute a QR decomposition of Xᵀ using a Tall‑Skinny QR (TSQR) algorithm, which processes X in small chunks and only retains the upper‑triangular factor R. Because Rᵀ R = X Xᵀ, the singular vectors of W Rᵀ are identical to those of W X, allowing the problem to be solved by a single SVD on the much smaller matrix W Rᵀ (Proposition 2). This “inversion‑free” pipeline avoids the catastrophic loss of precision that occurs when tiny singular values of X cause the Gram matrix to become nearly singular.

The authors also extend the method to a regularized setting, adding a term μ‖W − W′‖_F² to the loss to prevent over‑fitting when calibration data are scarce. Proposition 3 shows that the regularized problem is equivalent to the unregularized one applied to an augmented data matrix X′ =