Budget Constraints as Riemannian Manifolds

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.

• The paper introduces a novel geometric formulation by modeling budget constraints as a Riemannian manifold in logit space to ensure exact constraint enforcement.
• It employs tangent projection and binary search retraction, eliminating hyperparameter tuning while maintaining feasibility at every optimization step.
• Empirical results on knapsack problems and LLM compression demonstrate superior precision and efficiency over traditional penalty and Lagrangian methods.

Budget Constraints as Riemannian Manifolds: A Formal and Technical Overview

Introduction and Motivation

The optimization of discrete assignments subject to a total cost constraint is prevalent in machine learning—especially in contexts such as mixed-precision quantization, non-uniform pruning, and expert selection within large models. The non-decomposability of the objective across assignment groups introduces complexity: combinatorial methods enforce budgets exactly but only optimize proxy objectives, while penalty-based gradient methods offer efficiency but violate the constraint and require hyperparameter tuning. This paper introduces a geometric framework, parameterizing softmax-relaxed assignments and treating the budget constraint as a smooth Riemannian manifold in the logit space. This formulation enables exact budget enforcement with first-order optimization via Riemannian Constrained Optimization (RCO), obviating hyperparameter tuning and providing efficient, scalable algorithms for model compression tasks (2605.00649).

Geometry of the Budget Manifold

The core theoretical contribution establishes that the level set $$M = \{\boldsymbol{\alpha} : C(\boldsymbol{\alpha}) = B\}$$, where $C(\boldsymbol{\alpha})$ is the softmax-parameterized expected cost, forms a well-behaved Riemannian submanifold in logit space. The normal vector at each point is closed-form: $(\nabla C)_{ik} = w_i\, p_{ik}(c_k - E_{p_i}[c])$. This geometry ensures that tangent projection, monotonic retraction (via binary search), and vector transport reduce to computationally trivial steps—each requiring only inner products or scalar root-finding. These operations enable the optimizer to maintain feasibility after every iteration, in sharp contrast to penalty/Lagrangian schemes.

Figure 1: $M$ is preserved at each step by tangent projection, Adam update, and retraction via binary search along cost vector directions.

The monotonicity of the expected cost when shifting along the cost vector simplifies retraction to a scalar monotone root-finding problem, permitting highly efficient budget corrections. Extensions are provided for inequality constraints (via slack variables) and multiple simultaneous constraints (projecting out multiple normals), with the same geometric rationale and computational simplicity.

Algorithmic Framework: Riemannian Constrained Optimization (RCO)

RCO wraps tangent projection, monotonic retraction, and momentum transport around standard Adam steps. Discrete assignments are produced via Gumbel-STE sampling and budget-constrained dynamic programming (DP), ensuring discrete feasibility in the forward pass. The backward pass projects gradients onto the budget manifold, ensuring continuous feasibility and eliminating constraint-gradient bias before it can accumulate in optimizer momentum.

Specifically, each step:

• Projects the loss gradient onto the tangent space of the manifold.
• Applies an Adam update, which may leave the iterate off-manifold due to curvature or non-uniform scaling.
• Retracts back to the manifold via binary search along the cost vector.
• Transports optimizer momentum to the new tangent space.
• Handles discrete assignment feasibility using DP and differentiable STE surrogates.

This framework exactly enforces the budget at every step, with no hyperparameters controlling constraint strength or forgiveness, and negligible computational overhead.

Empirical Validation

The paper benchmarks manifold constraint handling against Lagrangian and penalty methods on synthetic multiple-choice knapsack problems (MCKP) with known optima. Manifold projection achieves exact constraint satisfaction ($|C(\boldsymbol{\alpha}) - B| < 10^{-8}$) throughout optimization, while Lagrangian methods exhibit persistent violations (up to $10^{-1}$). On hard instances, manifold optimization converges to the DP optimum up to $8\times$ closer than augmented Lagrangian approaches.

Figure 2: Manifold optimization outperforms penalty approaches in gap to DP optimum and exact constraint satisfaction on huge knapsack scenarios ($N=1000, K=32$).

Figure 3: Per-constraint violation traces for $m=16$ simultaneous constraints highlight orders-of-magnitude superiority in feasibility for RCO versus Lagrangian baselines.

Figure 4: Adversarial MCKP scenario demonstrates robustness of manifold projection under challenging cost structures.

Figure 5: Boundary-clustered scenario shows RCO attaining low-budget solutions where Lagrangian methods plateau.

In LLM compression (mixed-precision quantization, MoE expert pruning), RCO matches or exceeds evolutionary search baselines (e.g., EvoESAP, EvoPress) at $3$–$C(\boldsymbol{\alpha})$0 lower wall-clock cost across diverse configurations. For example, at high compression on Qwen3-8B, RCO improves FineWeb perplexity by 36% over HIGGS and 4% over EvoPress, with marked efficiency gains. In MoE expert pruning, nonuniform RCO allocation recovers up to 97% of HumanEval as compared to only 55% for uniform allocation, with sharp trade-offs observed depending on calibration domain.

Technical Implications and Claims

Key claims and results substantiated in the paper include:

• Exact budget enforcement: Unlike Lagrangian methods, the RCO maintains constraint satisfaction to within floating-point precision throughout optimization, independent of gradient estimator bias or optimizer state.
• Superior performance on non-decomposable objectives: On MCKP and LLM compression problems, RCO achieves closer convergence to combinatorial optima and outperforms surrogate methods, particularly at extreme compression where layer interactions are nontrivial.
• Hyperparameter-free constraint handling: The manifold projection does not require $C(\boldsymbol{\alpha})$1 or schedule tuning, which is an Achilles heel of penalty approaches.
• Efficient scaling to large parameter spaces: The modularity and trivial computational cost of the geometric operations enable scaling to hundreds or thousands of groups or constraint dimensions.
• Robustness to discrete and combinatorial assignment variance: Averaging Gumbel samples per step enhances optimization stability and quality, a unique facilitation enabled by the manifold structure.

Theoretical and Practical Implications

Practically, the approach enables efficient, scalable, and exact optimization in large-scale model compression and resource allocation tasks. Theoretically, it demonstrates the power of exploiting manifold geometry for constraint satisfaction, providing a blueprint for extending constrained optimization to other domains by identifying or constructing well-behaved constraint manifolds. In settings with multiple resource constraints or nonconvex discrete assignment landscapes, the geometric approach offers substantial improvement in feasibility, robustness, and solution quality.

The separation of discrete feasibility (via DP with Gumbel sampling) from continuous feasibility (via manifold projection) yields optimization architectures that are both modular and analytically tractable. The method's robustness to gradient bias and potential for modular composability with modern optimizers provides avenues for extension to constrained fine-tuning, reinforcement learning with hard constraints, and beyond.

Future work may include manifold-based enforcement of nonlinear constraints (e.g., inference latency), combinatorial optimization in stochastic settings, and rigorous convergence proofs in presence of biased gradient estimators such as STE.

Conclusion

The analysis and empirical results in "Budget Constraints as Riemannian Manifolds" (2605.00649) rigorously establish that the expected-cost level set under softmax is a structurally clean Riemannian manifold enabling exact and efficient budget enforcement. The RCO algorithm leverages this geometry, delivering strong empirical performance and scalability without the limitations of proxy objectives or hyperparameter dependency endemic to penalty methods. It sets a technical foundation for manifold-based constraint handling in large-scale optimization problems in machine learning and opens further avenues into combinatorial and geometric optimization for practical AI deployments.