Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GAGA: Gaussianity-Aware Gaussian Approximation for Efficient 3D Molecular Generation

Jingxiang Qu, Wenhan Gao, Ruichen Xu, Yi Liu

TL;DR

A principled method to improve generation efficiency without sacrificing training granularity or inference fidelity of GPPGMs, and identifies a characteristic step at which molecular data attains sufficient Gaussianity, after which the trajectory can be replaced by a closed-form Gaussian approximation.

Abstract

Gaussian Probability Path based Generative Models (GPPGMs) generate data by reversing a stochastic process that progressively corrupts samples with Gaussian noise. Despite state-of-the-art results in 3D molecular generation, their deployment is hindered by the high cost of long generative trajectories, often requiring hundreds to thousands of steps during training and sampling. In this work, we propose a principled method, named GAGA, to improve generation efficiency without sacrificing training granularity or inference fidelity of GPPGMs. Our key insight is that different data modalities obtain sufficient Gaussianity at markedly different steps during the forward process. Based on this observation, we analytically identify a characteristic step at which molecular data attains sufficient Gaussianity, after which the trajectory can be replaced by a closed-form Gaussian approximation. Unlike existing accelerators that coarsen or reformulate trajectories, our approach preserves full-resolution learning dynamics while avoiding redundant transport through truncated distributional states. Experiments on 3D molecular generation benchmarks demonstrate that our GAGA achieves substantial improvement on both generation quality and computational efficiency.

GAGA: Gaussianity-Aware Gaussian Approximation for Efficient 3D Molecular Generation

TL;DR

A principled method to improve generation efficiency without sacrificing training granularity or inference fidelity of GPPGMs, and identifies a characteristic step at which molecular data attains sufficient Gaussianity, after which the trajectory can be replaced by a closed-form Gaussian approximation.

Abstract

Gaussian Probability Path based Generative Models (GPPGMs) generate data by reversing a stochastic process that progressively corrupts samples with Gaussian noise. Despite state-of-the-art results in 3D molecular generation, their deployment is hindered by the high cost of long generative trajectories, often requiring hundreds to thousands of steps during training and sampling. In this work, we propose a principled method, named GAGA, to improve generation efficiency without sacrificing training granularity or inference fidelity of GPPGMs. Our key insight is that different data modalities obtain sufficient Gaussianity at markedly different steps during the forward process. Based on this observation, we analytically identify a characteristic step at which molecular data attains sufficient Gaussianity, after which the trajectory can be replaced by a closed-form Gaussian approximation. Unlike existing accelerators that coarsen or reformulate trajectories, our approach preserves full-resolution learning dynamics while avoiding redundant transport through truncated distributional states. Experiments on 3D molecular generation benchmarks demonstrate that our GAGA achieves substantial improvement on both generation quality and computational efficiency.

Paper Structure

This paper contains 46 sections, 4 theorems, 46 equations, 2 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Gaussian Probability Path based Generative Models
  4. Learning the Reverse Process
  5. Zero-Mean Invariance
  6. Gaussianity-Aware Gaussian Approximation
  7. Gaussian Approximation
  8. Gaussian Approximation and Initial Data Distribution
  9. Evaluating Gaussianity: Data Dependency and Distributional Similarity
  10. Data Dependency Decay.
  11. Distributional Similarity.
  12. Estimation of $T^*$.
  13. Merits of GAGA.
  14. Experiments
  15. Experimental Setup
  16. ...and 31 more sections

Key Result

Proposition 3.1

Given $t\in[0,T)$ and $K\ge3$, and the Gaussianity evaluation functional where $\mathbf{1}_{\{\cdot\}}$ is the indicator function, $\beta>0$ and $w_k>0$ ($k\ge3$). $\mathsf{D}:=\{\operatorname{Diag}(v):v\in\mathbb{R}^d\}$ is the diagonal subspace and are the orthogonal projections. $Cov(\cdot)$ and $C^{(k)}(X)$ are the covariance calculator and the $k$-th cumulant tensor, respectively. Let $A, B

Figures (2)

  • Figure 1: The flowchart of the GAGA, where the forward process is discretized to $T$ steps. In such a case, when the noised data distribution ${\bm{x}}_t$ has attained sufficient Gaussianity at timestep $t$, we approximate it with a reference Gaussian $\mathcal{N}(\tilde{\mu}_t, \tilde{v}_t)$. Therefore, the length of the generative trajectory can be reduced from $T$ steps to $t$ steps.
  • Figure 2: Comparisons of the forward noising process across different data modalities. (a) shows a continuous-valued image matrix, while (b) and (c) illustrate the distribution of molecular data consisting of one-hot vectors for atom types and 3D Euclidean coordinates for atom positions, respectively. The same schedule for gaussian probability path is applied across all modalities, with the number of forward timesteps $T$ up to 1000. Despite identical signal-to-noise ratios, molecular data obtains sufficient Gaussianity for significantly fewer steps compared to image data. This comes from the different Gaussianity across initial data distributions, as empirically quantified in Appendix \ref{['Appendix:Initial_gaussianity']}.

Theorems & Definitions (7)

  • Proposition 3.1
  • Lemma C.1: Variance preserving propagation of moments/cumulants and contraction of $\mathcal{H}^{(K)}$
  • proof
  • Lemma C.2: $\theta$-decomposition via prefix sums
  • proof
  • Lemma C.3: Order preservation under prefix dominance
  • proof