Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport

• URL: http://arxiv.org/abs/2405.15252v2
• Date: Sun, 02 Mar 2025 14:10:09 GMT
• Title: Accelerating 3D Molecule Generation via Jointly Geometric Optimal Transport
• Authors: Haokai Hong, Wanyu Lin, Kay Chen Tan,
• Abstract summary: This paper proposes a new 3D molecule generation framework, called GOAT, for fast and effective 3D molecule generation.<n>We formulate a geometric transport formula for measuring the cost of mapping multi-modal features between a base distribution and a target data distribution.<n>We show that GOAT enjoys the efficiency of solving geometric optimal transport, leading to a double speedup compared to the sub-optimal method.
• Score: 25.89095414975696
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: This paper proposes a new 3D molecule generation framework, called GOAT, for fast and effective 3D molecule generation based on the flow-matching optimal transport objective. Specifically, we formulate a geometric transport formula for measuring the cost of mapping multi-modal features (e.g., continuous atom coordinates and categorical atom types) between a base distribution and a target data distribution. Our formula is solved within a joint, equivariant, and smooth representation space. This is achieved by transforming the multi-modal features into a continuous latent space with equivariant networks. In addition, we find that identifying optimal distributional coupling is necessary for fast and effective transport between any two distributions. We further propose a mechanism for estimating and purifying optimal coupling to train the flow model with optimal transport. By doing so, GOAT can turn arbitrary distribution couplings into new deterministic couplings, leading to an estimated optimal transport plan for fast 3D molecule generation. The purification filters out the subpar molecules to ensure the ultimate generation quality. We theoretically and empirically prove that the proposed optimal coupling estimation and purification yield transport plan with non-increasing cost. Finally, extensive experiments show that GOAT enjoys the efficiency of solving geometric optimal transport, leading to a double speedup compared to the sub-optimal method while achieving the best generation quality regarding validity, uniqueness, and novelty. The code is available at https://github.com/WanyuGroup/ICLR2025-GOAT.

Related papers

• Towards Optimal Multi-draft Speculative Decoding [102.67837141152232]
  Multi-Draft Speculative Decoding (MDSD) is a recent approach where, when generating each token, a small draft model generates multiple drafts. This paper discusses the dual of the optimal transport problem, providing a way to efficiently compute the optimal acceptance rate.
  arXiv  Detail & Related papers  (2025-02-26T03:22:44Z)
• Conditional Synthesis of 3D Molecules with Time Correction Sampler [58.0834973489875]
  Time-Aware Conditional Synthesis (TACS) is a novel approach to conditional generation on diffusion models. It integrates adaptively controlled plug-and-play "online" guidance into a diffusion model, driving samples toward the desired properties.
  arXiv  Detail & Related papers  (2024-11-01T12:59:25Z)
• Dynamical Measure Transport and Neural PDE Solvers for Sampling [77.38204731939273]
  We tackle the task of sampling from a probability density as transporting a tractable density function to the target. We employ physics-informed neural networks (PINNs) to approximate the respective partial differential equations (PDEs) solutions. PINNs allow for simulation- and discretization-free optimization and can be trained very efficiently.
  arXiv  Detail & Related papers  (2024-07-10T17:39:50Z)
• Bisimulation Metrics are Optimal Transport Distances, and Can be Computed Efficiently [14.262270388108112]
  We propose a new framework for formulating optimal transport distances between Markov chains. We show that calculating optimal transport distances in the full space of joint distributions can be equivalently formulated as solving a linear program.
  arXiv  Detail & Related papers  (2024-06-06T13:25:14Z)
• Interaction-Force Transport Gradient Flows [45.05400562268213]
  This paper presents a new gradient flow dissipation geometry over non-negative and probability measures. Using a precise connection between the Hellinger geometry and the maximum mean discrepancy (MMD), we propose the interaction-force transport (IFT) gradient flows.
  arXiv  Detail & Related papers  (2024-05-27T11:46:14Z)
• Unified Generative Modeling of 3D Molecules via Bayesian Flow Networks [19.351562908683334]
  GeoBFN naturally fits molecule geometry by modeling diverse modalities in the differentiable parameter space of distributions. We demonstrate that GeoBFN achieves state-of-the-art performance on multiple 3D molecule generation benchmarks in terms of generation quality. GeoBFN can also conduct sampling with any number of steps to reach an optimal trade-off between efficiency and quality.
  arXiv  Detail & Related papers  (2024-03-17T08:40:06Z)
• Space-Time Diffusion Bridge [0.4527270266697462]
  We introduce a novel method for generating new synthetic samples independent and identically distributed from real probability distributions. We use space-time mixing strategies that extend across temporal and spatial dimensions. We validate the efficacy of our space-time diffusion approach with numerical experiments.
  arXiv  Detail & Related papers  (2024-02-13T23:26:11Z)
• Equivariant Flow Matching with Hybrid Probability Transport [69.11915545210393]
  Diffusion Models (DMs) have demonstrated effectiveness in generating feature-rich geometries. DMs typically suffer from unstable probability dynamics with inefficient sampling speed. We introduce geometric flow matching, which enjoys the advantages of both equivariant modeling and stabilized probability dynamics.
  arXiv  Detail & Related papers  (2023-12-12T11:13:13Z)
• Diffusion Bridge Mixture Transports, Schr\"odinger Bridge Problems and Generative Modeling [4.831663144935879]
  We propose a novel sampling-based iterative algorithm, the iterated diffusion bridge mixture (IDBM) procedure, aimed at solving the dynamic Schr"odinger bridge problem. The IDBM procedure exhibits the attractive property of realizing a valid transport between the target probability measures at each iteration.
  arXiv  Detail & Related papers  (2023-04-03T12:13:42Z)
• Sampling with Mollified Interaction Energy Descent [57.00583139477843]
  We present a new optimization-based method for sampling called mollified interaction energy descent (MIED) MIED minimizes a new class of energies on probability measures called mollified interaction energies (MIEs) We show experimentally that for unconstrained sampling problems our algorithm performs on par with existing particle-based algorithms like SVGD.
  arXiv  Detail & Related papers  (2022-10-24T16:54:18Z)
• Theory and Approximate Solvers for Branched Optimal Transport with Multiple Sources [12.139222986297263]
  Branched Optimal Transport (BOT) is a generalization of optimal transport in which transportation costs along an edge are subadditive. We show how to efficiently find the best geometry of a BOT network for many sources and sinks, given a topology.
  arXiv  Detail & Related papers  (2022-10-14T10:51:16Z)
• Manifold Interpolating Optimal-Transport Flows for Trajectory Inference [64.94020639760026]
  We present a method called Manifold Interpolating Optimal-Transport Flow (MIOFlow) MIOFlow learns, continuous population dynamics from static snapshot samples taken at sporadic timepoints. We evaluate our method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.
  arXiv  Detail & Related papers  (2022-06-29T22:19:03Z)
• GeoMol: Torsional Geometric Generation of Molecular 3D Conformer Ensembles [60.12186997181117]
  Prediction of a molecule's 3D conformer ensemble from the molecular graph holds a key role in areas of cheminformatics and drug discovery. Existing generative models have several drawbacks including lack of modeling important molecular geometry elements. We propose GeoMol, an end-to-end, non-autoregressive and SE(3)-invariant machine learning approach to generate 3D conformers.
  arXiv  Detail & Related papers  (2021-06-08T14:17:59Z)
• CO-Optimal Transport [19.267807479856575]
  Optimal transport (OT) is a powerful tool for finding correspondences and measuring similarity between two distributions. We propose a novel OT problem, named COOT for CO- Optimal Transport, that simultaneously optimize two transport maps between both samples and features. We demonstrate its versatility with two machine learning applications in heterogeneous domain adaptation and co-clustering/data summarization.
  arXiv  Detail & Related papers  (2020-02-10T13:33:15Z)
• A Near-Optimal Gradient Flow for Learning Neural Energy-Based Models [93.24030378630175]
  We propose a novel numerical scheme to optimize the gradient flows for learning energy-based models (EBMs) We derive a second-order Wasserstein gradient flow of the global relative entropy from Fokker-Planck equation. Compared with existing schemes, Wasserstein gradient flow is a smoother and near-optimal numerical scheme to approximate real data densities.
  arXiv  Detail & Related papers  (2019-10-31T02:26:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.