Reproduce “Solving Attention Kernel Regression Problem via Pre-conditioner” on OpenResearch

About this paper

Zhao Song, Junze Yin, Lichen Zhang

Researchers from Adobe Research, Boston University, and MIT introduce efficient algorithms to solve novel regression problems that act as proxies for the computationally intensive attention mechanism in large language models. Their methods achieve significant runtime improvements by leveraging sketching and preconditioning for both matrix exponential powers of Gram matrices and entry-wise exponential Gram matrix regressions.

Key insights

• Complex matrix function regressions related to attention can be decomposed into a series of more fundamental linear and PSD regression problems.
• Combining subspace embedding sketches and preconditioning transforms ill-conditioned regression problems into well-conditioned ones, enabling rapid convergence of iterative solvers like gradient descent.
• Tensor sketching techniques can effectively approximate high-degree polynomial kernels, which is crucial for handling the Taylor series expansion of the entry-wise matrix exponential.

Summary by alphaXiv.