基于黎曼流形的精确预算约束模型压缩
机器学习
2026-05-08 v2
摘要
在有效AI中，为每个N个组分配一个K选项，总成本受预算约束，这是一种常见问题，包括混合精度量化、非均匀剪枝和专家选择等。 目标函数（通常为模型损失）在所有分配之间存在联合依赖，无法跨组分解，从而阻止组合求解器直接优化真实目标，迫使依赖代理公式。 诸如进化搜索等方法评估实际损失但缺乏梯度信息，而惩罚方法则对预算仅实现近似约束，往往需要 extensive hyperparameter调优。 我们提出一种新方法，表明在softmax松弛下，预算约束在logit空间中定义一个光滑的黎曼流形，其几何结构异常简单。 正切向量 admits闭式表达式，沿成本向量移动logit会单调改变期望成本，向量运输简化为单个内积。 基于这些性质，我们提出黎曼约束优化（RCO），其增强标准Adam步骤，包括切向投影、二分查找撤退和动量运输。 结合Gumbel直通估计和用于离散可行性的预算受限动态规划，RCO能够在无需引入约束特定超参数的情况下，对实际损失进行一阶优化，严格满足预算约束。 在合成基准和真实LLM压缩场景下，RCO在匹配或超越当前最先进方法的同时，往往需要显著更少的计算时间。 源代码已公开于https://github.com/IST-DASLab/RCO。
引用
@article{arxiv.2605.00649,
title = {Model Compression with Exact Budget Constraints via Riemannian Manifolds},
author = {Michael Helcig and Dan Alistarh},
journal= {arXiv preprint arXiv:2605.00649},
year = {2026}
}