📊 statistics

Sample-Efficient Optimisation over the Outputs of Generative Models

This paper introduces O3, a model-agnostic method that enables sample-efficient black-box optimization over continuous-variable generative models by leveraging low-dimensional, training-free surrogate latent spaces to find superior task-specific samples without retraining the underlying model.

Original authors: Samuel Willis, Paul Duckworth, Jack Simons, Aleksandra Kalisz, Krisztina Sinkovics, Noam Ghenassia, Shikha Surana, Henry T. Oldroyd, Alexandru I. Stere, Dragos D Margineantu, Carl Henrik Ek, Henry Mos

Published 2026-05-14

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Samuel Willis, Paul Duckworth, Jack Simons, Aleksandra Kalisz, Krisztina Sinkovics, Noam Ghenassia, Shikha Surana, Henry T. Oldroyd, Alexandru I. Stere, Dragos D Margineantu, Carl Henrik Ek, Henry Moss, Erik Bodin

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

The Big Problem: Finding a Needle in a Haystack

Imagine you have a magical, infinite library (a Generative AI Model) that can create any picture, song, or protein structure you can imagine. You ask it to "draw a cat," and it does. You ask for "a cat wearing a hat," and it does that too.

But what if you don't just want any cat? What if you need a cat that is specifically orange, sleeping on a red rug, and looking at the moon?

If you just keep asking the library to "draw a cat" and hoping one of those millions of random cats matches your specific description, you might wait forever. This is called sampling. It's like throwing darts at a giant board hoping to hit a tiny bullseye. It's slow, expensive, and inefficient.

The Old Way: Trying to Navigate the Library

Previously, scientists tried to solve this by building a map of the library's "back room" (the Latent Space). They thought, "If we can find the right coordinates in the back room, we can walk straight to the orange cat."

However, modern AI libraries (like Diffusion models) are so huge and complex that their "back rooms" are confusing mazes. If you try to walk through them blindly, you often end up in a wall or a room that produces garbage (like a black screen or a distorted blob).

The New Solution: O3 (The "Surrogate Map")

The authors of this paper invented a new method called O3. Think of O3 as building a small, manageable, 2D map (a "Surrogate Latent Space") that sits on top of the giant, confusing library.

Here is how it works, step-by-step:

1. The "Seed" Ingredients

Instead of trying to understand the whole library, O3 starts with just a few specific examples. Let's say you want that orange cat on the red rug.

You find three pictures that are close to what you want (maybe one has an orange cat, one has a red rug, and one has a moon).
These are your "Seed Latents." They are like three specific ingredients you have in your kitchen.

2. The "Blender" (The Surrogate Space)

O3 creates a small, simple control panel (a low-dimensional grid) where you can mix these three seeds together.

Imagine a blender where you can adjust the knobs to mix 50% of the orange cat, 30% of the red rug, and 20% of the moon.
This "blender" is the Surrogate Space. It is tiny and easy to navigate, unlike the giant library.

3. The "Magic Translator"

The paper's secret sauce is a special mathematical trick (called a Surrogate Chart) that translates your simple knob-turning in the blender into a valid request for the giant library.

When you turn the knobs, O3 doesn't just mash the images together randomly. It uses a precise formula to ensure the library understands the request and produces a realistic image.
Crucially: You don't have to retrain the library or teach it anything new. It just works with the library you already have.

Why This is a Game-Changer

The paper claims three main benefits, explained simply:

1. It's a "No-Training" Shortcut
Most other methods require you to teach the AI a new skill (fine-tuning) to find your specific cat. That takes a lot of time and money. O3 is like using a universal remote control. You don't need to rebuild the TV; you just need the right remote to find the channel you want instantly.

2. It Finds Better Results Faster
The authors tested this on images and protein designs (the building blocks of life).

The Test: They asked the AI to find images that matched very specific, hidden descriptions.
The Result: Using O3, they found "perfect" matches in just a few tries. Using the old "throw darts" method, they would have needed thousands of tries to find the same thing.
The Analogy: If finding the perfect image is like finding a specific book in a library, the old way is reading every book on the shelf. O3 is like having a librarian who knows exactly which three books to mix together to write your new book instantly.

3. It Works on Anything
The method isn't picky. They showed it works on:

Images (drawing cats and cars).
Audio (making music).
Video (creating short clips).
Proteins (designing molecules for science).
It's like a universal key that fits many different locks.

The "Magic" Math (Simplified)

The paper mentions some complex math about "spheres" and "cosine similarity." Here is the simple version:

Imagine the library's "back room" is a giant sphere.
If you move a tiny bit on the surface of that sphere, the picture changes just a tiny bit.
O3 flattens that curved sphere into a flat, square piece of paper (the Surrogate Space) so you can use standard tools to search for the best spot. It's like flattening a globe into a map so you can draw a straight line to your destination.

Summary

O3 is a tool that lets you take a few "good examples" and use them to build a small, easy-to-search map. This map guides a powerful AI to create exactly what you want without needing to retrain the AI or wait for it to randomly guess. It turns a "search for a needle in a haystack" problem into "mixing three ingredients to make a perfect cake."

What the paper doesn't claim:

It does not claim to create new types of AI.
It does not claim to fix safety issues (if the AI can make bad things, O3 can help find those bad things faster, too).
It does not claim to work on every possible problem, but it works on the specific types of "black-box" problems where you can't see the math inside the AI, only the results.

Technical Summary: Sample-Efficient Optimisation over the Outputs of Generative Models (O3)

Problem Statement

Modern generative AI models, particularly diffusion and flow-matching models, excel at sampling from rich data distributions. However, many scientific and engineering applications require more than random sampling; they necessitate searching within the model's distribution to find samples that optimize specific, often expensive, black-box objectives (e.g., protein stability, prompt alignment).

Traditional approaches to this "Latent Space Optimisation" (LSO) face significant hurdles when applied to modern stochastic interpolant models:

Dimensionality: Unlike Variational Autoencoders (VAEs), diffusion and flow-matching models often operate in latent spaces with dimensions comparable to the data itself, making direct optimization intractable.
Statistical Mismatch: Naive linear combinations of latent vectors often violate the statistical properties required by the generative model, leading to poor or implausible outputs.
Retraining Costs: Existing methods for optimizing these models often require fine-tuning the generative model itself or training auxiliary models, which is computationally expensive and limits generalizability.

Methodology: O3 (Optimisation Over the Outputs of Generative Models)

The paper proposes O3, a framework for sample-efficient black-box optimization that operates directly on the outputs of pre-trained continuous-variable generative models without retraining or fine-tuning. The core innovation is the construction of surrogate latent spaces.

1. Surrogate Latent Spaces

O3 constructs a low-dimensional, bounded Euclidean search space $U$ that maps to valid, high-quality generations. The process involves three key steps:

Seed Latents: The method begins by selecting $K$ "seed" latent vectors ( $z_1, \dots, z_K$ ) from the model's latent distribution. These seeds can be chosen randomly, filtered based on high scores, or obtained via inversion of known data points.
Linear Combination with Statistical Preservation: To ensure generated samples remain within the model's reliable operating regime, O3 employs Latent Optimal Linear combinations (LOL). Instead of arbitrary linear combinations, it constructs new latents $z^*$ as unit-norm linear combinations of the seeds:
$z^* = \sum_{i=1}^K w_i z_i \quad \text{where} \quad \sum w_i^2 = 1$
This constraint ensures that the resulting latent $z^*$ maintains the statistical distribution (e.g., $N(0, I)$ ) expected by the generative model, preventing the "mismatch" that plagues methods like REMBO.
Surrogate Chart ( $\phi_w$ ): To avoid the redundancy of representing the same latent with different scalar multiples of weights and to facilitate standard optimization, the unit-norm weights are mapped from a bounded Euclidean space $U = [0, 1]^{K-1}$ $U = [0, 1]^{K - 1}$ to the positive orthant of the unit hypersphere $S^{K-1}_+$ $S_{+}^{K - 1}$ .
- The paper utilizes a variant of the Knothe–Rosenblatt (KR) chart. This bijection flattens the hypersphere, allowing standard black-box optimizers (like Bayesian Optimization or CMA-ES) to operate in a simple Euclidean space while preserving the geometric structure (cosine similarity) of the latent space.

2. Workflow

Initialization: Select $K$ seed latents.
Optimization: Run a standard optimizer (e.g., BO, CMA-ES) in the surrogate space $U$ .
Generation: For each candidate $u \in U$ , map to weights $w = \phi_w(u)$ , compute the latent $z^* = \sum w_i z_i$ , and decode $x = g(z^*)$ .
Evaluation: Score $x$ using the black-box objective.

Key Contributions

Model-Agnostic Framework: O3 is applicable to any continuous-variable generative model with a latent distribution amenable to interpolation (e.g., diffusion, flow-matching), requiring no modification to the pre-trained model.
Statistical Validity: By enforcing unit-norm constraints on seed combinations, O3 guarantees that search proposals adhere to the generative model's latent statistics, a requirement identified by Bodin et al. [5] for reliable generation.
Surrogate Geometry: The introduction of the KR chart provides a bounded, unconstrained Euclidean search space that preserves the similarity structure of the latent manifold, enabling the use of standard optimizers.
Zero Retraining Cost: Unlike methods requiring fine-tuning (e.g., DPO, Adjoint Matching) or auxiliary models, O3 incurs negligible computational overhead beyond standard generation.

Experimental Results

The authors validate O3 across image, audio, video, and protein design tasks:

Image Generation (Flux & Stable Diffusion):
- O3 significantly outperforms standard "Best-of-N" sampling and methods that do not preserve latent statistics (like REMBO).
- In high-dimensional surrogate spaces (e.g., 89D), O3 with Bayesian Optimization finds substantially higher-scoring images than lower-dimensional alternatives or random search.
- O3 achieves performance comparable to fine-tuned models (e.g., Importance Fine-tuning, DPOK) without the associated training costs.
Protein Design (BOLTZ-2 & RFDIFFUSION):
- Conformation Biasing: When biasing BOLTZ-2 toward a specific calmodulin conformation (1CLL), O3 outperforms FK-steering and DPO at low oracle budgets ( $N \le 100$ ) and remains competitive at higher budgets, all without fine-tuning.
- De Novo Design: For generating 600-residue proteins with RFDIFFUSION, standard sampling almost never yields designs meeting the recovery criteria (RMSE < 2.0 Å). O3, using CMA-ES in the surrogate space, reliably finds successful designs (median RMSE ~1.08 Å) even with random seed latents.
Generalization: The method successfully constructs smooth, low-dimensional manifolds for audio (StableAudio2.0) and video (HunyuanVideo), demonstrating cross-modal applicability.

Significance and Claims

The paper positions O3 as a general, "plug-and-play" layer for optimizing pre-trained generative models. Its primary significance lies in decoupling the search process from model training.

Efficiency: It enables highly sample-efficient optimization for expensive, non-differentiable objectives (common in science and engineering) without the massive compute costs of retraining or fine-tuning large models.
Accessibility: By removing the need for task-specific model adaptation, O3 lowers the barrier for users to steer generative models toward specific scientific or engineering goals.
Robustness: The method relies on the geometric properties of the latent space rather than learning new parameters, ensuring that the optimization remains within the "safe" regime of the pre-trained model.

The authors conclude that surrogate latent spaces provide a tractable, structured manifold for optimization across diverse modalities, offering a practical alternative to fine-tuning for black-box optimization tasks. Future work is suggested in adaptive seed selection and multi-objective optimization.