stat.MLJul 21, 2026

Boltzmann-Expected Molecular Design with Decoupled Annealing Flows

Authors: Selma MoqvistRichard BeckmannRoss IrwinRocío MercadoSimon Olsson

Organizations: Department of Computer Science and Engineering Chalmers University of Technology & University of Gothenburg SE-41296 Gothenburg, Sweden · Molecular AI Discovery Sciences, R&D AstraZeneca SE-43183 Mölndal, Sweden

Abstract

Most 3D properties relevant to molecular design, including free energies and shape descriptors, are expectations\textit{expectations} over the Boltzmann distribution over 3D configurations of a molecular graph. However, existing property-guided generative models tie each property to a single structure, ignoring the underlying ensemble. We recast 3D molecular design as Boltzmann-expected design\textbf{Boltzmann-expected design} and realise it with DECAF\textbf{DECAF} (Decoupled Annealing Flows), which factorise the joint distribution over graphs and coordinates into two conditional flow models: a graph-conditioned flow p(xG)p(x\mid\mathcal{G}), acting as a Boltzmann emulator\textit{Boltzmann emulator}, and a coordinate-conditioned flow p(Gx)p(\mathcal{G}\mid x), proposing new graphs from 3D information. By alternating the two flows, DECAF optimises molecular graphs with a simulated-annealing acceptance rule whose scoring function is evaluated on ensembles drawn from p(xG)p(x\mid\mathcal{G}), making ensemble statistics, not single-conformer properties, the design target. The resulting loop requires no retraining to change objectives. On GEOM-Drugs, we show that ensemble-aware optimisation produces graphs whose mean radius of gyration and solvent-accessible surface area consistently shift toward targets, while single-conformer optimisation degrades on larger drug-like molecules where Boltzmann distributions are broadest. DECAF extends to multi-objective trade-offs and, uniquely among 3D generative models, to higher-moment design\textbf{higher-moment design}: jointly optimising an ensemble property's variance and skewness to produce flexible molecules biased to a prescribed conformational regime: we verify the conformational distributions of these higher-moment designs with all-atom MD simulations.

Explore similar work

CardsList
  1. Ensemble-Conditioned Molecular Design

    Sep 14, 2026Ross Irwin, Alessandro Tibo, Jon Paul Janet +1Molecular GenerationConformational Ensembles

  2. Generating Physically Consistent Molecules with Energy-Based Models

    May 18, 2026Christoph Griesbacher, Lea Bogensperger, Andreas Habring +1Molecular GenerationFree-Energy Landscape