On scalable and efficient training of diffusion samplers

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Kim, Minkyu, Seong, Kiyoung, Woo, Dongyeop, Ahn, Sungsoo, Kim, Minsu
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866914139945828352
author Kim, Minkyu
Seong, Kiyoung
Woo, Dongyeop
Ahn, Sungsoo
Kim, Minsu
author_facet Kim, Minkyu
Seong, Kiyoung
Woo, Dongyeop
Ahn, Sungsoo
Kim, Minsu
contents We address the challenge of training diffusion models to sample from unnormalized energy distributions in the absence of data, the so-called diffusion samplers. Although these approaches have shown promise, they struggle to scale in more demanding scenarios where energy evaluations are expensive and the sampling space is high-dimensional. To address this limitation, we propose a scalable and sample-efficient framework that properly harmonizes the powerful classical sampling method and the diffusion sampler. Specifically, we utilize Monte Carlo Markov chain (MCMC) samplers with a novelty-based auxiliary energy as a Searcher to collect off-policy samples, using an auxiliary energy function to compensate for exploring modes the diffusion sampler rarely visits. These off-policy samples are then combined with on-policy data to train the diffusion sampler, thereby expanding its coverage of the energy landscape. Furthermore, we identify primacy bias, i.e., the preference of samplers for early experience during training, as the main cause of mode collapse during training, and introduce a periodic re-initialization trick to resolve this issue. Our method significantly improves sample efficiency on standard benchmarks for diffusion samplers and also excels at higher-dimensional problems and real-world molecular conformer generation.
format Preprint
id arxiv_https___arxiv_org_abs_2505_19552
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle On scalable and efficient training of diffusion samplers
Kim, Minkyu
Seong, Kiyoung
Woo, Dongyeop
Ahn, Sungsoo
Kim, Minsu
Machine Learning
We address the challenge of training diffusion models to sample from unnormalized energy distributions in the absence of data, the so-called diffusion samplers. Although these approaches have shown promise, they struggle to scale in more demanding scenarios where energy evaluations are expensive and the sampling space is high-dimensional. To address this limitation, we propose a scalable and sample-efficient framework that properly harmonizes the powerful classical sampling method and the diffusion sampler. Specifically, we utilize Monte Carlo Markov chain (MCMC) samplers with a novelty-based auxiliary energy as a Searcher to collect off-policy samples, using an auxiliary energy function to compensate for exploring modes the diffusion sampler rarely visits. These off-policy samples are then combined with on-policy data to train the diffusion sampler, thereby expanding its coverage of the energy landscape. Furthermore, we identify primacy bias, i.e., the preference of samplers for early experience during training, as the main cause of mode collapse during training, and introduce a periodic re-initialization trick to resolve this issue. Our method significantly improves sample efficiency on standard benchmarks for diffusion samplers and also excels at higher-dimensional problems and real-world molecular conformer generation.
title On scalable and efficient training of diffusion samplers
topic Machine Learning
url https://arxiv.org/abs/2505.19552