Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps

GenAI Diffusion ML POA Generative Models DM AI4S Other AI4S
近期基于流的生成建模研究取得了重要进展,涌现出一批仅需少量函数求值即可输出高质量样本的模型。然而,目前在模型似然估计方面尚缺乏类似突破。具体而言,现有主流方法要么依赖于结构受限的架构(以支持似然的精确计算),要么采用随机近似方法(例如Hutchinson迹估计器),而后者会引入显著的方差。本文提出一种可扩展的流映射似然蒸馏方法——SCALLOP(SCAlable LikeLihood distillation of flOw maPs)。SCALLOP建立在近期提出的F2D2模型基础之上:F2D2是一种似然流映射模型,可在极少量函数求值下同步生成样本及其对应密度。尽管F2D2在训练过程中采用Hutchinson估计器,我们则设计了一种替代性的、更具可扩展性的似然蒸馏目标函数——该目标完全无需Hutchinson估计,且支持向量化实现。实验结果表明,SCALLOP作为分子科学领域的玻尔兹曼生成器效果卓著;我们还在图像数据集上进一步验证了其有效性。相较于F2D2,SCALLOP显著降低了训练过程中的方差与耗时,同时持续提升模型性能;其综合表现与当前最优方法相当,且推理速度相较最快基线最高可提升达10倍。
Recent progress in flow-based generative modeling has led to models that output high-quality samples while using only a small number of function evaluations. However, at present, there is a lack of similar advances in estimating the model likelihood. In particular, most existing methods either rely on restrictive architectures that enable exact calculations, or use stochastic approximations such as Hutchinson's trace estimator that introduce substantial variance. In this work, we introduce SCAlable LikeLihood distillation of flOw maPs (SCALLOP). SCALLOP builds on the recently proposed F2D2, a likelihood flow map model that can generate samples and their densities in a small number of function evaluations. While F2D2 uses Hutchinson's estimator during training, we introduce an alternative and more scalable likelihood distillation objective that is Hutchinson-free and admits a vectorized formulation. Empirically, we demonstrate the effectiveness of SCALLOP as a Boltzmann generator in molecular science, and further validate its benefit on image datasets. SCALLOP significantly reduces both training variance and training time while consistently improving performance compared to F2D2, and is competitive with the state-of-the-art while achieving up to 10x inference speedup over the fastest baseline.
许愿