RCO: Riemannian Constrained Optimization
A reference implementation of an optimizer for budget-constrained discrete model compression, which relaxes a per-layer choice of compression option onto a Riemannian manifold so a fixed overall budget, such as total bits, is satisfied exactly rather than only approximately.
RCO addresses budget-constrained discrete assignment: choosing one of several options for each of many groups, such as a bit-width per model layer, to minimize an objective that does not decompose across groups, subject to an exact equality budget. It relaxes the discrete choice into softmax probabilities over per-group logits, so the set of assignments meeting a target budget forms a smooth constraint manifold.
Its optimizer wraps a standard Adam step with three manifold operations: a tangent projection that keeps each step budget-preserving using a closed-form gradient of the budget constraint, a retraction that corrects any drift off the constraint manifold by a binary search along the cost direction, and a vector transport that keeps Adam's momentum term tangent to the manifold across steps.
The discrete forward pass combines Gumbel noise with an exact knapsack solver to pick the best feasible assignment, with straight-through gradients flowing through the same perturbed logits; the tangent projection removes gradient components that would otherwise violate the constraint.