Model Compression with Exact Budget Constraints via Riemannian Manifolds

TL;DR

This paper introduces Riemannian Constrained Optimization (RCO), a novel method for model compression under exact budget constraints that leverages the geometry of a smooth manifold to enable efficient first-order optimization.

Contribution

The paper proposes RCO, a new optimization approach that enforces exact budget constraints in model compression by exploiting Riemannian manifold geometry, avoiding hyperparameter tuning.

Findings

RCO matches or exceeds state-of-the-art in synthetic and LLM benchmarks.

RCO often requires less wall-clock time than existing methods.

The method enables direct optimization of the true loss under exact budget constraints.

Abstract

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.