Budget Constraints as Riemannian Manifolds

2026-05-01

AI Summary

This work addresses the challenge of allocating K discrete options across N groups under a total cost budget when the objective function—such as model loss—is non-decomposable and thus difficult to optimize directly. The authors propose the first approach that models the budget constraint as a smooth Riemannian manifold in logit space, leveraging its geometric structure to enable first-order optimization that exactly satisfies the constraint. Their method integrates Riemannian gradients, tangent-space projection, bisection-based retraction, momentum transport, and Gumbel straight-through estimation, while incorporating dynamic programming to guarantee discrete feasibility—all without requiring hyperparameter tuning. Experiments demonstrate that the approach recovers optimal solutions on synthetic knapsack problems, outperforming penalty-based methods by 83%, and achieves performance on large language model compression tasks comparable to or better than evolutionary search, with 3–16× lower computational cost.

Abstract

Assigning one of K options to each of N groups under a total cost budget is a recurring problem in machine learning, appearing in mixed-precision quantization, non-uniform pruning, and expert selection. The objective (model loss) depends jointly on all assignments and does not decompose across groups, which prevents combinatorial solvers from optimizing the true objective directly and limits them to proxy objectives. Evolutionary search evaluates the actual loss but lacks gradient information, while penalty-based methods provide gradients but enforce the budget only approximately and require sensitive hyperparameter tuning. We observe that under softmax relaxation, the budget constraint defines a smooth Riemannian manifold in logit space with particularly simple geometry: the normal vector is available in closed form, shifting logits along the cost vector changes expected cost monotonically, allowing binary-search retraction, and vector transport reduces to a single inner product. Building on this structure, we propose Riemannian Constrained Optimization (RCO), which augments a standard Adam update 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 true objective under exact budget enforcement, without introducing constraint hyperparameters. On synthetic knapsack problems with known optima, the manifold-based constraint handling recovers optimal solutions, whereas penalty methods plateau at 83% of optimal. On LLM compression tasks, including mixed-precision quantization and MoE expert pruning, RCO matches or exceeds evolutionary search methods while requiring 3x to 16x lower wall-clock cost on the evaluated configurations.

Problem

• budget constraint
• combinatorial optimization
• non-decomposable objective
• Riemannian manifold
• discrete assignment

Innovation

• Riemannian manifold
• budget-constrained optimization
• softmax relaxation
• first-order optimization
• discrete feasibility