# A Multi-stage Constrained Optimization Framework for Data-driven Problems

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 45 of its 62 references that are in the map.
- First posted on arXiv 26 Jul 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Intriguing properties of neural networks](https://mlsanity.com/paper/1312.6199/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations](https://mlsanity.com/paper/1811.12359/) (2018)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Disentangling by Factorising](https://mlsanity.com/paper/1802.05983/) (2018)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Recent Advances in Adversarial Training for Adversarial Robustness](https://mlsanity.com/paper/2102.01356/) (2021)
- [Loss landscapes and optimization in over-parameterized non-linear systems and neural networks](https://mlsanity.com/paper/2003.00307/) (2020)
- [Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders](https://mlsanity.com/paper/1611.02648/) (2016)
- [Optimization of Molecules via Deep Reinforcement Learning](https://mlsanity.com/paper/1810.08678/) (2018)
- [Bayesian Deep Convolutional Encoder-Decoder Networks for Surrogate Modeling and Uncertainty Quantification](https://mlsanity.com/paper/1801.06879/) (2018)
- [Diagnosing and Enhancing VAE Models](https://mlsanity.com/paper/1903.05789/) (2019)
- [Constrained Graph Variational Autoencoders for Molecule Design](https://mlsanity.com/paper/1805.09076/) (2018)
- [DC3: A learning method for optimization with hard constraints](https://mlsanity.com/paper/2104.12225/) (2021)
- [End-to-End Constrained Optimization Learning: A Survey](https://mlsanity.com/paper/2103.16378/) (2021)
- [From Variational to Deterministic Autoencoders](https://mlsanity.com/paper/1903.12436/) (2019)
- [Disentangled Representation Learning](https://mlsanity.com/paper/2211.11695/) (2022)
- [Constrained Generation of Semantically Valid Graphs via Regularizing Variational Autoencoders](https://mlsanity.com/paper/1809.02630/) (2018)
- [Lagrangian Duality for Constrained Deep Learning](https://mlsanity.com/paper/2001.09394/) (2020)
- [To Compress or Not to Compress- Self-Supervised Learning and Information Theory: A Review](https://mlsanity.com/paper/2304.09355/) (2023)
- [Self-Supervised Primal-Dual Learning for Constrained Optimization](https://mlsanity.com/paper/2208.09046/) (2022)
- [Posterior Collapse and Latent Variable Non-identifiability](https://mlsanity.com/paper/2301.00537/) (2023)
- [Improving black-box optimization in VAE latent space using decoder uncertainty](https://mlsanity.com/paper/2107.00096/) (2021)
- [Generative Modeling of Regular and Irregular Time Series Data via Koopman VAEs](https://mlsanity.com/paper/2310.02619/) (2023)
- [Inverse Optimization: Theory and Applications](https://mlsanity.com/paper/2109.03920/) (2021)
- [Denoising Criterion for Variational Auto-Encoding Framework](https://mlsanity.com/paper/1511.06406/) (2015)

## Abstract (arXiv)

Variational autoencoders (VAEs) transform high-dimensional, often noisy data into a compact latent representation, making downstream optimization more tractable. Three challenges persist in VAE-based constrained optimization: (i) sampling effectively within the latent space, (ii) identifying the active decision variables that actually influence the objective and constraints, and (iii) enforcing constraints without destabilizing training. We propose a Multi-stage Constrained Optimization Framework (MCOF). First, an entropy-constrained VAE (EC-VAE) coupled with a feature selector embeds objective and constraint information into a designated subset of latent variables, so that optimization proceeds over a low-dimensional subspace while the remaining coordinates supply solution diversity. Second, a Uniform Transformation (UT) module applies a per-dimension probability integral transform, replacing the irregular aggregate posterior with a uniform distribution over a bounded box and mitigating posterior collapse and Gaussian mixture bias. Third, a constraint-priority filter method (CPFM) solves the resulting surrogate problem by alternating violation-reduction and objective-reduction steps under a filter acceptance test, returning solutions that are feasible for the learned surrogate to a specified tolerance without requiring multiplier estimation. Finally, unselected latent coordinates are resampled to generate diverse decodings of a single optimized solution. We validate MCOF on a synthetic problem, where we ablate each stage and recover the analytic optimum, and on a ZINC250k drug design task, where the generated molecules satisfy the imposed constraints and are entirely novel relative to the training set.