Verifier-Constrained Flow Expansion for Discovery Beyond the Data

TL;DR

This paper introduces Flow Expander, a method to adapt pre-trained flow models for scientific discovery by expanding their sampling domain beyond the training data, guided by verifiers to ensure sample validity.

Contribution

It proposes a verifier-constrained flow expansion framework with theoretical guarantees, enabling models to generate more diverse valid samples beyond the original data distribution.

Findings

Successfully expands molecular design space while maintaining validity.

Theoretically guarantees convergence of the flow expansion process.

Demonstrates increased diversity in generated samples in empirical tests.

Abstract

Flow and diffusion models are typically pre-trained on limited available data (e.g., molecular samples), covering only a fraction of the valid design space (e.g., the full molecular space). As a consequence, they tend to generate samples from only a narrow portion of the feasible domain. This is a fundamental limitation for scientific discovery applications, where one typically aims to sample valid designs beyond the available data distribution. To this end, we address the challenge of leveraging access to a verifier (e.g., an atomic bonds checker), to adapt a pre-trained flow model so that its induced density expands beyond regions of high data availability, while preserving samples validity. We introduce formal notions of strong and weak verifiers and propose algorithmic frameworks for global and local flow expansion via probability-space optimization. Then, we present Flow Expander (FE), a scalable mirror descent scheme that provably tackles both problems by verifier-constrained entropy maximization over the flow process noised state space. Next, we provide a thorough theoretical analysis of the proposed method, and state convergence guarantees under both idealized and general assumptions. Ultimately, we empirically evaluate our method on both illustrative, yet visually interpretable settings, and on a molecular design task showcasing the ability of FE to expand a pre-trained flow model increasing conformer diversity while preserving validity.