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
Abstract:Modern generative AI models, such as diffusion and flow matching models, can sample from rich data distributions. However, many applications, especially in science and engineering, require more than drawing samples from the model distribution: they require searching within this distribution for samples that optimise task-specific criteria. In this work, we propose O3 (Optimisation Over the Outputs of Generative Models), a method for sample-efficient black-box optimisation over continuous-variable diffusion and flow-matching models. O3 is built around surrogate latent spaces: low-dimensional Euclidean embeddings that can be extracted from a generative model without additional training. The resulting representations have controllable dimensionality and support the direct application of standard optimisation algorithms. We show, on image and protein design tasks, that surrogate-space optimisation finds substantially higher-scoring samples than standard sampling or optimisation in the original latent space. Our method is model- and optimiser-agnostic, incurs negligible additional cost over standard generation, and requires no retraining or fine-tuning of the generative model.
| Subjects: | Machine Learning (stat.ML); Machine Learning (cs.LG) |
| Cite as: | arXiv:2509.23800 [stat.ML] |
| (or arXiv:2509.23800v3 [stat.ML] for this version) | |
| https://doi.org/10.48550/arXiv.2509.23800 arXiv-issued DOI via DataCite |
Submission history
From: Samuel Willis [view email]
[v1]
Sun, 28 Sep 2025 10:50:06 UTC (45,935 KB)
[v2]
Mon, 15 Dec 2025 23:50:21 UTC (41,435 KB)
[v3]
Wed, 13 May 2026 11:18:10 UTC (36,142 KB)