零权重是 SSR 利用的资源
三值权重网络把每个权重限制为 −1、0 或 +1。论文指出,与 FP8 量化相比,这带来 4 倍压缩,并把 GEMM 中的乘法替换为加法和减法。零值不只是存储上的便利:零权重对矩阵乘积没有任何贡献,因此挂在它上面的每一次乘加都是白费的计算。论文报告,三值大语言模型和一般三值权重网络的稀疏度经常超过 50%,而在实践中可以达到 90–95% 稀疏度且没有明显精度损失。
论文对先前工作的批评正源于这一点。BitNet、冗余段归约(RSR)及其改进版 RSR++ 是为三值权重设计的,但按作者的说法,它们没有利用稀疏结构。CSR、JDS、COO 等传统稀疏格式则相反:它们处理稀疏性,却忽略三值特性。SSR 被定位为同时利用这两种性质的方法。
SSR 先按相同位模式分组再求和
SSR 先把三值权重矩阵 W 拆成两个二值矩阵 W1 和 W2,满足 W = W1 − W2:W1 标记 +1 元素,W2 标记 −1 元素。随后把每个二值矩阵切成每 L 列的连续列块。在一个块内,行被重新排序,使其 L 位模式按整数升序排列。模式全为零的行会从这个排序中删除,得到缩减置换序列,而分段列表记录每个剩余模式的累计计数。
推理时,输入行向量按这些分段求和,于是同一块内共享相同位模式的输入元素被一次性聚合。接着,递归树计算按位模式把这些部分和组合起来。由于排除了全零模式,分段求和数组少一个条目,论文称这在每一树层至少节省 M·N/L 次加法,并缩短了置换数组和分段列表。预处理对固定权重只运行一次,推理反复复用生成的结构。
渐近优势从 50% 稀疏度以上开始
论文推导了三种方法的操作数。BitNet.cpp 列为 2MNK/3 次操作,RSR++ 在给出的下界形式中为 2MNK/log2(K),SSR 的界则包含与稀疏度 p 相关的因子 log2(1/(1−p))。关键一步是:在独立性假设下,排除全零模式使每行每块的加法次数从 K 降到 K − p^L·K。取 L = ⌊log2 K / log2(1/(1−p))⌋ 即得到所述上界。
当 p > 0.5 时,K 的指数严格小于 RSR++ 的指数,两者在 p = 0.5 时重合。SSR 的常数因子更大,因此论文给出明确的交叉点:对 (0.5, 1] 内的稀疏度,当 K ≥ (2c)^(c/(c−1)) 且 c = log2(1/(1−p)) 时,SSR 更高效。这是关于渐近行为和矩阵维度的陈述,并不保证每个具体形状都成立。
各测试稀疏度下都出现实测加速
摘要报告,在 45–95% 稀疏度范围内,SSR 在三值 GEMM 上比 RSR++ 快 2.1–11.3 倍。这一范围比理论阈值更宽:论文指出实际评估在所有稀疏度水平上都显示改进,而渐近优势只在 50% 以上才被主张。这两种说法不应合并——一个是作者实现上的实测范围,另一个是带交叉条件的复杂度界。
讨论部分还加入与通用稀疏库的比较:SSR 据报告比 PyTorch Sparse CSC 快最多 75.8%。论文还用自身实验中的一个观察作为动机:在 CPU 上,基于 CSC 的 PyTorch Sparse GEMM 在 90% 或更低稀疏度时比 PyTorch Dense GEMM 更慢,远未达到稀疏 GEMM 约 10 倍的理论加速。这些都是作者设置下的 CPU 测量。
端到端增益小于内核增益
在 Llama-3 1B 模型上,论文报告比 RSR++ 端到端快 3.5–6.3 倍,同时总内存减少 4.9%。端到端范围比 2.1–11.3 倍的内核级范围更窄,当推理时间只有一部分花在三值 GEMM 内核上时,这是预期之中的;论文没有拆解这一差距,因此应把差异理解为两种测量的范围不同,而不是矛盾。
内存也在两个层面报告。讨论部分指出,SSR 的紧凑数据格式相比 RSR++ 最多减少 15% 的内存大小,而 Llama-3 1B 的端到端结果是总内存节省 4.9%。15% 描述的是格式,4.9% 描述的是整模型推理。它们回答不同问题,不能互换。
实现边界是仅限 CPU
论文指出,SSR 算法与 BitNet 和 RSR++ 一样不依赖底层硬件架构,可以运行在 CPU 和 GPU 上。然而,基准测试只在 CPU 上提供,因为作者提到 GPU 实现需要很高的工程投入。因此,报告的加速描述的是用循环展开、AVX2 SIMD 和多线程优化的 C++ 实现,而不是 GPU 内核。
声明的限制是缺少 GPU 实现,也缺少在更多大语言模型架构上的演示。鉴于 BitNet 和 RSR++ 在 GPU 上的强劲表现,作者预期 GPU 专用优化会带来进一步提升,但这一预期并未在本文中测量。块大小 L 也按机器经验选择,报告的最优值较小、严格低于 log2 N,并随稀疏度升高而减小、随矩阵规模增大而增大。
稀疏度水平和矩阵规模决定收益
论文自身的决策结构很清晰。在 50% 稀疏度及以下,SSR 的渐近指数在 p = 0.5 时与 RSR++ 相同,更大的常数因子对它不利;超过 50% 后,一旦 K 大到满足所述交叉条件,更小的指数就占主导。实测的 45–95% 范围显示,在作者的 CPU 实现上即使低端也有增益,因此实际阈值并不简单等于理论阈值。
对判断这是否重要的读者来说,论文最有用的区分是:相对 RSR++ 的内核级加速与 Llama-3 1B 上的端到端加速,以及格式级内存减少与整模型内存节省。证据支持的最强主张是:在作者的 CPU 设置下,利用三值 GEMM 中的零模式在测试的稀疏度范围内带来收益;它能否迁移到 GPU 或其他大语言模型架构,则被明确留作开放问题。
本文引用的文献
SSR: Sparse Segment Reduction for Ternary GEMM Acceleration
SSR 通过紧凑数据格式和随稀疏度缩放的计算树利用三值零权重,报告在 45–95% 稀疏度的三值 GEMM 上比 RSR++ 快 2.1–11.3 倍,并在 Llama-3 1B 推理上实现 3.5–6.3 倍端到端加速和 4.9% 内存节省,测量在 CPU 上完成。
