arXiv:2605.00649cs.LG2026-05被引 3

用黎曼几何精确控制模型压缩预算,实现高效优化。

Model Compression with Exact Budget Constraints via Riemannian Manifolds

论文配图:Model Compression with Exact Budget Constraints via Riemannian Manifolds
图 1 · 摘自论文原文
  • 将预算约束转化为对数空间中的光滑黎曼流形,简化优化结构。
  • 在真实损失函数下实现精确预算约束,速度优于现有方法。
  • 适合需要精准预算控制的模型压缩场景,如大模型量化与剪枝。

在高效AI中,为N个组分配K种选项并满足总成本预算是一类常见问题,涵盖混合精度量化、非均匀剪枝和专家选择等任务。目标函数(通常为模型损失)依赖于所有分配结果,无法按组分解,导致组合优化器无法直接优化真实目标,只能依赖代理形式。现有方法如进化搜索虽可评估真实损失但缺乏梯度信息,惩罚法仅近似满足预算且需大量超参数调优。本文提出新方法:通过Softmax松弛,发现预算约束在对数空间中定义了一个具有特殊简单几何结构的光滑黎曼流形。其法向量有闭式表达,沿成本向量调整对数概率单调改变期望成本,向量传输退化为单个内积。基于此,我们提出黎曼约束优化(RCO),在标准Adam步骤中加入切向投影、二分搜索回缩和动量传输。结合Gumbel直通估计与预算约束动态规划保证离散可行性,实现无需引入特定约束超参数的第一阶优化,精确满足预算约束并最小化真实损失。在合成基准与真实大模型压缩设置中,RCO性能达或超越当前最优,且耗时显著更短。代码已开源:https://github.com/IST-DASLab/RCO。

原文摘要 · Abstract (English)

Assigning one of K options to each of N groups under a total cost budget is a recurring problem in efficient AI, including mixed-precision quantization, non-uniform pruning, and expert selection. The objective, typically model loss, depends jointly on all assignments and does not decompose across groups, preventing combinatorial solvers from directly optimizing the true objective and forcing reliance on proxy formulations. Methods such as evolutionary search evaluate the actual loss but lack gradient information, while penalty-based approaches enforce the budget only approximately and often require extensive hyperparameter tuning. We present a new approach by showing that, under softmax relaxation, the budget constraint defines a smooth Riemannian manifold in logit space with unusually simple geometry. The normal vector admits a closed-form expression, shifting logits along the cost vector changes expected cost monotonically, and vector transport reduces to a single inner product. Building on these properties, we propose Riemannian Constrained Optimization (RCO), which augments a standard Adam step with tangent projection, binary-search retraction, and momentum transport. Combined with Gumbel straight-through estimation and budget-constrained dynamic programming for discrete feasibility, RCO enables first-order optimization of the actual loss under exact budget enforcement without introducing constraint-specific hyperparameters. Across both synthetic benchmarks and realistic LLM compression settings, RCO matches or exceeds state-of-the-art methods while often requiring substantially less wall-clock time. Source code is available at https://github.com/IST-DASLab/RCO.

模型压缩黎曼优化预算约束大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。