一文读懂 NVIDIA 与 Sakana AI 联手打造的稀疏 LLM 黑科技

在 AI 圈，Scaling Law 几乎被奉为真理：模型越大，性能越强。但这个“强”是有代价的。动辄数千亿参数的模型，不仅让推理成本（VRAM、延迟）变成天文数字，更让环境足迹和能源消耗成为不可忽视的问题。

长期以来，研究者们一直盯着 Transformer 结构中那个最“臃肿”的部分——前馈网络（FFN）。在现代 LLM 中，FFN 占据了超过三分之二的参数量和 80% 以上的计算量。

有趣的是，生物大脑其实非常“省电”，只有一小部分神经元在特定时刻活跃。大模型其实也具备这种潜力：在采用 ReLU 激活函数的模型中，对于任何给定的输入，其实只有极少数神经元会被激活。

这就引出了一个困扰工业界多年的“稀疏悖论”：既然理论上大部分计算是 0，为什么在 GPU 上跑稀疏算子反而比密集算子更慢？

近日，Sakana AI 与 NVIDIA 发布的最新研究《Sparser, Faster, Lighter Transformer Language Models》正式正面硬刚这一难题。

他们不仅证明了模型可以达到 99% 以上的稀疏度且性能无损，更从底层 CUDA 内核层面设计了一套名为 TwELL 的新格式，真正把理论上的“稀疏”变成了实打实的“加速”。

1. 为什么“稀疏”在 GPU 上跑不快？

要理解这项创新的价值，得先看懂老办法错在哪。

GPU 的核心优势是极致的并行计算，它是为矩阵乘法（GEMM）这种规整、密集的活儿设计的。传统的稀疏格式（如 ELLPACK）在处理稀疏矩阵时，需要记录非零元素的索引和位置。

“由于硬件和软件栈已针对密集计算模式进行了重度优化，异构工作负载以及物化和管理稀疏索引带来的开销，一直是阻碍通用计算节省的关键挑战。”

在 gated FFN（如 Llama 采用的结构）中，稀疏模式是由 Gate 层的激活值决定的。如果想用传统稀疏算子，你得先跑一遍 Gate，数出哪些是非零的，重新排列索引，然后再跑后续的矩阵乘法。这个“重新排列”的过程（转换开销）往往比省下来的计算时间还要长。换句话说，GPU 在等待指令和搬运零散数据上浪费了太多生命。

2. TwELL 格式：为 GPU 磁贴而生的“拼图”

为了打破这个悖论，作者提出了 TwELL（Tile-wise ELLPACK）。

这是一个非常精妙的工程改进。既然 GPU 喜欢按“磁贴”（Tile）来处理任务，那我们就把稀疏化也限制在磁贴内部。TwELL 不再试图对整个矩阵进行全局的压缩和索引重排，而是在每个磁贴内部独立地进行非零元素的收集。

这种设计的核心优势在于算子融合（Operator Fusion）：

• 它允许在执行 Gate 矩阵乘法的同一个 CUDA Kernel 的末尾（Epilogue），直接物化成 TwELL 格式。

• 无需全局同步，无需中间内存读写。

• 后续的 Up 投影和 Down 投影算子可以被融合进同一个流水线，直接读取这些局部对齐的稀疏数据。

更直白地说，TwELL 就像是在生产线上直接把零件分类装好了，而不是等产品全部生产完再停掉流水线去分拣。这种“磁贴级”的局部操作，完美契合了现代 NVIDIA GPU 的硬件特性。

3. 推理与训练的双重进化

除了推理侧的 TwELL 融合算子，论文在训练侧也祭出了大招。

大模型训练时，显存（VRAM）是最大的瓶颈。中间层激活值如果不压缩，会占据巨量空间。但稀疏训练有一个巨大的坑：非均匀性。有些 Token 可能会激活 500 个神经元，有些可能只激活 5 个。如果你按 500 个的空间来预留，那就没省下内存；如果你按平均值预留，那激活多的 Token 就会产生溢出。

作者设计了一种 Hybrid（混合）格式 来解决这个问题：

• 大部分符合稀疏规律的行被存储在紧凑的稀疏矩阵中。

• 极少数激活异常高的“长尾”行则被导向一个密集的备份缓冲区。

• 这种方案既利用了 Tensor Core 的密集计算能力处理“重活”，又利用了定制的稀疏内核处理“轻活”，实现了内存与速度的双赢。

4. 实验结果：99% 稀疏度的奇迹

作者通过 正则化来诱导模型产生稀疏性。实验结果令人振奋：

从数据中我们可以看到几个核心结论：

1. 规模效应：模型越大，稀疏化的收益越明显。2B 模型的推理速度提升了 20.5%，训练速度提升了 21.9%。

2. 内存极简：1B 模型的训练峰值内存从 44.5GB 降至 33.1GB，降幅达 25.5%。

3. 性能保真：在引入轻微 正则化后，模型的平均任务准确率几乎没有下降（例如 1B 模型从 44.6% 反而微升至 44.7%）。

“我们提供了一项关于 LLM 稀疏性的定量研究，证明简单的 正则化可以诱导超过 99% 的稀疏度，且对下游性能的影响微乎其微。”

这意味着，我们不仅可以让模型跑得更快，还能在更便宜、显存更小的显卡（如 RTX 6000）上训练原本跑不动的模型。

5. 稀疏性背后的逻辑：模型也懂得“抓重点”

有趣的是，论文还揭示了 LLM 到底在什么地方变得稀疏。

作者发现，稀疏度与输入的信息熵高度相关。对于那些高度可预测的 Token（比如网址中的 doi、gov，或者常见的缩写），模型分配的激活神经元极少。而对于那些携带关键上下文信息的词（比如特定的地点 Vermont、复杂的化学名词 formaldehyde），模型的激活程度显著增加。

此外，序列位置也有影响。序列开头的前几个 Token 往往需要最多的神经元来建立上下文，随着序列增长，稀疏度会呈指数级上升。这证明了稀疏 LLM 实际上学会了动态分配计算资源，把算力花在刀刃上。

6. 局限性与未来

当然，这项技术并非没有代价。目前这套内核高度针对 NVIDIA 的 Hopper 架构（如 H100） 进行了优化，特别是利用了 Tensor Memory Accelerator (TMA) 等新特性。这意味着在旧硬件或非 NVIDIA 硬件上，收益可能打折扣。此外， 正则化系数的选择需要精细调优，过高的正则化会导致“死神经元”问题。

但瑕不掩瑜，Sakana AI 和 NVIDIA 的这项工作为我们指明了一个方向：大模型的未来不一定是无止境的堆砌算力，而是走向更精细化的计算分配。

通过开源这些代码和内核，作者希望稀疏性能够成为现代基础模型设计的一个新维度。当我们可以用更少的能量、更小的内存、更快的速度跑出同样的效果时，Scaling Law 才算真正进化到了下一个阶段。