arXiv:2608.08770stat.MLcs.LG2026-08

提出分布级控制扩散模型的理论框架,实现更精准的生成调控。

A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models

论文配图:A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models
图 1 · 摘自论文原文
  • 基于均场理论构建分布级奖励的粒子加权机制
  • 在低维任务中准确实现目标分布采样,高维蛋白构象任务也表现稳定
  • 兼容点级奖励,适合需要群体调控的应用场景

扩散模型作为可调控采样器,在推理阶段可根据奖励函数引导生成。尽管现有方法多基于单样本奖励,但许多应用需基于分布层面的奖励进行调控,如与群体信息对齐或提升多样性。直接将奖励梯度融入动态虽有效,却缺乏理论保证。此前针对点级奖励的研究已建立基于粒子重加权的合理框架,但分布级奖励尚无相应理论支撑。本文在均场框架下,将推理阶段的分布级控制建模为靶向倾斜测度问题,并推导出一种加权交互粒子方案,实现原理严谨的分布调控。该框架可还原点级奖励控制为特例,同时为现有批量级调控方法提供理论基础。实验表明,在可解析的低维设置中,该方法能正确收敛至目标分布;在更高维度的蛋白质构象任务中亦表现出良好行为。

原文摘要 · Abstract (English)

Diffusion models are increasingly used as controllable samplers, whose generations can be steered at inference time according to a chosen reward function. While such rewards are typically defined on individual samples, for many applications it is desirable to steer according to distribution-level rewards, for example to calibrate with population-level information or to encourage diversity. In both cases, simply incorporating the reward gradient into the dynamics, while often effective, comes with few theoretical guarantees on the sampled distribution. For pointwise rewards, recent work has therefore sought to develop a principled framework for targeting a prescribed tilted distribution using particle reweighting. However, an analogous theoretically-grounded approach for distributional rewards is currently lacking. In this work, we formulate inference-time distributional control as targeting a tilted measure under a mean-field framework, and derive a weighted interacting particle scheme to target it in a principled manner. Our framework recovers pointwise-reward steering as a special case, while providing a theoretical foundation for existing batch-level steering methods. Empirically, we verify that the procedure correctly targets the prescribed distribution in tractable low-dimensional settings, and investigate its behaviour in higher-dimensional protein conformation tasks.

扩散模型分布控制均场理论生成建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。