Diffusion Alignment as Variational Expectation-Maximization

Each language version is independently generated for its own context, not a direct translation.

这篇论文介绍了一种名为 DAV（Diffusion Alignment as Variational Expectation-Maximization，变分期望最大化扩散对齐）的新方法。

为了让你轻松理解，我们可以把训练一个能画好图或设计好 DNA 的 AI 模型，想象成教一个天才画家去画“完美”的画作。

1. 背景：天才画家遇到了瓶颈

现在的扩散模型（Diffusion Models，比如 Midjourney 或 Stable Diffusion）就像是一个天赋异禀的画家。他看过无数名画，能画出非常逼真、风格多样的作品。

但是，如果你让他画“一张既符合人类审美、又能被压缩成极小文件、且看起来像某种特定生物”的画，他可能会犯两个错误：

过度优化（Reward Over-optimization）： 为了讨好你，他开始疯狂画某种特定的“完美”图案，结果画出来的东西千篇一律，失去了多样性（比如所有画里的猫都长得一样）。
模式崩溃（Mode Collapse）： 他为了追求高分，只敢画最安全、最平庸的东西，不敢尝试新风格，导致作品变得枯燥乏味。

以前的方法（比如强化学习 RL 或直接反向传播）就像是严厉的考官，直接告诉画家：“你这里画错了，扣分！”或者“你那里画得好，加分！”。但这往往导致画家变得焦虑，要么只敢画一种东西，要么为了拿高分而画出了怪异的“作弊”作品。

2. 核心创意：DAV 的“两步走”策略

DAV 提出了一种更聪明的教学方法，它把训练过程分成了两个交替进行的步骤，就像**“探索”和“总结”**的循环。

第一步：E-Step（探索者）—— 在脑海中“试错”

想象画家在正式动笔前，先拿出一堆草稿纸，进行**“试画”**。

怎么做？ 画家利用一种特殊的“搜索”技巧（Test-time Search），在生成图片的过程中，像侦探一样尝试各种可能性。他会问自己：“如果我稍微改一下这个笔触，奖励（分数）会不会更高？”
关键点： 这一步不需要修改画家的核心能力，而是利用额外的计算力，在生成过程中主动寻找那些既高分又多样的“完美草稿”。这就像是在大海里撒网，专门捞那些最漂亮的鱼。
比喻： 这就像是一个**“试吃员”**，在正式上菜前，先尝遍所有可能的口味组合，找出最好吃的那几道。

第二步：M-Step（总结者）—— 把经验“内化”

一旦“试吃员”找到了那些完美的草稿，真正的**“教学”**就开始了。

怎么做？ 我们把刚才找到的那些“完美草稿”收集起来，让画家模仿这些草稿重新学习。
关键点： 这里有一个精妙的设计。以前的方法只让画家模仿“最高分”的那张图（导致他只会画一种东西）。但 DAV 告诉画家：“你要模仿所有找到的好草稿，包括那些风格各异的。”
比喻： 这就像老师给学生看了一堆**“优秀作业集”**，并告诉学生：“不要只抄第一名，要把所有优秀作业里的优点都学过来，保持你原本丰富的风格，同时提升水平。”

3. 为什么 DAV 这么厉害？

DAV 的核心在于它把“搜索”和“学习”分开了，并且用一种数学上的**“前向 KL 散度”**（Forward KL）来指导学习。

传统方法（后向 KL）： 像是强迫画家只画“最像第一名”的画。结果：画家变成了复印机，只会画一种东西。
DAV 方法（前向 KL）： 像是告诉画家：“你要覆盖所有好画的风格。”结果：画家既学会了高分技巧，又保留了原本丰富多彩的风格。

简单总结就是：

先搜后学： 先花点力气在生成时多尝试（E 步），找到好样本。
批量吸收： 再把这些好样本教给模型（M 步），让它学会“既要分高，又要花样多”。

4. 实际效果：不仅会画画，还会设计 DNA

论文在两个领域验证了这种方法：

文生图（连续数据）： 让 Stable Diffusion 画出的图，既符合人类审美（分数高），又不会千篇一律（多样性好）。以前的方法画久了，图里的猫都长得一样；DAV 画出来的猫，每只都不同，但都很美。
DNA 序列设计（离散数据）： 让 AI 设计 DNA 片段。以前的方法设计的 DNA 虽然活性高，但可能全是同一种序列，缺乏多样性。DAV 设计出的 DNA 既活性高，又保留了自然界中 DNA 的多样性和自然感。

5. 总结

这篇论文就像给 AI 画家发明了一套**“先试错、后总结、保多样”**的独门心法。它解决了 AI 在追求高分时容易“走火入魔”（变得单一、怪异）的问题，让 AI 既能拿高分，又能保持创造力。

一句话概括： DAV 让 AI 在“追求完美”和“保持多样”之间找到了完美的平衡点，不再为了拿高分而变成只会画一种东西的“死板机器”。

Each language version is independently generated for its own context, not a direct translation.

1. 研究背景与问题 (Problem)

背景：
扩散模型（Diffusion Models）在图像生成、机器人控制和计算生物学等领域表现出色。然而，许多实际应用不仅需要生成高似然的样本，还需要生成针对特定下游目标（如图像美学质量、DNA 增强子活性）进行优化的样本。

现有挑战：
现有的扩散模型对齐（Alignment）方法主要分为两类，但都存在显著缺陷：

基于强化学习（RL）的微调： 通常使用反向 KL 散度（Reverse-KL）作为目标。这容易导致模式寻求（Mode-seeking）行为，引发奖励过优化（Reward Over-optimization）和模式坍塌（Mode Collapse），即模型生成单一的高分样本，丧失了多样性和自然性。
直接反向传播（Direct Backpropagation）： 虽然样本效率高，但依赖于可微分的奖励函数，且梯度信号往往尖锐且脆弱，同样容易导致严重的奖励过优化。
测试时推理（Test-time Inference）： 如引导（Guidance）或搜索方法，虽然无需微调模型，但计算开销巨大，且难以在推理阶段维持长期的多样性，或者存在欠优化问题。

核心问题：
如何设计一个微调框架，既能有效最大化奖励，又能保留预训练扩散模型的多样性和自然性，同时适用于连续（如图像）和离散（如 DNA 序列）数据模态，且不需要奖励函数可微分？

2. 方法论：DAV (Methodology)

作者提出了 DAV (Diffusion Alignment as Variational Expectation-Maximization)，将扩散对齐问题形式化为一个变分期望最大化（Variational EM）的迭代过程。该方法交替进行两个互补的阶段：**E 步（探索）**和 M 步（摊销/蒸馏）。

2.1 核心框架

DAV 的目标是最大化优化变量 $O$ （表示样本是否最优）的边缘似然 $\log p_\theta(O=1)$ 。由于直接优化不可行，作者引入变分分布 $\eta(\tau)$ 来近似后验分布，并推导出了证据下界（ELBO）。

E 步 (Exploration / 后验推断)：
- 目标： 从变分后验分布 $\eta^*(\tau)$ 中采样，寻找高奖励、多样化的轨迹。
- 机制： 利用测试时搜索（Test-time Search）。
  - 首先定义最优策略 $\eta^*$ ，其形式为奖励倾斜分布（Reward-tilted distribution）： $\eta^*(x_{t-1}|x_t) \propto p_\theta(x_{t-1}|x_t) \exp(Q^*_{soft}/\alpha)$ 。
  - 由于直接采样困难，采用两阶段局部搜索：
    1. 提案分布构建： 如果奖励函数可微，利用梯度引导（Gradient Guidance）构建提案分布；对于离散数据或不可微奖励，则直接采样候选粒子。
    2. 重要性采样修正： 对候选粒子进行重要性采样（Importance Sampling），修正提案分布与真实后验分布的偏差，从而获得高质量且多样化的样本。
- 特点： 这是一个模块化的组件，可以集成任何先进的测试时搜索算法。
M 步 (Amortization / 模型更新)：
- 目标： 将 E 步中发现的轨迹知识蒸馏回扩散模型参数 $\theta$ 。
- 机制： 通过最小化**前向 KL 散度（Forward KL Divergence）**来更新模型。
  - 更新目标： $\theta_{k+1} = \arg\min_\theta D_{KL}(\eta^*_k(\tau) || p_\theta(\tau))$ 。
  - 这等价于最大化 E 步中采样轨迹的对数似然。
- 关键创新： 与 RL 方法常用的反向 KL（导致模式坍塌）不同，DAV 使用前向 KL。前向 KL 是一种**模式覆盖（Mode-covering）**目标，鼓励模型覆盖 E 步中发现的所有多样化模式，从而防止多样性丧失。
- DAV-KL 变体： 为了防止模型偏离预训练分布太远，引入 KL 正则化项，平衡对齐专家策略与保留预训练知识。

2.2 理论支撑

折扣因子引入： 针对扩散过程早期步骤对最终结果影响较小的特点，引入了折扣因子 $\gamma$ ，优化了信用分配（Credit Assignment）。
软 Q 函数近似： 利用 Tweedie 公式近似软 Q 函数，使得在离散和连续空间中都能有效计算引导信号。

3. 主要贡献 (Key Contributions)

新框架提出： 首次将扩散模型对齐问题形式化为变分 EM 过程，通过交替的“测试时搜索（E 步）”和“前向 KL 蒸馏（M 步）”解决对齐问题。
解决模式坍塌与过优化： 通过 E 步的多样化搜索和 M 步的前向 KL 优化，成功在最大化奖励的同时保留了样本的多样性和自然性，有效缓解了奖励过优化问题。
通用性与灵活性：
- 适用于连续（图像生成）和离散（DNA 序列设计）两种数据模态。
- 不依赖奖励函数的可微性，能够处理黑盒奖励（如图像压缩性、生物活性预测器）。
模块化设计： E 步的搜索组件是模块化的，可以无缝集成未来的测试时搜索技术进步。

4. 实验结果 (Results)

作者在两个主要领域进行了验证：

4.1 连续扩散：文本到图像合成 (Text-to-Image)

基准： Stable Diffusion v1.5，优化美学分数（Aesthetic Score）。
对比方法： DDPO (RL), TDPO, DRaFT (直接反向传播), DAS (测试时搜索)。
结果：
- 奖励与多样性平衡： DAV 在美学分数（8.04）上显著优于 DDPO (6.83) 和 DRaFT (7.22)，同时保持了极高的 ImageReward (0.95) 和多样性 (LPIPS-A 0.53)。相比之下，DDPO 和 DRaFT 在优化过程中出现了严重的多样性下降和 ImageReward 崩溃。
- 定性分析： DAV 生成的图像背景丰富多样，而基线方法（如 DDPO）出现了重复的背景纹理（模式坍塌）。
- 消融实验： 证明了 E 步中测试时搜索的重要性，跳过搜索仅使用重加权（Reweight）会导致性能大幅下降。

4.2 离散扩散：DNA 序列设计 (DNA Sequence Design)

任务： 设计具有高增强子活性（Enhancer Activity）的 DNA 序列。
基准： DRAKES, VIDD, DDPO。
指标： 预测活性、ATAC-Acc（生物有效性）、3-mer 相关性（自然性）、Levenshtein 距离（多样性）。
结果：
- DAV 在保持高多样性（87.91）和自然性的同时，达到了比基线更高的预测活性（7.71 vs 7.51）。
- 基线方法（如 DDPO）虽然活性高，但 ATAC-Acc 和多样性大幅下降，表明发生了过优化。
- DAV Posterior（仅使用测试时搜索，不微调模型）甚至达到了最高的活性（9.24）和有效性（0.920），展示了该框架在推理阶段的强大潜力。

4.3 非可微奖励

在图像压缩性/不可压缩性优化任务中，DAV 成功优化了黑盒奖励，且仅需 6 个 epoch 的训练，性能优于训练 100 个 epoch 的 DDPO。

5. 意义与影响 (Significance)

理论突破： 将扩散对齐从单纯的 RL 或梯度下降范式，转变为基于变分推断的 EM 范式，从理论上解释了为何前向 KL 能更好地维持多样性。
实际应用价值： 为生物设计（如蛋白质、DNA 设计）和创意生成（如图像、视频）提供了一种稳健的对齐工具，特别是在奖励函数不可微或容易过优化的场景下。
效率与效果的权衡： 虽然 E 步引入了额外的计算开销（测试时搜索），但通过 M 步的蒸馏，模型在推理阶段无需额外搜索即可保持高性能，且总训练成本在许多情况下低于需要大量 epoch 才能收敛的 RL 方法。
未来方向： 论文指出 E 步的计算开销是主要限制，未来可通过集成更高效的搜索算法或蒸馏技术（减少去噪步数）来进一步优化。

总结： DAV 通过巧妙结合测试时搜索的探索能力和前向 KL 蒸馏的覆盖能力，成功解决了扩散模型对齐中的“奖励过优化”与“多样性丧失”这一长期存在的矛盾，为多模态扩散模型的下游任务优化提供了新的通用范式。