Model Compression with Exact Budget Constraints via Riemannian Manifolds

Assigning one of K options to each of N groups under a total cost budget is a recurring problem in efficient AI, for instance in mixed-precision quantization, non-uniform pruning, and expert selection. The objective (model loss) depends on all assignments jointly and does not decompose across groups, which means that combinatorial solvers can only optimize proxy objectives. Methods such as evolutionary search evaluate the actual loss but lack gradients, while penalty-based methods enforce the budget only approximately and can require heavy hyperparameter tuning. We show that under softmax relaxation, the budget constraint defines a smooth Riemannian manifold in logit space with unusually clean geometry. Specifically, the normal vector is available in closed form, shifting logits along the cost vector changes expected cost monotonically, and the vector transport reduces to a single inner product. Building on this, we propose Riemannian Constrained Optimization (RCO), which wraps tangent projection, binary-search retraction, and momentum transport around a standard Adam step. Combined with Gumbel straight-through estimation and budget-constrained dynamic programming for discrete feasibility, RCO provides first-order optimization of the actual loss with the expected budget enforced exactly at every iterate, and no constraint-related hyperparameters. The same construction handles several simultaneous budgets, whenever an assignment must respect more than one resource limit, with no added coefficients. RCO exceeds or matches the performance of state-of-the-art methods in both synthetic problems and realistic LLM compression settings, often at considerably lower wall-clock cost.