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

2026-07-26

AI Summary

This work addresses three key challenges in constrained optimization with variational autoencoders: inefficient latent space sampling, difficulty in identifying critical variables, and training instability. To overcome these issues, the authors propose a Multi-stage Constrained Optimization Framework (MCOF) that integrates an entropy-constrained VAE to identify salient latent variables, applies a probability integral transform to uniformize the posterior distribution, and employs a constraint-priority filtering strategy to alternately optimize objectives and constraints. Diversity of solutions is preserved through resampling of non-critical variables. MCOF innovatively combines feature selection, posterior regularization, and multiplier-free constraint handling, enabling efficient optimization within a low-dimensional subspace while avoiding posterior collapse and Gaussian mixture bias. Experiments demonstrate that the method exactly recovers analytical optima on synthetic problems and generates fully novel, constraint-compliant molecular structures in the ZINC250k drug design task.

Abstract

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.

Problem

• constrained optimization
• variational autoencoders
• latent space sampling
• active decision variables
• constraint enforcement

Innovation

• Multi-stage Constrained Optimization
• Entropy-constrained VAE
• Uniform Transformation
• Constraint-priority Filter Method
• Latent Space Optimization