潜空间奖励寄存器:让扩散模型在生成中理解偏好

📅 2026/8/27 8:31:56
潜空间奖励寄存器:让扩散模型在生成中理解偏好
先讲一个我在日常图像生成中反复遇到的场景在 stable diffusion 里输入一大段经过精心挑选的提示词生成第一版结果看起来构图、光影、色彩都还行但人物五官不对、透视别扭、或者风格跟预期差很远。于是只能继续加大关键词换随机种子反复抽卡甚至有人为了稳定效果开始研究 ControlNet、LoRA、负面提示词、以及 webui 和 comfyui 的差异。这其实是很多生成工具用户的共同体感模型并不是“画得不够好”而是“不知道在你心里什么才算好”。Prompt 只是文字指令它没法承载你对构图、人物比例、光影氛围、审美倾向的细腻判断。于是大家默认把“抽卡”当成工作流的一部分。但抽卡是概率游戏。真正高的壁垒是让生成过程本身理解偏好并在生成中途不断修正方向。这正是“偏好对齐”要解决的问题。而最近我注意到一个听起来很技术性的方向——Latent Reward Registers for Diffusion Preference Alignment把奖励信号写进潜空间做成一种可以参与采样的“寄存器”。这篇想聊聊我自己的理解它解决了什么问题为什么比“最后再打分”更难以及如果我们要跟进这个方向应该从哪里开始。1. 扩散模型的隐性门槛不是不会画而是不知道“什么叫好”1.1 扩散模型生成的本质是分布采样不是目标搜索扩散模型的基本逻辑是从一个噪声分布开始通过多步去噪逐步恢复出符合训练数据分布的结构。理论上只要训练数据足够丰富模型就能生成高质量图像。但高质量是统计意义上的。模型看到的是“大量图片的平均感觉”无法理解你个人对“好”的定义。举个例子如果训练集里大量图片是低饱和、胶片质感的模型就会倾向于生成这种风格哪怕你写的是“色彩鲜艳”它也只会往鲜艳方向轻微偏移。因为模型的核心任务是最大似然估计不是满足某一个人的主观偏好。所以我们在实际使用中会发现同一组提示词换个随机种子结果可以完全不同。这不是随机性问题而是模型在统一分布里随机抽样。你想要的是一张“符合一定审美标准的图”模型只是在“所有可能的图”里随机取一个点。1.2 事后评估只能筛选不能纠正生成路径传统上我们会用 FID、CLIP Score、美学评分模型或者人工打分来评估图像。但这些评估都发生在生成结束之后。这种“事后评估”思路有一个天然问题它只能告诉我们“这张图可能不好”但没法告诉生成过程“在哪一步、哪个特征维度上跑偏了”。就像餐厅里客人吃完后说“这道菜太咸”厨师知道咸度出了问题却不知道是下盐的哪一步放多了更不能在炖煮过程中自行判断、及时调整。很多生成应用的工具链也是这种模式。写 prompt → 生成 → 人工判断 → 调 prompt → 再生成。工具在这里只负责“生成”判断和修正全靠人。所以问题的关键不是模型不能画而是它缺少一个“内部质检系统”在实际生成过程中实时评估当前状态是否符合偏好并利用这个信号调整后续去噪方向。1.3 偏好对齐是“让模型知道偏好可以参与生成”在语言模型领域我们已经看到 RLHF、DPO 等方法通过人类偏好数据来调整生成策略。但图像生成有自己的特殊性图像没有唯一正确答案偏好表达不稳定而且生成过程是一个多步迭代潜变量更新的过程不是一次性的字符预测。因此把偏好信号引入扩散模型不能简单照搬文本模型的做法需要一种更适合潜空间迭代结构的形式。Latent Reward Registers 这个方向本质上就是想在潜空间里维护一套“可读写的偏好状态”让奖励信息不只在最终输出时出现而是能介入中间步骤。2. 为什么文本模型的对齐方案在扩散模型里会失效2.1 文本模型的决策是离散的扩散模型是连续的迭代过程语言模型生成每个 token 是一个离散决策RLHF 可以把每一步 token 选择看作一个状态并基于奖励模型给出即时或最终回报。这个过程虽然延迟但至少每一步的决策有明确的逻辑归属。扩散模型不一样。从 pure noise 到 final image每一步都在一个连续的高维潜空间中做微小修正。你不容易说清楚是哪一步的噪声预测改变了人物的眼睛形状或者哪一步导致光线不自然。如果把奖励模型的最终分数直接反传回所有步梯度信号会被这个多步迭代过程稀释很难精确控制。简单说文本对齐像教师批改一道数学题可以指出第一步错在哪里而扩散模型像教师看到一份长达几十步的草稿只知道最后答案不对却没法指出哪一步出了问题。2.2 稀疏奖励无法指导密集去噪过程奖励模型通常对完整图像输入打分输出一个标量。这个标量是全局的整体构图是否合理、色彩是否协调、风格是否符合 prompt。但去噪过程每一步都涉及很多潜变量在不同尺度、不同通道上这些局部特征和最终偏好之间的关系非常复杂。一个全局标量无法提供细粒度的“在哪个特征通道、哪个空间位置、哪个尺度上需要加强或减弱”。这种“稀疏奖励”在强化学习里已经有大量研究典型解法是设计 dense reward 或 reward shaping。Latent Reward Registers 在我看来就是一种 reward shaping 的实现思路通过额外训练一组潜空间状态表示把全局奖励信号转译成每个时间步、每个 token 位置可以使用的局部指导信号。2.3 直接反传梯度会遇到尺度与稳定性的工程问题如果尝试直接把奖励模型接到扩散模型输出上再对 U-Net 或 DiT 的所有参数做梯度回传会出现两类问题奖励模型对图像质量的评分函数通常已经饱和梯度在优化一段时间后会变得非常小尤其是当你已经用了 CFG、负面提示词等技巧后评分差异更微弱对扩散模型的去噪网络同时更新会导致生成质量下降因为它不是只在学习“朝奖励高的方向走”还会被奖励模型中的噪声和偏置带偏。因此需要一种机制能在潜空间中缓存并更新奖励信号在每一步去噪时只做轻量修正。这比端到端联合优化更稳定也更容易控制。3. Latent Reward Registers 的直觉把奖励信号搬进潜空间里3.1 “寄存器”这个词的启发寄存器在计算机体系结构里是一种高速临时存储单元用于在计算过程中保存中间值供后续指令使用。它不是一个静态的存储而是在每个周期被读取、更新、再写入。借用这个概念Latent Reward Registers 可以理解为在扩散模型的潜空间特征层上放置一组额外的可学习向量这些向量被当作“偏好状态寄存器”。它们会在每个去噪步骤中读取当前中间特征根据一个奖励模型或偏好模型更新自身状态并参与下一个去噪步骤的注意力计算。这就是整个方法最核心的部分奖励信号不再只是最终输出上的一个数字而是存在于潜空间中、随着生成过程动态变化的“一组状态”。3.2 一种可能的架构思路如果从工程实现角度来推演一个比较自然的方案是在 U-Net 或 DiT 的某个中间注意力层中学习一些新增的 token类似 class token 或 register token 的扩展。假设输入是扩散模型的中间隐变量 z_t条件特征是 c。我们可以在注意力层里增加 K 个可学习的寄存器向量 r_1, ..., r_K初始值可以为零或可学习参数。在每一步将 z_t、条件 c 和寄存器向量一起送入注意力层注意力机制会让寄存器和图像 token 互相关注再额外引入一个奖励模型 R在关键时间步把当前潜变量的某个投影输出与寄存器状态一起送入 R得到一个奖励预测根据预测更新寄存器的状态让下一轮去噪时模型能够感知到偏好偏差。这样寄存器起到一个“中间状态的偏好评测员”的作用。它不只是被动地记录还会根据奖励模型的反馈主动调整自己从而影响生成。这个方向不同于简单地把奖励分数拼进条件向量因为寄存器是在生成过程中不断更新的更像一个内置的闭环控制。3.3 与 CFG、DPO 等方法的关系有人可能会问稳定扩散里已经有无分类器引导Classifier-Free Guidance通过拉大条件生成与无条件生成的差距来增强 prompt 服从性为什么还需要 register我的理解是CFG 本质上是把“是否跟随条件”作为一个控制维度它并不理解“跟随条件”之后的结果是否符合人类偏好。它只会让模型更严格地按照文本条件走但 prompt 本身表达不了审美所以 CFG 只是在服从指令层面做强化。DPO直接偏好优化则是先准备人类偏好的成对数据直接优化策略模型使生成结果更偏向“被喜欢”的一侧。但它仍然是在训练阶段把偏好编码进模型参数。一旦模型训练好生成时就是一个固定策略。Latent Reward Registers 若被设计为带推理时引导的方法它的价值在于模型可以先保持通用能力在使用时根据奖励模型动态调整偏好方向。这意味着同一个模型可以通过更换 reward model 或调整寄存器更新权重来适配不同审美偏好而不是每次重新训练模型。如果这个假设成立那么它真正改变的不只是生成质量而是“偏好定制”的成本结构。4. 从概念到落地这条路线需要面对的四层边界4.1 偏好数据的颗粒度决定寄存器能学到什么训练一个有效的 reward register至少需要两类信号成对偏好数据同一 prompt 下的两组图片哪一组更符合需求中间过程状态数据生成到第几步、哪个中间特征有助于预测最终偏好。很多公开偏好数据集只提供最终图像的排序缺少中间潜变量的标注。没有中间信号寄存器很难学到“在第几步应该调整什么”。我实际做算法实验时的第一反应是先从已有的 final-preference 数据里蒸馏出一个偏好打分模型再用它给中间潜变量打分生成回归标签。这样虽然会引入噪声但比没有强。真正要落地可能还需要采集更多用户对中间状态的反馈比如“这个阶段很好不要过头了”这种 step-level 标注这类数据目前非常稀缺。4.2 不是所有潜空间位置都适合做奖励寄存器扩散模型的中间特征在不同层、不同时间步上的语义差异很大。浅层靠近噪声主要是低频结构深层靠近图像主要是高频细节。如果在不合适的位置强制插入寄存器可能不但不会引导偏好反而破坏生成分布。更稳妥的做法是在几个关键层上分别做小规模对照实验比如在注意力层、交叉注意力层、以及最底层特征上各放一组寄存器然后比较生成效果。这里不能想当然“越深越好”。从工程经验看最值得优先尝试的可能是中低层因为这些位置既保留了足够的全局结构又还没有过度受到纹理细节影响。偏好信号往往与构图、风格、氛围有关这些信息在中低层更有表达能力。4.3 计算代价会在推理阶段显著增加加入寄存器不是零成本。每多一组寄存器 token注意力矩阵的计算量就会增加如果寄存器在每一步都要读取奖励模型并更新状态推理时间会明显上升。对于需要实时或批量生成的应用这可能成为拦路石。如果目标是移动端或 Web 端体验很可能需要在推理效率和引导强度之间做折中。比如每 N 步才更新一次寄存器或者在较低分辨率阶段启用寄存器在高分辨率细化阶段关闭。在这个方向没有成熟优化前不要把“0 额外开销”作为预期。前期做实验时可以用 50 步采样中的最后 20 步使用寄存器看能不能获得大部分收益。4.4 评估难度比方法本身更大怎么判断“寄存器有没有用”最简单的做法是看最终图像的人类偏好率。但这无法区分是否是“奖励过拟合”导致的可能寄存器只是让模型记住了奖励模型的审美好恶而不是真正的语义偏好。为了避免这种假阳性测试集需要包含训练阶段没有见过的 prompt、不同风格、不同主体并且最好同时评估稳定性、多样性、以及关键特征的偏差程度。我建议多关注多样性指标。很多偏好对齐方法会让生成结果收敛到少数几个“最安全样式”上表面看符合偏好但实际失去了多样性。如果一个方法提高了偏好率但显著降低了多样性长期使用体验未必好。5. 给普通开发者的参考这个方向值得现在跟吗5.1 先判断你处于哪个角色场景如果你是在产品里使用 stable diffusion、midjourney 或类似生成 API 的开发者现阶段不建议把 Latent Reward Registers 作为依赖。因为你更需要的其实是把 webui/comfyui 的工程流程做扎实比如建立负面提示词库、用 ControlNet 固定结构、用 LoRA 控制风格这些手段已经可以解决大部分偏好问题。奖励寄存器可能在未来成为“推理时插件”但目前大概率还停留在研究或小规模验证阶段。普通开发者更适合关注方法论文是否发布、有没有开源权重、有没有完整的工程实现而不是自己从零复现。5.2 如果要在自己的模型上做验证我建议的最少实验路径假设你已经有一个开源的 diffusion 模型比如 Stable Diffusion 或 DiT 系列想快速验证这个方向是否对你有价值可以考虑以下步骤准备 2k 到 5k 条偏好数据从现有偏好数据集中筛选或自己用奖励模型打分生成先训练一个轻量级偏好奖励模型输入图像或中间潜变量输出偏好分数在模型中间注意力层插入少量可学习寄存器 token初始化为 pad token只训练寄存器 token 和少量归一化层冻结原模型主体生成时对比有/无寄存器的结果并统计偏好率、CLIP Score、多样性。这套路径的意图是“最小改动验证”避免一开始就训练全模型。寄存器 token 数量可以从 1 开始。如果单个 token 就能带来可感知的变化再考虑扩展到多个。5.3 一个评估决策框架当我们面对一个新对齐方法时可以从五个维度做判断维度关注点适用阶段训练时更新、还是推理时也能更新成本边界额外参数、推理延迟、数据需求可解释性寄存器状态意味着什么能否定位偏差收益范围只提升审美一致性还是也能提升 prompt 服从性工程友好度是否有 PyTorch 实现、权重、基准评测如果一种方法只在特定数据集上有效且无法给出中间状态分析我会更谨慎。如果一个方向能同时启发“推理时引导”和“反馈信号进入潜空间”这两件事那么不管具体论文是否能复现它都值得持续跟踪。结束前的一点判断Latent Reward Registers 这个名字表面上像是在讨论一种模型结构往深看它代表了一个更普适的趋势生成模型不能只学会“怎么画”还要学会在生成过程中不断衡量“画得对不对”。这种思路真正会改变的是我们和生成模型的协作方式。过去我们只能通过 prompt 表达偏好通过抽卡碰运气未来模型也许可以内置一个“偏好状态中心”每走一步都知道当前结果距离你想要的还有多远然后自己修正。不过我也提醒自己任何潜空间里的额外状态都只是辅助。生成质量的底座仍然是数据和模型容量。寄存器能做的是让已经学到的能力更容易被你指挥。如果你也想跟进这个方向我更建议先从一次小实验开始找一个偏好数据集在你的基础模型上加一个寄存器 token看看它能不能在一次生成里让图像从“一般”变成“更接近你想要的”。这样试过一次之后你再看那些论文就不会觉得抽象了。