← 返回 2026-05-08

MARBLE：面向扩散模型强化学习的多维度奖励均衡框架 MARBLE: Multi-Aspect Reward Balance for Diffusion RL

Canyu Zhao, Hao Chen, Yunze Tong, Yu Qiao, Jiacheng Li, Chunhua Shen 📅 2026-05-07 👍 40 2026-07-13 08:36
图像生成 多奖励优化 强化学习 扩散模型 梯度调和
arXiv HuggingFace PDF

梯度空间均衡多奖励，单模型同时提升扩散RL所有质量维度

前置知识

扩散模型与强化学习微调

扩散模型通过迭代去噪生成图像，RL微调（如FlowGRPO、DiffusionNFT）以奖励模型评估样本质量并更新参数。DiffusionNFT使用无噪声训练目标 $\ell(\theta;x,t)=r\cdot L^++(1-r)\cdot L^-$，其中$r$由优势$A(x)$线性映射得来。

MARBLE建立在DiffusionNFT的affine loss结构之上，需要理解扩散RL的基本流程和NFT损失中优势与梯度之间的线性关系，才能抓住论文的核心推导。

多任务学习中的梯度均衡

经典方法MGDA、CAGrad等在梯度空间解决任务冲突：将多个任务的梯度$\{g_k\}$归一化后在单纯形$\Delta_K$上求最小范数点。MARBLE正是把这一思想迁移到扩散RL的奖励维度上。

读者必须理解'梯度空间调和'与'损失空间加权'的本质区别，才能明白为什么加权求和$R(x)=\sum_k w_k R_k(x)$会失败、而MARBLE能成功。

专家样本现象（specialist sample phenomenon）

指单条样本通常只在部分奖励维度上信息丰富，例如OCR奖励高的样本在美学维度可能平庸。在加权求和下，该样本对无关维度的优势被稀释，整体优势不再反映其真正擅长的方向。

这是论文提出的核心诊断概念，解释了为什么传统标量化聚合会导致梯度方向与最难奖励负相关，是阅读3.2节和Figure 2的钥匙。

研究动机

扩散模型的RL微调需要同时对齐多种质量维度（美学、文本-图像对齐、OCR准确性、组合正确性等），但现有方法各有明显缺陷。第一种是'一个奖励一个专家模型'：FlowGRPO对PickScore单独训练得22.51分，但迁移到其他奖励时显著下降，需要维护多个模型且无法跨奖励泛化。第二种是加权求和$R(x)=\sum_k w_k R_k(x)$：作者在SD3.5 Medium上测量发现，加权求和的更新方向在80%的mini-batch中与'最难奖励'的单奖励梯度呈负相关，即优化过程主动在与最难目标相反的方向上走。第三种是顺序多阶段训练（DiffusionNFT†）：需要手工设计阶段调度，例如'800步奖励1、300步奖励2、200步奖励1、200步奖励2、100步奖励3'，超参敏感且新增奖励时必须回访旧奖励以缓解遗忘，调度规模随奖励数量线性膨胀。

本文的目标是本文提出MARBLE（Multi-Aspect Reward BaLancE），目标是给扩散RL一个无需手工权重调节、无需多模型维护、无需分阶段调度的统一方案。具体包含四个子目标：(1)让单一扩散模型在一个训练run中同时优化多个奖励信号的所有维度；(2)显式避免'专家样本稀释'导致的梯度方向偏移；(3)通过梯度空间调和确保对最难奖励的余弦相似度始终为正；(4)在摊销设计下把每步训练开销控制在接近单奖励基线水平，让多目标对齐从理论可行走向工程实用。

与已有工作不同的是，MARBLE的独特切入角度是把多奖励平衡问题从'损失/奖励空间'转移到'梯度空间'。已有的多任务学习文献（MGDA、CAGrad、PCGrad）证明了在梯度空间解决任务干扰的有效性，但扩散RL领域此前从未系统地做过这一迁移。论文同时利用DiffusionNFT损失关于优势$A$的affine结构（Proposition 1），将多奖励反向传播从$K+1$次降为接近1次，使梯度均衡在算力上变得可行，这是其他多任务方法没有利用的扩散专属结构。

核心方法

MARBLE的整体思路分三层：第一层'奖励维度解耦'——不再把$K$个奖励聚合为标量，而是为每个奖励$R_k$维护独立的优势估计器$A_k(x)$，并通过反向传播得到$K$个独立策略梯度$g_k$。第二层'梯度空间调和'——把所有$g_k$归一化到单位球$\hat{g}_k=g_k/\|g_k\|$，然后在单纯形$\Delta_K=\{\alpha\geq 0:\sum_k\alpha_k=1\}$上求解凸二次规划$\min_\alpha\|\sum_k\alpha_k\hat{g}_k\|^2$，得到均衡更新方向$d^*=\sum_k\alpha_k^*\hat{g}_k$。第三层'算力分摊与稳定化'——利用NFT损失的affine结构，每$N$步才执行完整$K$次反向传播刷新$\alpha^*$，中间$N-1$步用缓存系数做单次反向传播，再对$\alpha^*$做EMA平滑（$\rho=0.7$）。直觉上就像团队决策：每个奖励'投票'、QP当'调解员'、EMA防止朝令夕改。

MARBLE与已有方法的本质区别有两点。第一，加权求和在'奖励空间'做聚合，混入了样本-奖励错配的噪声；MARBLE在'梯度空间'做调和，每个样本只在它真正擅长的维度上贡献监督信号。第二，MARBLE首次利用DiffusionNFT的affine结构$\nabla_\theta\ell(\theta;\bar{A})=\sum_k\alpha_k\nabla_\theta\ell_k(\theta;A_k)$（当clamp未激活时精确成立）实现摊销，把多奖励训练的每步开销从$K+1$次反向传播降到接近1次，使梯度均衡方案在算力上真正可行。这两个创新共同保证了'既有效率又能多目标对齐'。

方法步骤详情

MARBLE流程分四步。第一步：共享模型采样prompt batch的图像，$K$个奖励模型独立打分。第二步：对每个奖励$R_k$做组内z-score归一化$A_k$，由NFT损失得插值系数$r_k=A_k/(2A_{max})+1/2$，独立反向传播得$K$个策略梯度$g_k$。第三步：归一化$g_k$到单位球，求解单纯形上的QP $\min_\alpha\|\sum_k\alpha_k\hat{g}_k\|^2$得$\alpha^*$，构造$d^*=\sum_k\alpha_k^*\hat{g}_k$，乘以平均梯度范数恢复步长，再加KL正则项更新参数。第四步（摊销与平滑）：每$N$步执行完整流程刷新$\alpha^*$，中间用$\bar{A}=\sum_k\bar{\alpha}_{t,k}A_k$做单次反向传播，并对$\alpha^*$做EMA平滑。实验中$A_{max}=5$、$N=10$、$\rho=0.7$。

技术新颖性

MARBLE的技术新颖性体现在三方面：诊断层面，首次形式化量化'专家样本现象'，并给出梯度级别的反相关证据（80% mini-batch的加权求和梯度与最难奖励负相关），把直觉问题变成可测量的指标；算法层面，把多任务学习的梯度调和（MGDA/CAGrad族）首次适配到扩散RL，针对NFT损失做per-reward advantage分解并独立求解QP；系统层面，摊销公式（Proposition 1）严格证明而非近似地把$K+1$次反向传播转化为接近1次，这是之前任何多任务方法都没有利用过的扩散专属结构。三层叠加使得MARBLE成为首个'同时满足多奖励联合训练、训练速度0.97×基线、单一模型、无需手工权重'的扩散RL方案。

Overview of MARBLE. Given a prompt batch, the shared model $\pi_\theta$ generates images that are scored by K reward models independently. Per-reward policy gradients $g_k=\nabla_\theta L_k$ are computed via separate backpropagation passes. The gradient harmonization finds a common descent direction $d$ that balances all reward objectives and the shared model is updated accordingly.
Figure 3: Overview of MARBLE. Given a prompt batch, the shared model $\pi_\theta$ generates images that are scored by K reward models independently. Per-reward policy gradients $g_k=\nabla_\theta L_k$ are computed via separate backpropagation passes. The gradient harmonization finds a common descent direction $d$ that balances all reward objectives and the shared model is updated accordingly.

实验结果

Table 1显示MARBLE在SD3.5 Medium+LoRA上联合优化PickScore/HPSv2/CLIPScore/OCR/GenEval五奖励，Composite列以+1.116位列第一。训练奖励GenEval=0.94、OCR=0.96、PickScore=22.83、CLIPScore=0.286、HPSv2.1=0.355均不弱于任何基线；Aesthetic=6.59、ImgRwd=1.53、UniRwd=3.52三项held-out全部最强，迁移性优异。对比FlowGRPO单专家跨奖励弱；对比DiffusionNFT顺序版MARBLE在OCR/Aesthetic反超且无需阶段超参；对比加权求和版GenEval 0.92 vs 0.94、OCR 0.91 vs 0.96、HPSv2.1 0.300 vs 0.355全维度领先。Table 2：摊销版速度0.97×、内存67G；非摊销仅0.56×。Table 3消融：去梯度归一化'FAIL'；固定$\alpha=0.2$使GenEval降至0.86；每步更新$\alpha$使PickScore跌至21.32。

Main results. Comparison of MARBLE with pre-trained diffusion models and RL fine-tuning methods. MARBLE jointly optimizes all in-domain rewards in a single run.
Table 1: Main results. Comparison of MARBLE with pre-trained diffusion models and RL fine-tuning methods. MARBLE jointly optimizes all in-domain rewards in a single run.
Training efficiency comparison on 8×H200. We report relative training speed and GPU memory, both normalized by the weighted-sum baseline.
Table 2: Training efficiency comparison on 8×H200. We report relative training speed and GPU memory, both normalized by the weighted-sum baseline.
Ablation study. Each row removes or replaces one component of MARBLE. All variants use the same 5-reward setup and training budget.
Table 3: Ablation study. Each row removes or replaces one component of MARBLE. All variants use the same 5-reward setup and training budget.
Visualizations of qualitative comparisons between MARBLE and Baselines.
Figure 4: Visualizations of qualitative comparisons between MARBLE and Baselines.
查看结构化数据
任务指标本文基线提升
组合正确性 GenEval 0.94 DiffusionNFT加权求和 0.92 / FlowGRPO专家 0.95 比加权求和基线+0.02，与最强单专家持平
文字渲染 OCR 0.96 DiffusionNFT加权求和 0.91 比加权求和基线+0.05，是所有方法中最高
人类偏好PickScore PickScore 22.83 DiffusionNFT加权求和 21.53 / DiffusionNFT顺序 23.80 比加权求和+1.30，仅略低于顺序调度方案
图文对齐CLIPScore CLIPScore 0.286 DiffusionNFT加权求和 0.267 / DiffusionNFT顺序 0.293 比加权求和+0.019，介于顺序与基线之间
美学迁移（held-out） HPSv2.1 0.355 DiffusionNFT加权求和 0.300 / SD3.5-L 0.288 比加权求和+0.055，比预训练SD3.5-L强23%
美学评分（held-out） Aesthetic 6.59 DiffusionNFT加权求和 6.15 / FLUX.1-Dev 5.71 比加权求和+0.44，超越闭源FLUX.1-Dev
综合排序（Composite z-score） Composite +1.116 DiffusionNFT顺序 +1.015 / 加权求和 +0.184 第一，比第二名+0.101，比加权求和+0.932
训练速度（8×H200） 相对速度 0.97× 加权求和基线 1.00× / 非摊销MARBLE 0.56× 摊销设计将速度从0.56×恢复到0.97×

局限与改进

作者明确承认的局限有两点：其一，MARBLE的实验目前集中在图像生成，未验证视频扩散与世界模型，而后者涉及时序一致性、运动合理性、物理一致性等更异质的奖励维度，平衡难度更大；其二，当前只在5个奖励上验证，未在大规模奖励集（如10+）上扩展到$N$和$\rho$如何随奖励数量自适应。本人观察到的额外限制包括：(1)MARBLE依赖DiffusionNFT的affine结构，对未来采用非affine损失（如基于PPO clip-ratio的目标）的方法不可直接迁移，需要重新证明等价性；(2)整体评估仍基于SD3.5 Medium的LoRA微调（rank 32），未在更大backbone（如FLUX.1-Dev或更大SD变体）或全参数微调下验证扩展性；(3)Composite指标是列内z-score，对个别极端样本敏感，且held-out奖励只有3个，泛化性证据偏少；(4)论文未深入分析EMA衰减$\rho=0.7$与摊销间隔$N=10$的耦合关系，仅放在附录；(5)对超长prompt或极端长尾奖励（如安全/有害性），未给出失败案例分析。

独立分析的弱点

独立分析可见的弱点有四。第一，QP求解对梯度归一化的数值稳定性高度敏感，消融中'w/o gradient normalization'直接'FAIL'，作者没有给出回退策略或鲁棒求解器（如信任域QP）替代方案。第二，摊销间隔$N=10$和EMA系数$\rho=0.7$都是超参，附录C.4/C.5虽给出敏感度但未提供自动调节机制，多奖励场景下需手工调。第三，论文用KL正则作为'调和之外的独立项'处理，但KL项$\beta_{\mathrm{KL}}\nabla_\theta D_{\mathrm{KL}}$本身可能与某些奖励梯度冲突，没有分析这种KL冲突是否会让均衡方向$d^*$偏移最优解。第四，主观质量评估仅做定性可视化（Figure 4），缺乏用户研究或大规模美学A/B测试，难以确认在真实人类偏好上是否优于顺序方案；改进方向：引入偏好模型做自动评估、把KL纳入QP约束、使用元学习动态调整$N$和$\rho$。

未来方向

作者明确提出的方向有二：把MARBLE从图像扩散扩展到视频扩散与世界模型，处理时序一致性、运动合理性、物理一致性等更复杂奖励；以及把奖励集规模扩展到10+个，研究$N$、$\rho$等超参随奖励数量的自适应。基于方法结构还可以延伸三个方向：(1)把'专家样本'诊断推广到LLM RLHF，研究GRPO在多偏好对齐上是否存在类似80%负相关现象；(2)用元梯度或上下文bandit在线学习$\alpha^*$的先验，让单纯形求解从一个batch扩展到序列级；(3)与直接偏好优化（DPO）类无RL方法结合，探索是否能在不显式采样的情况下实现多目标均衡；(4)把MARBLE的affine摊销思想推广到任何线性优势-损失映射的RL算法。

复现评估

论文在'Homepage and code repo: HERE'处给出代码链接（截至本次解读未访问确认），方法超参披露充分：优化器AdamW、学习率$3\times 10^{-4}$、LoRA rank 32 alpha 64、$A_{\max}=5$、EMA $\rho=0.7$、摊销间隔$N=10$（见附录）、KL强度$\beta_{\mathrm{KL}}$未显式给出但可从DiffusionNFT官方代码读取。算力方面，主实验使用16×H200 GPU，效率对比使用8×H200，门槛较高但非罕见配置。训练数据为常规文生图prompt集合（具体数据集未在主文披露，建议参考附录）。奖励模型全部开源/可用：PickScore、HPSv2、CLIPScore、OCR、GenEval、Aesthetic、ImageReward、UniReward。复现难度中等偏上：算法本身易于实现（核心是per-reward梯度计算+QP求解），但五奖励联合训练+8卡H200的算力门槛对个人研究者较高，最大的工程难点是确保$K$次反向传播之间的数值一致性。

论文图表

Comparison of multi-reward training paradigms. Left: Training one model per reward requires maintaining multiple models and cannot generalize across reward dimensions. Middle: Sequential multi-reward training produces a single model but demands extensive hyperparameter tuning and handcrafted stage schedules. Right: MARBLE trains a single model on all rewards simultaneously with minimal manual effort.
Figure 1: Comparison of multi-reward training paradigms. Left: Training one model per reward requires maintaining multiple models and cannot generalize across reward dimensions. Middle: Sequential multi-reward training produces a single model but demands extensive hyperparameter tuning and handcrafted stage schedules. Right: MARBLE trains a single model on all rewards simultaneously with minimal manual effort.

三个并列的概念图，左边是'一个奖励一个模型'（Model A/B/C各自专精PickScore/OCR/...），中间是'顺序多阶段训练'（一个模型在Stage 1-3依次训练各奖励），右边是MARBLE的'同步多奖励训练'（一个模型同时接受所有奖励监督），并标注三者的优缺点。

这是论文核心立论的'一张图说清动机'，直观展示MARBLE相对于两种现有范式的优势：避免维护多模型、避免手工调度、单一模型统一对齐。对理解整个论文定位至关重要。

Sample-level specialist structure. Each column denotes one sample, and each row shows its per-reward z-score advantage $A_k(x)$. High advantages are concentrated on source-specific rewards such as OCR or GenEval. Few samples achieve positive rewards across all dimensions.
Figure 2: Sample-level specialist structure. Each column denotes one sample, and each row shows its per-reward z-score advantage $A_k(x)$. High advantages are concentrated on source-specific rewards such as OCR or GenEval. Few samples achieve positive rewards across all dimensions.

左半部分是OCR来源样本的119个样本在OCR/PickScore/CLIPScore三行上的z-score优势热力图，可以看到OCR行的高分集中在少数列，而PickScore行的高分分散在不同列；右半部分是GenEval来源的238个样本在GenEval/PickScore/HPSv2/CLIPScore四行上的类似分布。

这是论文'专家样本现象'的实证基础，把抽象的'样本-奖励错配'变成可视化证据，是3.2节'为什么标量聚合失败'的图像化解释，必须看到才能理解MARBLE为何不在奖励空间做加权。