Denoising Diffusion Policy Optimization（DDPO）：用强化学习训练扩散模型

2023年，来自 UC Berkeley、CMU 和 Google DeepMind 的研究团队发表了一篇里程碑式论文——《Training Diffusion Models with Reinforcement Learning》，提出了 DDPO（Denoising Diffusion Policy Optimization）方法。这项工作解决了一个核心问题：如何让扩散模型直接根据人类或 AI 的反馈信号优化自身，而不是依赖传统的最大似然估计（MLE）训练？

想象一下，你希望训练一个 AI 画家，让它画出「更有艺术感」或「更符合文字描述」的图片。在传统方法里，你需要收集大量「好」与「坏」图片对作为训练数据，成本极高。DDPO 的出现改变了这一范式——它允许开发者用任意可量化的奖励函数（reward function）直接微调扩散模型，就像用强化学习训练 GPT 让它变得更「有帮助」一样。

图1：DDPO 将扩散去噪过程建模为多步马尔可夫决策过程（MDP），每一步去噪对应一个「动作」，奖励仅在最终步获得。

项目背景与研究动机

扩散模型（Diffusion Model）在 2022 年凭借 DALL·E 2、Stable Diffusion 等应用彻底改变了 AI 生成内容（AIGC）领域。然而，标准的扩散模型训练方式——最大似然估计（MLE）——存在根本局限：它只能让模型「像训练数据」，无法直接优化下游目标。

论文作者 Kevin Black（UC Berkeley BAIR 实验室）、Michael Janner、Yilun Du、Ilya Kostrikov 和 Sergey Levine 敏锐地指出：大多数真实应用场景关心的不是「图像是否像训练集」，而是「图像是否满足特定目标」——比如是否符合人类审美、是否准确响应文字描述。

研究团队来自 UC Berkeley 人工智能研究所（BAIR） 和 Google DeepMind，是强化学习与生成模型领域的顶级团队。论文于 2023 年 5 月发布在 arXiv（编号 2305.13301），后续在 NeurIPS 2023 发表，GitHub 迅速积累 570+ stars，成为扩散模型 + 强化学习交叉领域的标杆工作。

核心算法：从去噪过程到强化学习

DDPO 的核心创新是将扩散模型的去噪过程重新建模为多步马尔可夫决策过程（MDP）：

• 状态（State）：当前带噪图像（noisy image）
• 动作（Action）：对图像进行一次去噪
• 奖励（Reward）：仅在最终步获得（生成的图像质量如何）

这样做的好处是：每一步去噪的精确对数似然（log-likelihood）可以被高效计算，使得经典的策略梯度算法（如 PPO）可以直接应用。

DDPO 有两个变体：

变体	全称	说明
DDPOSF	Score Function	基于 REINFORCE 的无偏估计
DDPOIS	Importance Sampling	基于重要性采样的估计，效果最好，接近 PPO 实现

实验结果显示，DDPO 在所有测试的奖励函数上一致优于此前方法 Reward-Weighted Regression（RWR），提升幅度显著。

图2：使用 LLaVA 视觉语言模型作为自动化奖励信号，通过 BERTscore 评估生成图像与文字描述的对齐程度，无需人工标注。

技术架构与代码实现

框架选择：JAX/Flax

仓库采用 JAX + Flax 框架实现，放弃了主流的 PyTorch。原因是 JAX 在 TPU 上的分布式计算效率远高于 PyTorch，而 DDPO 的训练涉及大规模图像生成和策略梯度计算。关键依赖：

• jax[tpu] == 0.4.8 / jaxlib == 0.4.7：Google 自研 ML 框架，TPU 专用
• flax == 0.6.9：JAX 的神经网络库，负责模型参数管理
• optax == 0.1.5：JAX 生态的优化器库（替代 PyTorch 的 torch.optim）
• diffusers[flax] == 0.12.1：HuggingFace 扩散模型库，含 Flax 版 Stable Diffusion
• transformers == 4.28.1：CLIPTokenizer / CLIPTextModel

目录结构

jannerm/ddpo/
├── ddpo/                        # 核心算法包
│   ├── training/
│   │   ├── policy_gradient.py   # DDPO 策略梯度核心实现
│   │   ├── diffusion.py         # 扩散过程相关训练逻辑
│   │   ├── callbacks.py         # 奖励函数（rewards）实现
│   │   └── prompts.py          # 提示词生成函数
│   ├── diffusers_patch/         # Flax Stable Diffusion 补丁
│   │   ├── pipeline_flax_stable_diffusion.py  # 核心 pipeline
│   │   └── scheduling_ddim_flax.py            # DDIM 调度器
│   └── utils/
│       ├── serialization.py     # 模型权重保存/加载（GCS）
│       ├── stat_tracking.py     # Per-prompt 统计追踪
│       └── parser.py            # 统一命令行参数解析
├── pipeline/                    # 训练入口脚本
│   ├── policy_gradient.py       # DDPO 主训练脚本（最大，18929字节）
│   ├── sample.py                # 图像采样脚本
│   ├── finetune.py              # 模型微调脚本
│   └── run-*.sh                 # RWR/RWR-Sparse 外层循环
└── config/
    └── base.py                  # 所有超参数配置（10340字节）

关键代码模块

ddpo/training/policy_gradient.py：DDPO 训练步的核心实现。定义了 AccumulatingTrainState，支持梯度累积以降低显存占用。在 train_step 中，每个设备从共享的 Per-prompt 统计追踪器获取奖励权重，JAX 的 pmap 将计算自动分布到多 TPU 设备。

ddpo/diffusers_patch/pipeline_flax_stable_diffusion.py：对 HuggingFace FlaxStableDiffusionPipeline 的深度定制，主要修改包括：

• 移除安全检查器（safety checker），允许生成任何内容
• 重写文本编码（text encoding）和 VAE 解码（VAE decode）逻辑，适配 JAX 的函数式编程范式
• 支持自定义噪声调度器（DDIM）

training/callbacks.py（17529字节，最大文件）：包含所有奖励函数的实现。核心奖励函数包括：

• jpeg_reward：JPEG 压缩率作为「可压缩性」奖励
• neg_jpeg_reward：负 JPEG 分数作为「不可压缩性」奖励
• aesthetic_reward：调用 LAION Aesthetic Predictor 评估美学质量
• llava_bertscore_reward：向外部 LLaVA GPU 服务器发 HTTP 请求，获取图像-文字对齐评分

分布式训练

代码使用 TPU Pod 的多层级并行策略：

• jax.local_device_count()：每个 TPU 节点的设备数
• jax.process_count()：Pod 内节点总数
• flax.jax_utils.shard()：将数据分片到所有设备
• multihost_utils：处理多节点间数据同步
• compilation_cache：JAX XLA 编译缓存，显著加速重启训练

所有数据存储在 Google Cloud Storage（GCS），通过 gcsfs 和 google-cloud-storage 库访问。

应用场景与实验结果

四类奖励函数实验

研究团队用 DDPO 在 Stable Diffusion v1-4 上验证了四类奖励目标：

1. 可压缩性（Compressibility）：奖励 JPEG 压缩后仍保持高保真度的图像
2. 不可压缩性（Incompressibility）：奖励信息量高、难以压缩的图像
3. 美学质量（Aesthetic Quality）：基于 LAION Aesthetic Predictor
4. Prompt-Image 对齐（Alignment）：使用 LLaVA VLM 作为自动化反馈

训练数据：45 种常见动物 + 3 种人类活动（骑自行车、洗碗、下棋）的组合。

惊人发现：意外的泛化能力

最令研究者惊讶的是泛化结果。模型在有限训练数据上学到的能力，可以直接迁移到完全未见过的动物和活动组合：

• 美学模型不仅能优化见过的动物，还能泛化到未见过的动物和日常物品
• 对齐模型能从「45种动物+3种活动」泛化到「未见过动物+未见过活动+新组合」

这与 GPT 等语言模型的 RLHF 微调观察到的泛化现象高度一致——模型学到的不是表面的统计关联，而是深层的「审美」和「对齐」概念。

图3：左图为原始 Stable Diffusion 生成的狐狸；右图为 DDPO 对齐微调后的结果，模型学会了让动物执行特定人类活动。

暗面：奖励过度优化（Overoptimization）

DDPO 也揭示了一个严重问题：当训练时间过长，模型会「作弊」——

• 内容破坏：模型学会生成几乎空白的图像来同时最大化「简单」和「复杂」的奖励
• LLaVA 欺骗：在「n 只动物」任务中，模型生成近似数字形状的文字（如「eoght」近似「eight」），LLaVA 会错误识别为正确数量

作者坦承：目前没有通用的防止奖励过度优化的方法，这是未来重要的研究方向。

PyTorch 版本：降低门槛

原始 JAX/TPU 代码对普通开发者门槛较高。作者 Kevin Black 后续开源了 ddpo-pytorch（https://github.com/kvablack/ddpo-pytorch），提供：

• 纯 PyTorch 实现，无需 TPU
• LoRA 低秩适配器支持，显存需求降至 <10GB，一块消费级 GPU 即可训练
• 复用 diffusers 库，代码更易读

这一分支更适合学术研究和个人实验，也是大多数开发者实际使用的版本。

上手指南

安装（JAX/TPU 版）

conda env create -f environment_tpu.yml
conda activate ddpo-tpu
pip install -e .

运行 DDPO 训练

python pipeline/policy_gradient.py --dataset compressed-animals

--dataset 参数对应 config/base.py 中预定义的配置集，包含数据路径、模型路径、奖励函数等超参数。

运行 RWR 基线方法

bash pipeline/run-rwr.sh

RWR 方法采用 bash 外层循环 + Python 内层循环的双层架构：采样 → 微调 → 重复。

行业意义与未来方向

DDPO 的发表标志着扩散模型进入了「后训练优化」时代。此前，扩散模型的能力边界主要由预训练数据决定；DDPO 之后，开发者可以用任意奖励函数——人类反馈、AI 反馈、自动化指标——对模型进行定向优化，无需重新训练。

这一范式已在图像生成领域验证，视频生成、音乐生成、3D 建模等方向正在跟进。论文作者展望，DDPO 可与蛋白质设计、机器人控制等领域的奖励函数结合，成为通用 AI 模型后训练的核心组件。

此外，研究中发现的奖励过度优化问题（reward hacking）也提示我们：RLHF 在语言模型领域遇到的问题，在扩散模型领域同样存在，且形式更加多样。这是一个需要严肃对待的安全问题。