提出快速计算稀疏核矩阵乘法的新算法，适用于大模型注意力加速。
Improved Algorithms for Kernel Matrix-Vector Multiplication Under Sparsity Assumptions
- 基于核矩阵总和线性增长的假设，设计亚二次时间算法
- 在保持误差小于ε‖x‖₂的前提下，实现n的亚二次复杂度
- 对大模型中的注意力计算有实际应用价值
针对大模型中注意力矩阵的快速处理问题，研究了不对称高斯核矩阵 $K\in \mathbb{R}^{n\times n}$ 的快速矩阵-向量乘法。$K$ 的列由 $n$ 个键 $k_1,\ldots,k_n\in \mathbb{R}^d$ 索引，行由 $n$ 个查询 $q_1,\ldots,q_n\in \mathbb{R}^d$ 索引，其第 $i,j$ 项为 $K_{ij} = e^{-\|q_i-k_j\|_2^2/2σ^2}$，其中 $σ>0$ 为带宽参数。给定向量 $x\in \mathbb{R}^n$ 及误差参数 $ε>0$，目标是输出 $y\in \mathbb{R}^n$ 使得 $\|Kx-y\|_2\leq ε\|x\|_2$，且时间复杂度低于 $O(n^2)$，同时线性依赖于 $d$。算法依赖于矩阵 $K$ 的一个建模假设：其元素总和随 $n$ 线性增长，而非最坏情况下的二次增长。实验验证了该假设在多种场景（如大模型中的快速注意力）下成立。本文首次在该假设下实现了对任意向量的亚二次时间算法。
原文摘要 · Abstract (English)
Motivated by the problem of fast processing of attention matrices, we study fast algorithms for computing matrix-vector products for asymmetric Gaussian Kernel matrices $K\in \mathbb{R}^{n\times n}$. $K$'s columns are indexed by a set of $n$ keys $k_1,k_2\ldots, k_n\in \mathbb{R}^d$, rows by a set of $n$ queries $q_1,q_2,\ldots,q_n\in \mathbb{R}^d $, and its $i,j$ entry is $K_{ij} = e^{-\|q_i-k_j\|_2^2/2σ^2}$ for some bandwidth parameter $σ>0$. Given a vector $x\in \mathbb{R}^n$ and error parameter $ε>0$, our task is to output a $y\in \mathbb{R}^n$ such that $\|Kx-y\|_2\leq ε\|x\|_2$ in time subquadratic in $n$ and linear in $d$. Our algorithms rely on the following modelling assumption about the matrices $K$: the sum of the entries of $K$ scales linearly in $n$, as opposed to worst case quadratic growth. We validate this assumption experimentally, for Gaussian kernel matrices encountered in various settings such as fast attention computation in LLMs. We obtain the first subquadratic-time algorithm that works under this assumption, for unrestricted vectors.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。