🤖 machine learning

Budget Constraints as Riemannian Manifolds

本文提出了黎曼约束优化（RCO），这是一种新颖的框架，它将预算约束建模为光滑黎曼流形，从而在严格预算执行的前提下实现对不可分解目标的高效基于梯度的优化，在混合精度量化和专家剪枝等任务中，其解的质量和计算效率均优于现有的惩罚法和进化方法。

原作者： Michael Helcig, Dan Alistarh

发布于 2026-05-04

📖 1 分钟阅读☕ 轻松阅读

原作者： Michael Helcig, Dan Alistarh

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

想象一下，你是一家超大型高端餐厅的主厨。你当晚的预算严格受限，但菜单上却有数百道菜肴，且每道菜都有多种制作方式（例如，使用高级食材、标准食材或经济型替代品）。

你的目标是：为每道菜恰好选定一种制作方式，使得总成本严格控制在预算之内，同时让整餐的整体品质尽可能美味。

问题在于：餐食的整体品质并非仅仅是各道菜品质的简单相加。如果你选择了一块高级牛排，它可能需要搭配特定的葡萄酒，从而改变整桌的“风味轮廓”。这使得数学计算变得极其困难：你不能孤立地看待每道菜；你必须解决一个巨大而错综复杂的谜题，其中每一个选择都会影响其他所有选择。

这正是机器学习工程师在尝试压缩大型人工智能模型（例如驱动聊天机器人的那些模型）时所面临的难题。他们需要决定如何“缩小”或“剪枝”模型的不同部分，以符合尺寸限制（即预算），同时又不损害模型的智能（即品质）。

以下是该论文如何通过几个富有创意的类比来解决这一问题的：

1. 旧方法：猜测与惩罚

此前，工程师们主要尝试过两种方法，但都显得笨拙：

“惩罚”法：他们告诉计算机：“尽量保持在预算内，但如果超出预算，我就给你一笔巨额‘罚款’（惩罚分数）。”问题在于，计算机很难猜出正确的罚款额度。如果罚款太小，它就会无视预算；如果罚款太大，计算机就会畏缩不前，停止学习。这就像试图通过随机音量大喊“不！”来教狗坐下；狗永远学不会确切的规则。
“进化”法：他们让计算机尝试成千上万种随机组合，保留最好的那些，然后重复此过程。这种方法效果不错，但速度极慢。这就像试图通过烹饪世界上每一种可能的菜肴并逐一品尝，来寻找最佳食谱。这需要耗费永恒的时间。

2. 新构想：“预算流形”

作者们意识到，如果通过特定的数学透镜（使用所谓的"softmax"）来审视这个问题，预算约束就不再是一堵你必须撞上去的杂乱墙壁。相反，它是一条你可以行走的平滑曲面（流形）。

将预算想象成不是一道坚硬的栅栏，而是一条走钢丝。

曲面：想象一个巨大的、无形的、弯曲的蹦床，它只存在于你的总成本恰好等于预算的位置。
行走：计算机不需要跳出蹦床并指望它能落回上面。相反，它沿着曲面行走。

3. 新方法（RCO）如何运作

该论文提出了一种名为**黎曼约束优化（Riemannian Constrained Optimization, RCO）**的新算法。以下是它如何沿着那条钢丝行走的：

步骤 1：切向步（向前行走）：计算机计算出能让餐食更美味的方向（梯度）。但它不会直接朝那个方向走，而是将该方向投影到钢丝的曲面上。这确保它永远不会意外偏离预算线。
步骤 2：二分搜索（魔法滑梯）：有时，即使小心翼翼地行走，你也可能会轻微偏离线条。在其他方法中，你需要进行复杂的计算才能回到线上。而在这里，作者们发现了一个“魔法滑梯”。由于他们使用了特定的数学方法，他们只需通过一个旋钮（二分搜索）将整个餐食计划向上或向下滑动，就能完美地落回预算线上。这就像拥有一个能瞬间修复你平衡的遥控器。
步骤 3：动量（保持节奏）：当你在弯曲的表面上行走时，你的方向会发生变化。该算法有一个特殊的技巧来“传输”其动量（即它对自己行进方向的记忆），这样当它沿着曲线移动时，就不会感到眩晕或失去节奏。

4. 为何这意义重大

该论文声称，这种方法之所以是游戏规则的改变者，原因有二：

它是精确的：与旧有的“惩罚”方法（往往最终会略微超出或低于预算）不同，该方法在每一步都严格保持在预算线上。这就像一位从不摇晃的走钢丝者。
它很快：因为它使用梯度（数学方向）而非随机猜测，所以它能更快地找到最佳解决方案。
- 结果：在合成谜题的测试中，旧方法停留在最佳可能分数的 83%，而新方法找到了完美解。
- 现实世界：当他们在压缩大型人工智能模型（例如缩减“大语言模型”的规模）时进行测试，其结果与缓慢的“进化”方法相当或更优，但速度快了3 到 16 倍。

总结

该论文引入了一种解决人工智能中“预算”问题的新方法。他们不再将预算视为会破坏计算的硬性限制，而是将其转化为一条平滑、可行走的曲面。通过沿着这条曲面行走，计算机能够比以前更快、更准确地找到成本与品质之间的完美平衡，而无需猜测或调整棘手的参数。这之间的区别，就像是在黑暗的房间里跌跌撞撞地试图避开家具，与自信地走在一条灯火通明、铺设完美的道路上。

以下是 Michael Helcig 和 Dan Alistarh 的论文《预算约束作为黎曼流形》的详细技术总结。

1. 问题陈述

本文解决了预算约束离散分配问题，这是高效机器学习中的一个 recurring 挑战。目标是为 $N$ 个组（例如神经网络层）中的每一个分配 $K$ 个离散选项之一（例如量化位宽、剪枝决策或专家选择），以在总成本预算（例如模型大小或延迟）的约束下最小化全局目标（例如模型损失）。

关键挑战：

不可分解的目标函数： 模型损失依赖于所有分配的联合（例如层与层之间的交互），这意味着它不能写成独立各组项的总和。这导致标准的动态规划（DP）无法直接优化实际损失；DP 只能优化“代理”目标（例如逐层敏感性）。
现有方法的局限性：
- 组合求解器（DP）： 能精确处理预算，但需要可分解的目标函数，迫使使用不准确的代理分数。
- 进化搜索： 评估实际损失，但缺乏梯度，需要数百次前向传播，且计算成本高昂。
- 惩罚/拉格朗日方法： 在损失中添加惩罚项（例如 $\lambda(C-B)^2$ ）。这些方法仅近似地强制执行预算，经常发生振荡，且需要敏感的超参数调整。

2. 方法论：黎曼约束优化（RCO）

作者提出了一种几何方法，将预算约束视为对数空间（logits space）中的平滑黎曼流形，而非惩罚项。

A. 预算流形

作者使用对数 $\alpha_i \in \mathbb{R}^K$ 参数化每个组 $i$ 的分配分布，概率为 $p_i = \text{softmax}(\alpha_i)$ 。总期望成本定义为：
$C(\alpha) = \sum_{i=1}^N w_i \langle p_i, c \rangle$
其中 $w_i$ 是组权重， $c$ 是成本向量。

流形结构： 水平集 $\mathcal{M} = \{ \alpha \mid C(\alpha) = B \}$ 被证明是一个平滑的 $(NK-1)$ 维黎曼子流形。
几何性质：
1. 闭式法向量： 梯度 $\nabla C(\alpha)$ 具有闭式解且非零（假设成本不同），定义了流形的法向量。
2. 单调重入（Retraction）： 沿成本向量移动对数会单调地改变期望成本。这使得算法可以使用简单的二分搜索（或牛顿法）将任意点投影回流形 $\mathcal{M}$ ，而无需复杂的迭代投影。
3. 廉价的向量传输： 将动量从一个切平面移动到另一个切平面简化为单个内积运算。

B. RCO 算法

提出的算法将标准的一阶优化（Adam）包裹在三个几何操作中，以在每一步精确强制执行预算：

切向投影： 损失梯度通过移除沿法向量 $\nabla C(\alpha)$ 的分量，投影到流形的切空间上。这确保优化器只看到保持预算的方向。
重入（Retraction）： 在 Adam 步之后，迭代点可能会偏离流形。算法沿成本向量移动所有对数（ $\alpha' = \alpha + t \cdot c$ ），并使用二分搜索找到 $t$ ，使得 $C(\alpha') = B$ 。
动量传输： 优化器的动量（一阶矩）被投影到新的切空间以保持一致性。

C. 处理离散可行性（前向传播）

由于流形在连续对数空间中操作，本文将 RCO 与以下方法结合：

Gumbel-STE（直通估计器）： 用于生成前向传播的离散分配。
预算约束动态规划： 在前向传播期间精确求解离散分配问题（多选择背包问题），确保评估的配置满足预算。
机制： 前向传播使用 DP 根据受扰动的对数找到最佳离散配置 $z^*$ 。反向传播使用 STE 将梯度流通过受扰动对数的 softmax 概率，而流形投影确保这些梯度不违反预算约束。

3. 主要贡献

预算流形的识别： 作者指出，softmax 期望成本的水平集形成了一个具有独特且计算廉价几何性质（闭式法向量、单调重入）的黎曼流形。
RCO 算法： 一种新颖的优化框架，以零约束超参数（无需调整 $\lambda$ ）强制执行精确的预算约束。它结合了切向投影、二分搜索重入和围绕 Adam 的动量传输。
处理不等式约束： 通过松弛变量 $s$ ，将流形扩展到不等式约束（ $C(\alpha) \le B$ ），允许优化器自然地利用未使用的预算，而不会将其浪费在次优选项上。
离散 - 连续混合： Gumbel-STE 和预算约束 DP 的集成允许优化实际的不可分解模型损失，同时保持精确的离散可行性。

4. 实验结果

本文在合成和现实世界的大语言模型（LLM）压缩任务上验证了 RCO。

合成背包问题（MCKP）：
- 在困难实例上，RCO 恢复了动态规划（DP）的最优解。
- 拉格朗日方法在最优值的 83% 处达到平台期，并表现出显著的约束违反（ $10^{-1}$ 到 $10^{-2}$ ），而 RCO 将违反保持在机器精度附近（ $10^{-9}$ ）。
- RCO 通过利用松弛变量完美处理“廉价最优”情况（即最佳解的成本低于预算），而等式约束方法则浪费预算。
MoE 专家剪枝：
- 在 OLMoE-1B-7B 和 Qwen3-30B-A3B 上进行评估。
- 在多个基准测试（ARC、MMLU 等）中，RCO 的准确率匹配或超过了 EvoESAP（一种进化搜索基线）。
- 效率： RCO 以 3–16 倍 更低的挂钟时间实现这些结果（例如，EvoESAP 需 5.8 小时，RCO 仅需约 10 分钟）。
混合精度量化：
- 在 Qwen3-8B 上评估（为 252 层分配 2–8 位宽）。
- 在高压缩率（2.25 位）下，与 HIGGS（一种线性代理方法）相比，RCO 将困惑度降低了 36%，与 EvoPress（进化搜索）相比降低了 4%。
- RCO 在 2.5 位时收敛到与 EvoPress 相同的最优解，但时间缩短了 约 10 倍（62 分钟对比 11–14 小时）。

5. 意义与影响

精确的约束执行： 与惩罚方法不同，RCO 保证在每次迭代中预算都被精确满足，消除了对惩罚权重超参数调整的需求。
不可分解损失的优化： 通过将离散可行性（由前向传播中的 DP 处理）与连续优化（由反向传播中的流形处理）分离，RCO 允许在严格的预算约束下直接优化复杂的、不可分解的模型损失（如端到端准确率）。
计算效率： 几何操作（投影、重入、传输）计算廉价（本质上是内积和二分搜索），使该方法可扩展到进化搜索 prohibitively 昂贵的大型模型。
通用性： 该框架适用于任何涉及线性成本的 softmax 参数化问题，其应用范围从模型压缩扩展到资源受限的神经架构搜索和其他离散分配问题。

总之，本文将预算约束离散优化重新表述为在特定黎曼流形上导航的问题，提供了一种理论扎实、无超参数且高效的替代方案，优于基于惩罚和进化的方法。