通过黎曼流形进行具有精确预算约束的模型压缩

在总成本预算下将 K 个选项中的一个分配给 N 个组中的每一个是高效人工智能中反复出现的问题，出现在混合精度量化、非均匀剪枝和专家选择中。目标（模型损失）共同取决于所有分配，并且不会跨组分解，因此组合求解器只能优化代理目标。进化搜索评估实际损失，但缺乏梯度，而基于惩罚的方法仅近似地强制执行预算，并且通常需要大量的超参数调整。我们证明，在 softmax 松弛下，预算约束在 logits 空间中定义了一个平滑的黎曼流形，具有异常干净的几何形状：法线向量以封闭形式提供，沿着成本向量移动 logits 单调地改变预期成本，并且向量传输减少为单个内积。在此基础上，我们提出了黎曼约束优化（RCO），它围绕标准 Adam 步骤封装了切线投影、二分搜索回撤和动量传输。结合 Gumbel 直通估计和离散可行性的预算约束动态规划，RCO 使用一阶方法优化实际损失，在每次迭代时精确执行预期预算，并且不引入与约束相关的超参数。相同的结构可以处理多个同时预算，无需额外系数。 RCO 在综合问题和实际的 LLM 压缩设置上匹配或超过了最先进的方法，而且通常成本要低得多。源代码可在 https://github.com/IST-DASLab/RCO 获取。