🤖 machine learning

Budget Constraints as Riemannian Manifolds

This paper proposes Riemannian Constrained Optimization (RCO), a novel framework that models budget constraints as smooth Riemannian manifolds to enable efficient, gradient-based optimization of non-decomposable objectives under exact budget enforcement, outperforming existing penalty and evolutionary methods in both solution quality and computational efficiency for tasks like mixed-precision quantization and expert pruning.

Original authors: Michael Helcig, Dan Alistarh

Published 2026-05-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Michael Helcig, Dan Alistarh

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are the head chef of a massive, high-end restaurant. You have a strict budget for the night, but you also have a menu with hundreds of dishes, and each dish can be made in several different ways (e.g., using premium ingredients, standard ingredients, or budget-friendly substitutes).

Your goal is to pick exactly one version of each dish to serve so that the total cost stays exactly within your budget, while making the overall quality of the meal as delicious as possible.

The problem? The quality of the meal isn't just the sum of individual dishes. If you pick a fancy steak, it might pair better with a specific wine, changing the "flavor profile" of the whole table. This makes the math incredibly hard: you can't just look at each dish in isolation; you have to solve a giant, tangled puzzle where every choice affects every other choice.

This is the exact problem machine learning engineers face when trying to compress huge AI models (like the ones that power chatbots). They need to decide how much to "shrink" or "prune" different parts of the model to fit a size limit (the budget) without ruining the model's intelligence (the quality).

Here is how the paper solves this, using a few creative analogies:

1. The Old Way: Guessing and Penalizing

Previously, engineers tried two main approaches, both of which were clumsy:

The "Penalty" Method: They told the computer, "Try to stay under budget, but if you go over, I'll give you a big 'fine' (a penalty score)." The problem is that the computer is bad at guessing the right fine. If the fine is too small, it ignores the budget. If it's too big, the computer gets scared and stops learning. It's like trying to teach a dog to sit by shouting "No!" at random volumes; the dog never learns the exact rule.
The "Evolutionary" Method: They let the computer try thousands of random combinations, keep the best ones, and repeat. This works well but is incredibly slow. It's like trying to find the best recipe by cooking every possible meal in the world and tasting them one by one. It takes forever.

2. The New Idea: The "Budget Manifold"

The authors realized that if you look at the problem through a specific mathematical lens (using something called "softmax"), the budget constraint isn't a messy wall you have to bounce off of. Instead, it's a smooth, curved surface (a manifold) that you can walk on.

Think of the budget not as a hard fence, but as a tightrope.

The Surface: Imagine a giant, invisible, curved trampoline that exists only where your total cost equals your budget exactly.
The Walk: The computer doesn't need to jump off the trampoline and hope it lands back on it. Instead, it walks along the surface.

3. How the New Method (RCO) Works

The paper proposes a new algorithm called Riemannian Constrained Optimization (RCO). Here is how it moves along that tightrope:

Step 1: The Tangent Step (Walking Forward): The computer calculates the direction that makes the meal tastier (the gradient). But instead of just walking that way, it projects that direction onto the surface of the tightrope. This ensures it never accidentally steps off the budget line.
Step 2: The Binary Search (The Magic Slide): Sometimes, even when walking carefully, you might drift slightly off the line. In other methods, you'd have to do a complex calculation to get back. Here, the authors found a "magic slide." Because of the specific math they used, they can just slide the entire meal plan up or down a single knob (a binary search) to land perfectly back on the budget line. It's like having a remote control that instantly fixes your balance.
Step 3: The Momentum (Keeping the Rhythm): When you walk on a curved surface, your direction changes. The algorithm has a special trick to "transport" its momentum (its memory of where it was going) so it doesn't get dizzy or lose its rhythm as it moves along the curve.

4. Why It's a Big Deal

The paper claims this method is a game-changer for two reasons:

It's Exact: Unlike the old "penalty" methods that often end up slightly over or under budget, this method stays exactly on the budget line at every single step. It's like a tightrope walker who never wobbles.
It's Fast: Because it uses gradients (mathematical directions) instead of random guessing, it finds the best solution much faster.
- The Result: On tests with synthetic puzzles, the old methods got stuck at 83% of the best possible score, while this new method found the perfect solution.
- Real World: When they tested it on compressing huge AI models (like cutting down the size of a "Large Language Model"), it matched or beat the results of the slow "evolutionary" methods, but did it 3 to 16 times faster.

Summary

The paper introduces a new way to solve "budget" problems in AI. Instead of treating the budget as a hard limit that breaks your calculations, they turned it into a smooth, walkable surface. By walking along this surface, the computer can find the perfect balance between cost and quality much faster and more accurately than before, without needing to guess or tune tricky settings. It's the difference between stumbling through a dark room trying to avoid furniture and walking confidently down a well-lit, perfectly paved path.

1. Problem Statement

The paper addresses the budget-constrained discrete assignment problem, a recurring challenge in efficient machine learning. The goal is to assign one of $K$ discrete options (e.g., quantization bitwidths, pruning decisions, or expert selection) to each of $N$ groups (e.g., neural network layers) to minimize a global objective (e.g., model loss) subject to a total cost budget (e.g., model size or latency).

Key Challenges:

Non-Decomposable Objective: The model loss depends jointly on all assignments (e.g., interactions between layers), meaning it cannot be written as a sum of independent per-group terms. This prevents standard Dynamic Programming (DP) from optimizing the actual loss directly; DP can only optimize "proxy" objectives (like layer-wise sensitivity).
Limitations of Existing Methods:
- Combinatorial Solvers (DP): Handle the budget exactly but require decomposable objectives, forcing the use of inaccurate proxy scores.
- Evolutionary Search: Evaluates the actual loss but lacks gradients, requiring hundreds of forward passes and suffering from high computational costs.
- Penalty/Lagrangian Methods: Add a penalty term (e.g., $\lambda(C-B)^2$ ) to the loss. These enforce the budget only approximately, often oscillate, and require sensitive hyperparameter tuning.

2. Methodology: Riemannian Constrained Optimization (RCO)

The authors propose a geometric approach by treating the budget constraint not as a penalty, but as a smooth Riemannian manifold in the space of logits.

A. The Budget Manifold

The authors parameterize the assignment distribution for each group $i$ using logits $\alpha_i \in \mathbb{R}^K$ , with probabilities $p_i = \text{softmax}(\alpha_i)$ . The total expected cost is defined as:
$C(\alpha) = \sum_{i=1}^N w_i \langle p_i, c \rangle$
where $w_i$ is the group weight and $c$ is the cost vector.

Manifold Structure: The level set $\mathcal{M} = \{ \alpha \mid C(\alpha) = B \}$ is proven to be a smooth $(NK-1)$-dimensional Riemannian submanifold.
Geometric Properties:
1. Closed-form Normal: The gradient $\nabla C(\alpha)$ is available in closed form and is non-zero (assuming distinct costs), defining the normal vector to the manifold.
2. Monotonic Retraction: Shifting logits along the cost vector changes the expected cost monotonically. This allows the algorithm to project any point back onto the manifold $\mathcal{M}$ using a simple binary search (or Newton's method) rather than complex iterative projections.
3. Cheap Vector Transport: Moving momentum from one tangent plane to another reduces to a single inner product operation.

B. The RCO Algorithm

The proposed algorithm wraps standard first-order optimization (Adam) with three geometric operations to enforce the budget exactly at every step:

Tangent Projection: The loss gradient is projected onto the tangent space of the manifold by removing the component along the normal vector $\nabla C(\alpha)$ . This ensures the optimizer only sees directions that preserve the budget.
Retraction: After an Adam step, the iterate may drift off the manifold. The algorithm shifts all logits along the cost vector ( $\alpha' = \alpha + t \cdot c$ ) and uses binary search to find $t$ such that $C(\alpha') = B$ .
Momentum Transport: The optimizer's momentum (first moment) is projected onto the new tangent space to maintain consistency.

C. Handling Discrete Feasibility (Forward Pass)

Since the manifold operates in continuous logit space, the paper combines RCO with:

Gumbel-STE (Straight-Through Estimator): To generate discrete assignments for the forward pass.
Budget-Constrained Dynamic Programming: To solve the discrete assignment problem (Multiple-Choice Knapsack Problem) exactly during the forward pass, ensuring the evaluated configuration satisfies the budget.
Mechanism: The forward pass uses DP to find the best discrete configuration $z^*$ given perturbed logits. The backward pass uses STE to flow gradients through the softmax probabilities of the perturbed logits, while the manifold projection ensures these gradients do not violate the budget constraint.

3. Key Contributions

Identification of the Budget Manifold: The authors identify that the level set of softmax-expected cost forms a Riemannian manifold with unique, computationally cheap geometric properties (closed-form normals, monotonic retraction).
RCO Algorithm: A novel optimization framework that enforces exact budget constraints with zero constraint hyperparameters (no $\lambda$ to tune). It combines tangent projection, binary-search retraction, and momentum transport around Adam.
Handling Inequality Constraints: The manifold is extended to inequality constraints ( $C(\alpha) \le B$ ) via a slack variable $s$ , allowing the optimizer to naturally utilize unused budget without wasting it on inferior options.
Discrete-Continuous Hybrid: The integration of Gumbel-STE and budget-constrained DP allows for the optimization of the actual non-decomposable model loss while maintaining exact discrete feasibility.

4. Experimental Results

The paper validates RCO on synthetic and real-world Large Language Model (LLM) compression tasks.

Synthetic Knapsack Problems (MCKP):
- On hard instances, RCO recovers the Dynamic Programming (DP) optimal solution.
- Lagrangian methods plateau at 83% of the optimal value and exhibit significant constraint violations ( $10^{-1}$ to $10^{-2}$ ), whereas RCO maintains violations near machine precision ( $10^{-9}$ ).
- RCO handles "cheap optimal" cases (where the best solution costs less than the budget) perfectly by utilizing the slack variable, whereas equality-constrained methods waste budget.
MoE Expert Pruning:
- Evaluated on OLMoE-1B-7B and Qwen3-30B-A3B.
- RCO matches or exceeds the accuracy of EvoESAP (an evolutionary search baseline) across multiple benchmarks (ARC, MMLU, etc.).
- Efficiency: RCO achieves these results in 3–16× lower wall-clock time (e.g., ~10 minutes vs. 5.8 hours for EvoESAP).
Mixed-Precision Quantization:
- Evaluated on Qwen3-8B (assigning bitwidths 2–8 to 252 layers).
- At high compression (2.25 bits), RCO reduces perplexity by 36% compared to HIGGS (a linear surrogate method) and 4% compared to EvoPress (evolutionary search).
- RCO converges to the same optimum as EvoPress at 2.5 bits but in ~10× less time (62 minutes vs. 11–14 hours).

5. Significance and Impact

Exact Constraint Enforcement: Unlike penalty methods, RCO guarantees the budget is met exactly at every iteration, eliminating the need for hyperparameter tuning of penalty weights.
Optimization of Non-Decomposable Losses: By separating the discrete feasibility (handled by DP in the forward pass) from the continuous optimization (handled by the manifold in the backward pass), RCO allows for the direct optimization of complex, non-decomposable model losses (like end-to-end accuracy) under strict budget constraints.
Computational Efficiency: The geometric operations (projection, retraction, transport) are computationally cheap (essentially inner products and binary searches), making the method scalable to large models where evolutionary search is prohibitively expensive.
Generalizability: The framework is applicable to any problem involving softmax parameterization with linear costs, extending beyond model compression to resource-constrained neural architecture search and other discrete assignment problems.

In summary, the paper reframes budget-constrained discrete optimization as a problem of navigating a specific Riemannian manifold, providing a theoretically grounded, hyperparameter-free, and highly efficient alternative to penalty-based and evolutionary approaches.