摘要

本文介绍了一种基于最大熵逆强化学习（IRL）的方法，用于改进扩散生成模型的样本质量，特别是在生成时间步数较少的情况下。该方法通过训练扩散模型使用从训练数据估计的对数概率密度，类似于IRL基于专家演示学习奖励函数来训练策略。由于采用了基于能量模型（EBM）来表示对数密度，该方法简化为扩散模型和EBM的联合训练。提出的IRL框架名为Diffusion by Maximum Entropy IRL（DxMI），是一个达到平衡时两个模型都收敛到数据分布的最小最大问题。此外，还提出了Diffusion by Dynamic Programming（DxDP），一种用于扩散模型的新型强化学习算法，作为DxMI的子程序，通过将原始问题转化为最优控制问题，使扩散模型更新在DxMI中更高效。实验表明，使用DxMI微调的扩散模型可以在仅4到10步内生成高质量样本，并能训练EBM而无需MCMC，稳定EBM训练动态并提高异常检测性能。

原理

DxMI的核心在于通过最大熵IRL框架联合优化扩散模型和能量模型（EBM）。EBM提供估计的对数密度作为扩散模型的奖励信号，而扩散模型则被训练以最大化来自EBM的奖励同时最大化生成样本的熵。这种熵最大化促进了扩散模型的探索并确保了EBM的收敛。DxDP通过动态规划方法，将扩散模型的更新问题转化为最优控制问题，利用价值函数替代时间反向传播，从而解决了直接更新扩散模型时的边际熵估计难题和内存限制问题。

流程

DxMI的工作流程包括交替更新EBM和扩散模型以达到Nash均衡。EBM更新直接进行，而扩散模型更新则通过DxDP算法进行。DxDP首先通过优化原始目标的上界来缓解边际熵估计问题，然后通过解释目标为最优控制问题并应用动态规划使用价值函数来消除时间反向传播的需要。具体步骤包括：1) 输入数据集和模型参数；2) 初始化AVR；3) 对数据集中的每个样本，采样x0:T；4) 更新EBM参数；5) 通过动态规划更新价值函数；6) 更新扩散模型参数；7) 更新AVR。

应用

DxMI的应用前景广泛，特别是在需要快速且高质量样本生成的领域，如图像生成和异常检测。由于其能够在较少步骤内生成高质量样本，DxMI在实时应用和计算资源受限的环境中尤为有用。此外，DxMI在训练EBM时无需MCMC，这为EBM在各种高维数据处理任务中的应用提供了新的可能性。