🤖 machine learning

LaPrune: Controllable Differentiable Sparsity at Million Scale

本文介绍了 LaPrune，这是一种数学上精确预算的可微层，它通过使用 LapSum 障碍函数和归一化二阶矩约束，在实现硬性前 $k$ 选择的同时保持选择质量并确保梯度流，从而使百万级规模的模型实现可控的稀疏性。

原作者： Jakub Antczak, Joanna Wojciechowicz, Łukasz Struski, Jacek Tabor

发布于 2026-08-06

📖 1 分钟阅读☕ 轻松阅读

原作者： Jakub Antczak, Joanna Wojciechowicz, Łukasz Struski, Jacek Tabor

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你正在主持一场规模宏大、节奏极快的选秀节目，成千上万的选手正为了登上舞台的机会而激烈竞争。在人工智能的世界里，这些“选手”就是构成计算机大脑（神经网络）的微小部分，它们负责进行实际的思考。为了保持高效和快速，计算机不能让所有人同时发言；它需要为每项任务只挑选出最顶尖的几位专家。这被称为“稀疏计算”（sparse computation）。

难点在于如何教计算机做出这些选择。如果计算机做出了一个生硬、突然的决定（就像评委猛地敲下法槌说“你就入选了！”），数学逻辑就会崩溃，导致学习停止。但如果它做出了一个柔软、模糊的决定（就像评委说“也许是你，也许是你，也许也是你”），虽然学习效果很好，但可能会不小心让太多人登上舞台，从而浪费能量并降低速度。科学家们一直试图寻找一种兼顾两者的方案：一种既能平滑学习，又能精准挑选出恰好正确数量的获胜者，不多也不少的系统。这就是名为“LaPrune”的新论文所要解决的谜题。

问题所在：“太软”或“太硬”的困境

把神经网络想象成一个巨大的管弦乐团。当音乐开始时，指挥（AI）需要决定哪些乐器应该演奏。在一个“稀疏”乐团中，为了节省能量，任何时刻只允许少数乐器演奏。指挥使用“Top-k”规则，即“挑选出前 k 个声音最大的乐器”。

问题在于指挥如何学习这种行为。

硬方式（The Hard Way）： 如果指挥严格地指向前 k 名乐器，音乐在边界处会发生瞬间的变化。乐团无法学习如何变得更好，因为“梯度”（告诉他们如何改进的信号）被阻断了。这就像一位只接受完美答案，并拒绝给其他任何答案提供反馈的老师。
软方式（The Soft Way）： 如果指挥让每个人都带一点点声音参与，乐团的学习反馈会非常好。但现在，“预算”被打破了。原本应该只有 10 件乐器在演奏，结果可能变成了 15 件乐器在低音量播放。系统变得混乱且低效，无法满足稀疏乐团的严格规则。

以往的方法试图通过一个“温度”旋钮来修复这个问题。调高旋钮会让选择变得更软；调低旋钮则会让选择更硬。但这个旋钮非常棘手。它完全取决于乐器演奏的声音有多大。如果整个乐团的声音变大了，同样的旋钮设置会导致突然有太多人登台。这就像是试图根据室外阳光的热度而非室内实际温度来调节恒温器一样。

解决方案：LaPrune 的“归一化硬度”

于是，LaPrune（全称 Laplace Prune）登场了。作者 Jakub Antczak, Joanna Wojciechowicz, Łukasz Struski, 和 Jacek Tabor 引入了一种控制选择过程的新方法。他们没有使用一个含义随音量变化的温度旋钮，而是使用了一个归一化硬度参数（我们称之为 $\gamma$ ）。

想象 $\gamma$ 是调音台上的一个旋钮，范围从 0 到 1：

在 0 时： 旋钮设定为“等质量”（Equal Mass）。每个被选中的乐器都以完全相同的音量演奏。这是一种完美的、平滑的民主化选择。
在 1 时： 旋钮设定为“硬 Top-k”。前几名的乐器以全音量演奏，而其他人则完全静默。这是计算机在最终表演时所需的严格二进制选择。
在两者之间： 旋钮在两个极端之间创造了一条平滑的路径。

LaPrune 的魔力在于，无论你将这个旋钮设定在什么位置，**活跃乐器的总数（预算）**都保持完全一致。如果你告诉系统挑选 10 位专家，那么无论他们是都在轻声演奏，还是其中 10 位在响亮演奏而其余人保持静默，系统都会精准挑选出 10 位。

它是如何工作的：“二阶矩”的秘密

系统是如何知道该如何设置音量的？它利用了一个涉及“二阶矩”（second moment）的数学技巧。简单来说，这衡量了音量的“分布情况”。

如果所有人的音量都一样，分布就很低（低二阶矩）。
如果有些人很大声而有些人很安静，分布就会很高（高二阶矩）。

LaPrune 解决了一个复杂的数学难题，旨在找到完美的“温度”和“屏障”（一个截断点），从而同时达到你要求的精确预算和精确分布。这就像一位厨师，可以同时调整热度和食材，以确保汤的味道正好符合你要求的咸度，而不会改变锅里的总液体量。

研究发现：数据的证明

作者不仅是在猜测，他们还从多个方面证明了该方法有效：

具备百万级扩展性： 他们在一块拥有 1000 万个项目（ $n = 10^7$ ）的计算机芯片上进行了测试。LaPrune 高效地处理了这一庞大规模，耗时约 10.75 毫秒，占用 305 MB 内存。这至关重要，因为现实世界的 AI 模型规模巨大，在小规模测试组中有效的方法在面对数百万项数据时往往会崩溃。
具有尺度不变性（Scale-Invariant）： 他们测试了如果将所有输入得分放大 100 倍或缩小 10 倍会发生什么。使用旧方法（如带有固定温度的 LapSum）时，获胜者的数量会剧烈漂移。而使用 LaPrxne，如果你将硬度旋钮设为 0.9，无论输入声音多大，系统都会保持在 0.9。旋钮在任何情况下都代表相同的意义。
有助于学习： 在一项测试中，计算机必须从 200 个特征中找出 10 个隐藏的“信息特征”，LaPrune 帮助计算机成功找回正确特征的概率为 85.5%。这明显优于“软方法”（79.5%），也远好于“硬方法”（37.5%，因为梯度被阻断而无法学习任何东西）。
保持严格的预算： 他们在数学上证明了，该系统绝不会意外地让过多项目通过。即使在最坏的情况下，几乎为零（静默）的项目数量也会保证高于某个底线，从而确保系统保持稀疏。

为什么这很重要

论文指出，通过将“多少个”（预算）与“多硬”（硬度）分离，我们可以训练出既高效又智能的 AI 模型。作者展示了这种方法允许模型在训练期间（此时状态较软且灵活）进行有效学习，然后平滑过渡到现实世界使用时所需的严格、高效模式。

他们也指出，虽然这是一个强大的数学框架，但它并非解决所有问题的万灵药。该方法依赖于特定的数学原理（拉普拉斯分布），并且在系统接近完美二进制状态时，数值求解可能会变得复杂。然而，对于任何需要构建能够精准挑选赢家而不破坏规则的大规模、高效 AI 系统的开发者来说，LaPrune 提供了一个可靠且具有数学依据的新工具。它将一个混乱、依赖猜测的过程变成了一个精准、可控的旋钮。

技术摘要：LaPrune —— 百万级规模下的可控可微稀疏性

问题陈述

稀疏计算是现代神经架构（包括混合专家模型 MoE、长上下文记忆系统、检索机制和稀疏自编码器）的基础。这些系统依赖于一种基本操作：top-k 选择，即仅激活前 $k$ 个组件。

在推理过程中，这种选择是离散的。然而，在训练期间，硬 top-k 门控会在决策边界处阻断梯度，从而阻止模型学习保留哪些组件。连续松弛（软门控）虽然恢复了梯度，但往往无法解耦两个关键且独立的控制变量：

预算 ( $k$ )： 架构所需的活跃槽位（选择质量）的确切数量。
硬度 (Hardness)： 连续掩码逼近二值（0 或 1）掩码的程度。

现有的可微 top-k 松弛方法通常无法分离这些控制变量。某些方法返回不具备固定预算的软秩；其他方法使用温度或熵参数来改变稀疏性，但并未固定选择质量。因此，固定的温度可能会产生一个总质量远离预期 $k$ 值的掩码，或者硬度的变化会无意中改变预算。

方法论：LaPrune

作者引入了 LaPrune，这是一种可微稀疏选择层，它扩展了 LapSum 算子 (Struski et al., 2025)。LaPrune 在数学上强制执行精确预算，并引入了一个归一化的硬度参数。

核心机制

LaPrune 对得分向量 $r \in \mathbb{R}^n$ 和目标预算 $k$ 进行操作。它定义了一个掩码 $p$ ，其条目由屏障 $b$ 和温度 $t$ 决定：
$p_i = \sigma\left(\frac{r_i - b}{t}\right)$
其中 $\sigma$ 是拉普拉斯累积分布函数 (CDF)。屏障 $b$ 的选取需满足精确预算约束 $\sum p_i = k$ 。

为了在不改变预算的情况下控制硬度，LaPrune 使用归一化硬度参数 $\gamma \in [0, 1]$ 取代了与尺度相关的温度 $t$ 。这是通过施加一个二阶矩约束来实现的：
$\sum_{i=1}^n p_i^2 = \beta k, \quad \text{其中 } \beta = a + (1-a)\gamma, \quad a = k/n$

$\gamma = 0$ ： 对应于均匀掩码 ( $p_i = a$ )，代表弥散的分配。
$\gamma = 1$ ： 对应于硬 top-k 指示函数（二值掩码）。
$0 < \gamma < 1$ ： 该层求解唯一的有限温度 $t$ 和屏障 $b$ ，同时满足预算约束和二阶矩约束。

优化与梯度

对于 $0 < \gamma < 1$ ，LaPrune 使用牛顿法求解关于 $(b, \tau)$ 的二元非线性方程组（其中 $\tau = \log t$ ）。该过程仅需要逐元素操作以及每个迭代步求解一个 $2 \times 2$ 的线性系统。

至关重要的是，反向传播利用了隐函数定理。梯度通过直接对收敛解进行微分来计算，避免了展开牛顿迭代。这确保了反向传播具有内存效率，仅存储得分向量和收敛后的量。

主要贡献

LaPrune 层： 一种数学上精确预算的可微选择层，其归一化参数 $\gamma$ 通过二阶矩约束控制掩码硬度，且独立于得分尺度。
理论分析：
- 总体预测 (Population Prediction)： 一个预测得分分布下掩码饱和分数的模型，表明在接近二值极限时，非饱和分数随 $1-\gamma$ 线性下降。
- 最坏情况保证： 一个关于近零条目比例的无分布下界（定理 2），证明了无论输入得分分布如何，掩码几何结构都遵循严格的稀疏底线。
尺度不变性： 实证证据表明， $\gamma$ 对得分重缩放具有鲁棒性，而固定温度的方法（如 LapSum）在得分量级发生变化时会显著漂移。
可扩展性： 高效的实现方案，能够处理高达 $n = 10^7$ 的维度，同时严格执行选择预算。

实验结果

作者通过五个具体的断言验证了 LaPrune：

可扩展性： 在 NVIDIA A100 上，LaPrune 处理 $n=10^7$ 时运行时间为 10.75 ms，峰值内存为 305 MB。由于增加了第二个约束，它比 LapSum 稍慢，但与基于排序的松弛方法（如 NeuralSort, SoftSort）相比，显著降低了内存使用量，后者在较小维度下就会耗尽内存。
归一化硬度控制： 在高斯得分实验中， $\gamma$ 为任何固定预算提供了从等质量到二值端点的单调路径。与固定温度的 LapSum 不同，即使输入得分被缩放高达 100 倍，LaPrube 仍能保持一致的离散化比例（接近 0 或 1）。
特征恢复： 在一个合成特征选择任务（200 个特征，10 个信息特征）中，经过退火处理的 LaPrune 调度（ $\gamma: 0 \to 0.9$ ）实现了 0.855 的恢复 F1 值，显著优于 LapSum (0.795) 和带有直通估计器 (STE) 的硬 top-k (0.375)。
稀疏自编码器： 在针对 CIFAR-100 的 ResNet-18 激活值进行训练时，LaPrune 调度实现了优于 LapSum (FVU 7.906) 和硬 top-k (FVU 1.155) 的重建效果 (FVU 0.3462) 和探测准确率 (57.19%)。
最坏情况稀疏底线： 对 $2 \times 10^5$ 个得分进行的数值检查确认，LaPrune 掩码的近零比例始终位于定理 2 中推导出的理论无分布下界之上或与其相等。

重要性与主张

论文指出，LaPrune 建立了一个在不损害预设质量的前提下，控制可微 top-k 掩码“软度”的数学框架。其重要性在于：

解耦控制： 它成功地将预算约束与硬度约束分离，这是以往可微选择方法所缺失的能力。
理论严谨性： 它不仅提供了一种算法，还刻画了从稠密选择到离散选择的过渡，并给出了关于稀疏性的紧凑的最坏情况保证。
实用性： 它证明了在百万级维度下实现精确预算、可控的稀疏性在计算上是可行的，这使得它适用于像 MoE 和稀疏自编码器这样的大规模稀疏架构。

作者保持了谦逊的态度，指出这项工作侧重于数学构建及其数值实现，而非最大化特定下游基准测试的性能。他们承认了关于使用拉普拉斯-CDF 掩码以及在接近二值端点时面临数值调节挑战的局限性，并建议未来的工作可以扩展该框架到其他掩码族及依赖于训练的分布。