📊 statistics

Posterior Inference in Latent Space for Scalable Constrained Black-box Optimization

This paper proposes a scalable method for high-dimensional constrained black-box optimization that reformulates the problem as posterior inference within the latent space of flow-based generative models, utilizing outsourced diffusion models to efficiently sample promising candidates while avoiding mode collapse.

Original authors: Kiyoung Om, Kyuil Sim, Taeyoung Yun, Hyeongyu Kang, Jinkyoo Park

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kiyoung Om, Kyuil Sim, Taeyoung Yun, Hyeongyu Kang, Jinkyoo Park

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 trying to find the absolute best spot to build a new city in a massive, foggy wilderness. This wilderness represents a high-dimensional black-box optimization problem.

Here is the catch:

The Fog: You can't see the terrain clearly. You can only send out a scout to a specific spot, check if it's good, and get a report. You don't have a map.
The Traps: There are invisible "no-go zones" (constraints) scattered everywhere. If you build in one, the city collapses. These zones are hard to find and often take up most of the land.
The Goal: You want the highest possible hill (the best objective value) that isn't inside a no-go zone.

Traditional methods (like standard Bayesian Optimization) are like sending out scouts one by one, trying to guess the map. But when the wilderness is huge (high-dimensional) and the no-go zones are tricky, these scouts get lost, get stuck in small valleys, or waste time checking spots that are clearly bad.

Enter CiBO (Constrained Inference in Black-box Optimization).

The authors propose a new way to solve this using a clever two-step dance involving "generative models" (AI that learns patterns) and "latent space" (a simplified, hidden version of the world).

The Core Idea: The "Shadow World" Analogy

Imagine the real wilderness is a chaotic, jagged mountain range with deep, hidden caves (the no-go zones). Trying to navigate this directly is a nightmare.

The paper suggests creating a "Shadow World" (the latent space). In this shadow world, the jagged mountains are smoothed out into gentle, rolling hills. The scary, jagged edges of the no-go zones become soft, easy-to-detect curves.

CiBO works like this:

Phase 1: Learning the Map and the Rules

First, the AI builds two things based on the few spots the scouts have already checked:

The Shape Shifter (Flow-based Model): This learns the general shape of the "promising" areas. It understands where the scouts have been and what the terrain generally looks like.
The Oracle (Surrogate Models): These are smart predictors. One guesses how high the hill is at any spot, and others guess if a spot is a "no-go zone" (a constraint violation).

Crucially, the AI doesn't just look at the data; it reweights it. It pays extra attention to the spots that were high and safe, and ignores the ones that were low or dangerous.

Phase 2: The "Shadow" Search

Instead of sending scouts randomly or trying to climb the jagged real mountains directly, the AI does something smarter:

Translate to the Shadow World: It takes the "Shape Shifter" and projects the problem into the Shadow World. In this simplified space, the complex, jagged "no-go zones" and the "best spots" form a much smoother, easier-to-navigate landscape.
The Diffusion Sampler (The Guided Tour): The AI uses a special tool called a diffusion sampler (think of it as a very smart, guided tour guide). This guide doesn't just wander; it learns to walk specifically through the "Shadow World" to find the highest, safest hills.
- Why a Shadow World? Because in the real world, the "best spots" might be scattered in tiny, isolated islands (multi-modal). In the Shadow World, these islands are connected by smooth bridges, making it much easier for the AI to find them without getting stuck.
Project Back: Once the guide finds the best spot in the Shadow World, the AI translates that coordinate back to the real wilderness.
Filter and Repeat: The AI generates many potential spots, filters out the ones that look risky based on its "Oracle" predictions, and sends the best few real scouts to check them. Then, it updates its map and repeats the process.

Why is this a big deal?

Avoiding the "Trap": Traditional methods often get stuck in "mode collapse," where they keep finding the same small, safe spot and miss the truly amazing spots nearby. By searching in the smooth Shadow World, CiBO can explore widely without getting stuck.
Handling the Impossible: In high-dimensional problems (like designing a molecule with 200 parts or a robot policy with 100 variables), the "no-go zones" are so complex that standard math breaks down. CiBO bypasses this by doing the hard math in the simplified Shadow World.
Speed and Efficiency: The paper shows that CiBO finds better solutions with fewer "scout trips" (evaluations) than older methods, even in tasks with hundreds of dimensions and dozens of constraints.

The Bottom Line

The paper claims that by treating the search for the best solution as a posterior inference problem (finding the most likely good spot given what we know) and doing that search in a simplified, smoothed-out "Shadow World," we can solve complex engineering and scientific problems much faster and more reliably than before.

They tested this on synthetic math puzzles and real-world challenges like planning rover paths, designing robot movements, and optimizing vehicle mass. In every case, their method found better, safer solutions faster than the competition.

Technical Summary: Posterior Inference in Latent Space for Scalable Constrained Black-box Optimization

Problem Definition

The paper addresses constrained black-box optimization (CBBO), a fundamental task in scientific and engineering domains such as machine learning, drug discovery, and industrial design. The goal is to find an input $x^*$ that maximizes an objective function $f(x)$ subject to $M$ black-box constraints $g^{(m)}(x) \leq 0$ .

The authors highlight that CBBO is significantly harder than unconstrained optimization due to two main factors:

Intractable Feasible Regions: Feasible regions are analytically undefined and often occupy a tiny fraction of the search space.
High Dimensionality: Existing methods, particularly Bayesian Optimization (BO), struggle to scale to high-dimensional spaces (e.g., $D=200$ ). While BO handles constraints via new acquisition functions or Lagrangian relaxation, it suffers from poor scalability and difficulty in accurately estimating surrogate models in complex, constrained landscapes.

Generative models have emerged as an alternative, but sampling from the target distribution of feasible, high-scoring candidates is often intractable in high dimensions due to multi-modality and large flat regions. Standard approaches like Markov Chain Monte Carlo (MCMC) scale poorly, and fine-tuning generative models often leads to mode collapse.

Methodology: CiBO

The authors propose CiBO (Constrained Inference-based Black-box Optimization), a framework that reformulates candidate selection as posterior inference performed in the latent space of generative models. The method operates in an iterative loop consisting of two phases:

Phase 1: Training Models

Given a dataset $\mathcal{D}_r$ of observed inputs, objectives, and constraint violations:

Prior Training: A flow-based model $p_\theta(x)$ is trained to capture the distribution of the current dataset. This is done using Flow Matching with reweighted training, where data points are weighted based on a Lagrangian-augmented score $l(y, c) = y - \lambda \sum \max(0, c_m)$ . This focuses the prior on promising, feasible regions.
Surrogate Training: An ensemble of $K$ neural networks is trained to predict objective values ( $\mu_\phi, \sigma_\phi$ ) to quantify uncertainty. Separate proxies are trained for each constraint $g^{(m)}$ . These are also trained using the reweighted loss.

Phase 2: Sampling Candidates

Instead of fine-tuning the flow-based prior (which risks mode collapse), CiBO casts candidate selection as sampling from a posterior distribution:
$p_{post}(x) \propto p_\theta(x) \exp(\beta \cdot r_\phi(x))$
where $r_\phi(x)$ is a reward function combining the surrogate's mean, uncertainty bonus, and constraint penalty.

Key Innovation: Latent Space Amortized Inference
The authors observe that while $p_{post}(x)$ in data space is highly multi-modal and complex, the corresponding posterior in the latent space ( $z$ ) is significantly smoother.

Outsourced Diffusion Sampler: They train a separate diffusion model $p_\psi(z)$ to approximate the latent posterior $p_{post}(z) \propto p(z) \exp(\beta \cdot r_\phi(f_\theta(z)))$ .
Trajectory Balance (TB): The diffusion sampler is trained using the Trajectory Balance objective, which allows for off-policy training. This is crucial for covering multiple modes and avoiding mode collapse in the latent space.
Projection and Filtering: Samples are drawn from $p_\psi(z)$ , projected to data space via the deterministic mapping $x = f_\theta(z)$ , and filtered to select the top- $B$ candidates based on the reward function $r_\phi(x)$ .

Key Contributions

Reformulation as Posterior Inference: The paper formalizes constrained candidate selection as sampling from a posterior distribution, leveraging the smoothness of the latent space to bypass the difficulties of direct sampling in high-dimensional data space.
Latent Space Amortization: By decoupling the sampling process into a noise generation step in latent space and a deterministic projection, the method avoids the mode collapse issues associated with fine-tuning flow-based models directly on complex, multi-modal target distributions.
Scalability: The use of lightweight MLP-based surrogates and flow/diffusion models allows the method to scale to high dimensions (up to 180 dimensions in real-world tasks) where traditional BO methods fail.
Indicator Constraints: The framework is shown to handle "indicator constraints" (binary feasibility feedback only), a challenging setting where gradient-based or continuous constraint information is unavailable.

Experimental Results

The authors evaluate CiBO on three synthetic benchmarks (Rastrigin, Ackley, Rosenbrock at $D=200$ ) and four real-world tasks (Rover Planning, HalfCheetah, Mopta, Lasso DNA) with dimensions ranging from 60 to 180.

Performance: CiBO consistently outperforms state-of-the-art baselines, including BO variants (cEI, LogcEI, SCBO, PCAGP-SCBO, FuRBO), evolutionary algorithms (CMA-ES), and other generative approaches (DiffOPT).
Feasibility: In synthetic tasks, CiBO achieves near-perfect feasibility ratios within 5–10 batches, significantly faster than BO baselines. In real-world tasks, it is often the only method to find feasible solutions on difficult benchmarks like Mopta and DNA-180D.
Sample Efficiency: The method demonstrates superior sample efficiency, finding high-quality feasible solutions with fewer evaluations compared to competitors.
Ablation Studies: Experiments confirm that reweighted training, filtering, and off-policy training of the diffusion sampler are critical components. The method is robust to variations in hyperparameters like the Lagrange multiplier $\lambda$ and buffer size $L$ .
Runtime: Despite training multiple models per round, CiBO maintains a per-round runtime comparable to or lower than BO-based methods because it avoids the expensive Gaussian Process (GP) fitting required by traditional BO.

Significance and Claims

The paper claims that CiBO represents a significant step forward in scalable constrained black-box optimization. Its primary significance lies in:

Overcoming the "Curse of Dimensionality" in Constraints: By performing inference in the latent space, the method effectively navigates complex, multi-modal feasible regions that are intractable for standard BO or direct generative model fine-tuning.
Robustness: The framework is robust across diverse tasks, including those with binary constraint feedback and a large number of constraints.
Practical Applicability: The method is computationally efficient enough for practical deployment in real-world design and control problems where evaluation budgets are limited.

The authors modestly note limitations, specifically the need to retrain models at every round and the potential for future improvements by leveraging advancements in flow-based training and diffusion samplers. They emphasize that their work provides a foundation for more robust, constraint-aware optimization frameworks.