Gradient Preconditioning for Efficient and Reliable Reward-Guided Generation¶

会议: ICML 2026
arXiv: 2602.08646
代码: 待确认
领域: 图像生成 / 扩散模型 / Test-time 优化
关键词: reward-guided generation, 一步生成模型, 白高斯噪声约束, 梯度预条件, 谱域投影

一句话总结¶

通过把 reward 梯度投影到一个用 DFT 块状 \(\ell_1/\ell_2\) 范数刻画的"白高斯噪声可行集"上，作者把一步生成模型的 test-time latent 优化变得既快又稳：在 FLUX 上只用 30% 的 wall-clock 时间就追平 SOTA 正则化方法 MPGR 的 Aesthetic Score，并彻底避免 reward hacking。

研究背景与动机¶

领域现状：随着 shortcut / consistency 等蒸馏技术让扩散与流模型可以"一步出图"，在推理阶段直接对 latent 噪声 \(\bm{x} \in \mathbb{R}^N\) 做梯度上升以最大化某个 reward \(r(\mathcal{M}(\bm{x}))\) 成为热门方向（ReNO、MPGR、ORIGEN 等）。它不需要重新训练，可以即插即换 reward，是当下 reward-guided generation 最轻量的路线。

现有痛点：这套 test-time latent 优化在实践里有两个死结。其一是 reward hacking——latent 沿梯度跑着跑着就偏离了白高斯先验，模型开始产生伪影甚至崩坏的图，但 reward 数值反而被刷得很高。其二是慢——即使底层是一步模型，单张图常常也要做上百次梯度更新、耗费数十秒到数分钟。

核心矛盾：现有方法（ReNO、PRNO、MPGR）走的都是软正则路线，在目标函数里加一项 \(-\lambda \mathcal{L}_{\text{reg}}(\bm{x})\) 来鼓励 latent 保持高斯性质（\(\ell_2\) 范数、谱域块状 \(\ell_1\) 等）。但软正则有三重缺陷：(i) 不保证 latent 真的留在 noise-like 区域，只是"鼓励"；(ii) 需要手动调 \(\lambda\)，权重和学习率耦合在一起；(iii) 当优化器找到 shortcut（比如把某个频率分量拉爆）时，软惩罚根本拦不住。

本文目标：把"保持白高斯性质"从软约束升级为硬约束，同时不能牺牲速度（每步都要做投影，所以投影必须是闭式且至多 \(\mathcal{O}(N \log N)\)）。

切入角度：作者注意到 MPGR 的谱域块状 \(\ell_1\) 惩罚已经能很好刻画白噪声谱平坦性，那么是否可以直接把它升级成"硬集合"并对梯度做投影？直接对原 DFT 系数 \(\hat{\bm{x}} = \bm{F}\bm{x}\) 做投影行不通——实数信号的 DFT 满足 Hermitian 对称性，块之间彼此耦合、有的系数实、有的复，没有简单闭式解。但如果先剥掉 Hermitian 冗余，把独立自由度重组为一个紧凑的复向量 \(\bm{y} \in \mathbb{C}^{N/2}\)，问题就解耦成了 \(P\) 个独立的"\(\ell_1\) 球与 \(\ell_2\) 球交集"上的投影，而后者有已知闭式解。

核心 idea：用一个 bijection \(\mathcal{F}: \mathbb{R}^N \to \mathbb{C}^{N/2}\) 把白高斯先验映射到紧凑谱域，在那里定义"每个 size-\(B\) 块的 \(\ell_1\) 与 \(\ell_2\) 范数都等于其在 \(\mathcal{CN}(0,1)\) 下的期望"作为可行集 \(\mathcal{G}_{\mathbb{C}}\)，每次把 reward 梯度投影回 \(\mathcal{G}_{\mathbb{R}} = \mathcal{F}^{-1}(\mathcal{G}_{\mathbb{C}})\)，得到 noise-aligned 的更新方向。

方法详解¶

整体框架¶

输入是一个文本 prompt 与一个一步生成模型 \(\mathcal{M}: \mathbb{R}^N \to \mathbb{F}\)（论文用 FLUX-schnell，\(N = 65{,}536\)），输出是优化后的 latent \(\bm{x}^* \in \mathbb{R}^N\) 与对应图像 \(\mathcal{M}(\bm{x}^*)\)。流程是一个简洁的循环（Algorithm 1）：

repeat:
    J  ←  r(M(x))
    g  ←  ∇_x J
    g' ←  Proj_G(g)              # ← 本文的核心
    x  ←  Adam(x, g')

与软正则路线 \(\max_{\bm{x}} r(\mathcal{M}(\bm{x})) - \lambda \mathcal{L}_{\text{reg}}(\bm{x})\) 相比，这里没有 \(\lambda\)，且不要求 \(\bm{x}\) 本身落在 \(\mathcal{G}\) 里——只要每步的"更新方向"落在 \(\mathcal{G}\) 里就够了，这样既能避免 reward hacking，又能把搜索集中在与白噪声兼容的子空间。

整个流水线的难点全部集中在投影算子 \(\text{Proj}_{\mathcal{G}}\) 的设计与高效实现，下面三个关键设计就在解决这件事。

关键设计¶

紧凑谱域 bijection \(\mathcal{F}: \mathbb{R}^N \to \mathbb{C}^{N/2}\)：
- 功能：把实信号 DFT 中由 Hermitian 对称性带来的冗余完全剥掉，使后续投影问题在结构上解耦。
- 核心思路：对偶数维 \(N\)，DFT 系数 \(\hat{\bm{x}}\) 中只有 \(x^*_0, x^*_{N/2}\) 是实数，且 \(\hat{x}_k = \overline{\hat{x}_{N-k}}\)，所以独立自由度只有 \(N/2\) 个复数。作者定义 \(y_0 = \tfrac{\hat{x}_0}{\sqrt 2} + \tfrac{\hat{x}_{N/2}}{\sqrt 2} i\)、\(y_k = \hat{x}_k\)（\(k = 1, \dots, N/2-1\)），构成 \(\bm{y} = \mathcal{F}(\bm{x})\)。Proposition 4.1 证明 \(\mathcal{F}\) 是 \(\mathbb{R}^N \leftrightarrow \mathbb{C}^{N/2}\) 双射，且 \(\bm{z} \sim \mathcal{CN}(\bm 0, \bm I_{N/2})\) 当且仅当 \(\mathcal{F}^{-1}(\bm{z}) \sim \mathcal{N}(\bm 0, \bm I_N)\)；Proposition 4.2 进一步给出 \(\|\mathcal{F}^{-1}(\bm z)\|_2^2 = 2\|\bm z\|_2^2\)。这两条把"空间域上的高斯白噪声约束"等价转换为"紧凑谱域上的复高斯约束"。
- 设计动机：直接在 \(\hat{\bm{x}}\) 上写硬约束会因为 Hermitian 耦合导致不同块之间纠缠、没有闭式投影；剥掉冗余之后，约束自然按块解耦，下一步才能做高效投影。
块状 \(\ell_1/\ell_2\) 双范数可行集 \(\mathcal{G}_{\mathbb{C}}\)：
- 功能：在紧凑谱域 \(\mathbb{C}^{N/2}\) 上定义一个紧致集合，精确刻画白高斯噪声的统计特性，且比单一 \(\ell_2\) 或 \(\ell_1\) 约束更紧。
- 核心思路：将 \(\bm{y}\) 切成 \(P = N/(2B)\) 个 size-\(B\) 块（论文取 \(B = 16\)），对每块同时强制其 \(\ell_1\) 与 \(\ell_2\) 范数等于 \(\mathcal{CN}(0,1)\) 下的理论期望：\(\|\bm{y}^{(p)}\|_1 = \tfrac{\sqrt\pi}{2}B\) 且 \(\|\bm{y}^{(p)}\|_2^2 = B\)。空间域可行集定义为 \(\mathcal{G}_{\mathbb{R}} = \mathcal{F}^{-1}(\mathcal{G}_{\mathbb{C}})\)。这一来 \(\ell_2\) 约束保证整体能量与 \(\chi_N\) 分布的众数 \(\|\bm{x}\|_2^2 = N\) 完全一致（与 \(\ell_2\) norm 正则 \(\mathcal{L}_{\text{norm}}\) 的极小点几乎重合），\(\ell_1\) 约束则压抑任一单频率分量主导谱（理论上每块内 \(|y_j|^2\) 的最大值仅约 \(7.18\)，而总预算 \(N/2 \gg 10^4\)），从而对应白噪声的"无主导频率"性质。作者用 1.1M 个高斯样本验证：\(\bm{x} \sim \mathcal{N}(\bm 0, \bm I_N)\) 与其在 \(\mathcal{G}_{\mathbb{R}}\) 上的投影之间的余弦相似度最低也有 \(0.988\)，说明真实白噪声本就贴着这个可行集。
- 设计动机：比起仅匹配总 \(\ell_1\)、\(\ell_2\)（只两个全局等式），按块做 \(2P\) 个等式形成的可行集严格更小，是对白噪声分布更紧的刻画；比起 MPGR 仅惩罚 \(\ell_1\) 偏差的软正则，硬集合还额外卡住 \(\ell_2\) 能量，把"shortcut 解"彻底排除在外。
\(\mathcal{O}(N\log N)\) 闭式投影算子 \(\text{Proj}_{\mathcal{G}}\)：
- 功能：给定任意 \(\bm{x} \in \mathbb{R}^N\)（这里其实是 reward 梯度），在每个优化步用接近 FFT 的复杂度找到它在 \(\mathcal{G}_{\mathbb{R}}\) 上的最近点。
- 核心思路：先用 FFT 算 \(\bm{y} = \mathcal{F}(\bm{x})\)；由 Proposition 4.2 的等距关系 \(\|\mathcal{F}^{-1}(\bm y) - \mathcal{F}^{-1}(\tilde{\bm y})\|_2^2 = 2\|\bm{y} - \tilde{\bm y}\|_2^2\)，空间域的最小化问题在紧凑谱域里有相同最优解；进一步因为可行集按块独立，问题分解为 \(P\) 个独立的"\(\ell_1\) 球与 \(\ell_2\) 球交集"上的投影，每个用 Liu et al. (2020) 的算法以 \(\mathcal{O}(B\log B)\) 解决。具体地：把块内 \(|y_j|\) 降序排成 \(\bm{w}\)，算前缀和 \(S_{d,k} = \sum_{l=0}^k w_l^d\)（\(d = 1, 2\)），找唯一的 \(k^*\) 满足 \(w_{k^*+1} \le \lambda^{(k^*)} < w_{k^*}\)，其中 \(\lambda^{(k^*)} = \tfrac{S_{1,k^*}}{k^*+1} - \tfrac{\sqrt\pi}{2}\tfrac{\sqrt B}{k^*+1}\sqrt{\tfrac{(k^*+1)S_{2,k^*} - S_{1,k^*}^2}{k^*+1 - \tfrac{\pi}{4}B}}\)，再用一个 ReLU 软阈值表达式 \(\dot y_j = \tfrac{\sqrt\pi}{2}B \cdot \tfrac{\text{ReLU}(|y_j| - \lambda^{(k^*)})}{S_{1,k^*} - (k^*+1)\lambda^{(k^*)}} \cdot \tfrac{y_j}{|y_j|}\) 直接给出投影结果，最后逆 FFT 回空间域。整体 \(\mathcal{O}(N\log N)\)，在 FLUX 上实测仅占整次迭代 wall-clock 的 0.04%。
- 设计动机：每步都得做投影，速度是底线；FFT + 块内闭式解的组合让"硬约束"在工程上完全免费。相比之下，把 \(\mathcal{L}_{\text{power}}\) 的极小集当可行集做投影需要数百步内层梯度下降（MPGR 在 Figure 1 里就是这么做的），既慢又只是近似最优。

损失函数 / 训练策略¶

没有显式损失项，只有 reward \(r\) 与投影约束。优化器是 Adam（学习率 FLUX 0.02 / SDXL-Turbo 0.1），梯度裁剪 0.03；FLUX 跑 200 步、SDXL-Turbo 跑 50 步，整套实验都在单张 A6000 上完成。每步既投影梯度，也投影 latent 本身（论文附录 H 给了拆开二者的消融）。

实验关键数据¶

主实验¶

评测设置沿用 MPGR：以 Aesthetic Score / PickScore / HPSv2 / ImageReward 中的一个作为"被优化 reward"，其余三个作为 held-out 用于检测是否 reward hacking；prompt 集来自 animal dataset 与 T2I-CompBench++；底层一步模型主要用 FLUX-schnell。

方法	迭代数	Aesthetic Score (target) ↑	PickScore (held-out)	Wall-clock (s) ↓
No Opt.	0	5.99	0.219	—
ReNO	200	7.06	0.219	232.0
PRNO	200	7.02	0.218	255.4
MPGR (SOTA)	200	7.13	0.220	235.5
Ours (60 iters)	60	7.12	0.220	69.7
Ours (200 iters)	200	8.91	0.220	232.2

两个数字最值得记住：(1) 60 步、69.7 秒就追平 MPGR 200 步、235.5 秒的成绩——30% wall-clock 时间达到 SOTA；(2) 同样 200 步，本文 reward 飙到 8.91 而 baseline 卡在 7.13 一线，差距高达 1.8 分（Aesthetic 通常 6–9 这个区间）。Figure 2 还显示在四个 reward × 三个 held-out 的所有组合里，本文都形成最右上的 Pareto 前沿。

消融实验（与现有正则方法相比的"reward hacking 抗性"）¶

配置	现象	说明
No Reg.（无正则）	Aesthetic 飙高但图像崩坏	标准 reward hacking
\(\mathcal{L}_{\text{norm}}\)（\(\ell_2\) 正则）	余弦相似度 0.222	只约束总能量，空间相关性失控
MPGR（谱域 \(\ell_1\) 软正则）	余弦相似度 0.548；需数百步内层投影	软惩罚 + 慢
Ours（硬集合 + 闭式投影）	保持高余弦相似度且无 hacking；单步投影	Figure 1 中 latent 与初始噪声仍高度对齐，图像清晰真实

多样性消融（1,125 张图，Aesthetic 优化下）：本文 IS = 21.10 / Vendi = 6.97，与未优化基线（IS 21.57–22.33、Vendi 6.42–6.61）方差内一致，没有 mode collapse。

关键发现¶

投影开销可以忽略：FFT + 块内闭式解整体 \(\mathcal{O}(N\log N)\)，实测仅占总 wall-clock 的 0.04%，几乎不增加单步耗时。
"每步走得准"比"步数多"更重要：在 noise-aligned 子空间内更新让每步 reward 增益更大，因此 60 步就够、还不会 hack。
硬约束 > 软正则不止快，而且稳：Figure 3 的定性结果显示 PRNO/MPGR 偶尔仍会生成不合 prompt 或带伪影的图，本文每张都合 prompt（包括"skyline pierced the clouds"这种细节）。
可行集很贴近真实白噪声：1.1M 高斯样本的投影余弦相似度最低 0.988，说明这套硬约束并没有"勒死"先验，而是顺着它做剪裁。

亮点与洞察¶

优雅的"剥冗余"trick：用 \(\mathcal{F}\) 把 Hermitian 对称性这个"坏耦合"一次性消掉，把整个谱投影问题拆成 \(P\) 个独立的小问题——这是论文最 elegant 的一步，没有这一步就没有闭式投影。值得在任何要在"实信号谱域"做约束/正则/采样的工作里复用。
从"软正则"到"梯度预条件"的视角转换：传统正则把噪声性写进目标函数，本文把它写进优化几何。优化器仍然在做原问题的 reward ascent，但行进方向被强制限制在 noise-aligned 子空间——这是一种比加正则项更强、却又比硬约束 latent 更灵活的中间路线（latent 本身不必在 \(\mathcal{G}\) 里）。这个思路可以原样迁移到 RLHF/Inference-time scaling 等任何对 latent 做梯度优化的场景。
0 超参 + 0.04% 开销：相比所有 baseline 都要调的 \(\lambda\)，本文不引入任何新超参，且投影开销几乎免费——这种"减法式贡献"在工程上极具吸引力。
拿真实样本的余弦相似度反向验证可行集设计：1.1M 个 \(\mathcal{N}(\bm 0, \bm I)\) 样本与其投影的最小余弦相似度 0.988 这个实验非常巧妙，等于反过来证明"我设计的硬约束不会扭曲先验"，比单纯做消融更有说服力。

局限与展望¶

当前算法只验证了一步生成模型（FLUX-schnell、SDXL-Turbo、SANA-Sprint、SD-Turbo），多步扩散/流模型的中间状态如何套用这套硬约束（需要在哪一步、是否对所有 timestep 都投影）论文没回答。
可行集是用 \(\mathcal{CN}(0,1)\) 的理论期望做等式约束，但真实采样下每块的范数本身有方差；当 latent 维度较小、块数 \(P\) 较少时，"严格等于期望"可能比"允许小波动"更糟（论文用大 \(N\) 缓解了这点，但小模型场景可能要重新讨论）。
论文从头到尾都假设 latent 先验是各向同性高斯，对非高斯先验（categorical token、VQ index、Mixture of Gaussians 等）没有直接推广路径。
Block size \(B = 16\) 沿用 MPGR 设置，论文附录 G 讨论了影响但没给出"如何在不同模型上自适应选 \(B\)"的实操指南。
实验只在图像生成上做了完整对比，视频、3D、motion 等高维 latent 场景（噪声维度可能更大、Hermitian 结构更复杂）值得进一步验证。

评分¶

新颖性: ⭐⭐⭐⭐ 把谱域软正则升级为硬可行集 + 利用 Hermitian 对称性设计 bijection 拿到闭式投影，是干净漂亮的技术贡献，但整体仍站在 MPGR 的肩膀上。
实验充分度: ⭐⭐⭐⭐ 四个 reward × 多个一步模型 × 完整 baseline 对比 + 多样性 + 余弦相似度反向验证 + 投影时间分解，链路非常完整；扣一星是只做了图像生成。
写作质量: ⭐⭐⭐⭐⭐ 数学严谨、动机推导清晰、Figure 1 的可视化把方法优势讲得极透，附录把所有 proposition 都补全了。
价值: ⭐⭐⭐⭐⭐ 0 超参、0.04% 开销、30% 时间打平 SOTA、彻底解决 reward hacking——对所有做 test-time reward optimization 的人都是即插即用的硬升级。