# Model Compression with Exact Budget Constraints via Riemannian Manifolds

arXiv: https://arxiv.org/abs/2605.00649 · HTML: https://mlsanity.com/paper/2605.00649/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 45 references that are in the map.
- First posted on arXiv 1 May 2026.

## Builds on (most cited first, 30 of 33)

- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Constrained Policy Optimization](https://mlsanity.com/paper/1705.10528/) (2017)
- [MnasNet: Platform-Aware Neural Architecture Search for Mobile](https://mlsanity.com/paper/1807.11626/) (2018)
- [Poincar'e Embeddings for Learning Hierarchical Representations](https://mlsanity.com/paper/1705.08039/) (2017)
- [ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware](https://mlsanity.com/paper/1812.00332/) (2018)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [FBNet: Hardware-Aware Efficient ConvNet Design via Differentiable Neural Architecture Search](https://mlsanity.com/paper/1812.03443/) (2018)
- [HAQ: Hardware-Aware Automated Quantization with Mixed Precision](https://mlsanity.com/paper/1811.08886/) (2018)
- [Reward Constrained Policy Optimization](https://mlsanity.com/paper/1805.11074/) (2018)
- [HAWQ: Hessian AWare Quantization of Neural Networks with Mixed-Precision](https://mlsanity.com/paper/1905.03696/) (2019)
- [Responsive Safety in Reinforcement Learning by PID Lagrangian Methods](https://mlsanity.com/paper/2007.03964/) (2020)
- [HAWQ-V2: Hessian Aware trace-Weighted Quantization of Neural Networks](https://mlsanity.com/paper/1911.03852/) (2019)
- [Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets](https://mlsanity.com/paper/1903.05662/) (2019)
- [Differentiation of Blackbox Combinatorial Solvers](https://mlsanity.com/paper/1912.02175/) (2019)
- [Riemannian Adaptive Optimization Methods](https://mlsanity.com/paper/1810.00760/) (2018)
- [Outlier Weighed Layerwise Sparsity (OWL): A Missing Secret Sauce for Pruning LLMs to High Sparsity](https://mlsanity.com/paper/2310.05175/) (2023)
- [Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models](https://mlsanity.com/paper/2402.14800/) (2024)
- [First-order Methods for Geodesically Convex Optimization](https://mlsanity.com/paper/1602.06053/) (2016)
- [Learning with Differentiable Perturbed Optimizers](https://mlsanity.com/paper/2002.08676/) (2020)
- [REAP the Experts: Why Pruning Prevails for One-Shot MoE compression](https://mlsanity.com/paper/2510.13999/) (2025)
- [SPDY: Accurate Pruning with Speedup Guarantees](https://mlsanity.com/paper/2201.13096/) (2022)
- [Bridging Discrete and Backpropagation: Straight-Through and Beyond](https://mlsanity.com/paper/2304.08612/) (2023)
- [MxMoE: Mixed-precision Quantization for MoE with Accuracy and Performance Co-Design](https://mlsanity.com/paper/2505.05799/) (2025)
- [SDQ: Stochastic Differentiable Quantization with Mixed Precision](https://mlsanity.com/paper/2206.04459/) (2022)
- [HardCoRe-NAS: Hard Constrained diffeRentiable Neural Architecture Search](https://mlsanity.com/paper/2102.11646/) (2021)
- [EvoPress: Accurate Dynamic Model Compression via Evolutionary Search](https://mlsanity.com/paper/2410.14649/) (2024)

## Built on by (most cited first, 3 of 3)

- [Statistically-Lossless Quantization of Large Language Models](https://mlsanity.com/paper/2605.02404/) (2026)
- [Auditing Invisible Weight Updates with Reference Traces](https://mlsanity.com/paper/2607.09800/) (2026)
- [Q-PACE: Dynamic Precision Allocation for Quantization-Aware Training](https://mlsanity.com/paper/2610.09183/) (2026)

## Abstract (arXiv)

Assigning one of K options to each of N groups under a total cost budget is a recurring problem in efficient AI, arising 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, so combinatorial solvers can only optimize proxy objectives. Evolutionary search evaluates the actual loss but lacks gradients, while penalty-based methods enforce the budget only approximately and often 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: the normal vector is available in closed form, shifting logits along the cost vector changes expected cost monotonically, and 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 optimizes the actual loss with first-order methods, enforces the expected budget exactly at every iterate, and introduces no constraint-related hyperparameters. The same construction handles multiple simultaneous budgets with no additional coefficients. RCO matches or exceeds state-of-the-art methods on synthetic problems and realistic LLM compression settings, often at considerably lower wall-clock cost. Source code is available at https://github.com/IST-DASLab/RCO.