最近在处理 Stable Diffusion 系列模型的微调与生成效果优化时我反复遇到一个问题生成质量明明在提升但用户主观感受却变差了甚至有几次出现了“伪高质量”图像——局部细节异常丰富、整体布局却完全不符合提示词。查了一圈资料后发现这类问题背后往往指向一个更底层的方向扩散模型的偏好对齐Preference Alignment。再往深走就会看到 Latent Reward Registers 这一类新方法。本文围绕 Latent Reward Registers for Diffusion Preference Alignment 展开先解释扩散模型偏好对齐的价值再拆解 Reward Hacking 的根因最后结合 diffusers 与 PEFT 给出一套可运行的实战示例。如果你正在做 Stable Diffusion 模型微调、LoRA 训练或者想理解 DPO、RLHF 这类对齐方法在扩散模型中的落地方式这篇文章会比较合适。1. 背景与核心概念1.1 为什么扩散模型需要偏好对齐扩散模型Diffusion Model在图像生成领域已经成为主流Stable Diffusion、Midjourney 以及各种开源微调模型都依赖它完成文字到图像的生成。早期大家的优化重点集中在“生成是否清晰”“分辨率是否够高”“细节是否丰富”上但随着模型能力增强一个新的问题开始凸显模型生成的结果到底是不是用户想要的这就是偏好对齐要解决的问题。偏好对齐的意思是让模型的输出尽可能符合人类偏好。比如用户输入“一只坐在沙发上的橘猫电影感光线”不同模型会给出不同结果。技术指标上可能都合格但有的结果更符合用户预期有的则显得机械或失真。偏好对齐就是通过奖励信号、排序数据或人工反馈让模型学到“什么是更好的生成结果”。这个方向和传统意义上的“提高图像质量”不同。质量提升关注的是清晰度、纹理、细节偏好对齐关注的是语义符合度、构图合理性、审美偏好等更主观的维度。1.2 Latent Reward Registers 是什么Latent Reward Registers 是一个相对前沿的概念翻译过来是“潜在奖励寄存器”。要理解它需要先拆成两个部分。Reward奖励在强化学习和偏好优化中奖励是一个分数用来表示某个生成结果的好坏。奖励可以来自人类打分也可以来自一个奖励模型。Register寄存器这个词借鉴了 Transformer 结构中的 register token。在视觉 Transformer 中研究者发现额外的 register token 可以吸收特征图中的异常高范数 token从而提升模型表现。Latent Reward Registers 的思想类似在扩散模型的潜在空间中加入一批“奖励寄存器”专门用来吸收奖励信号中的异常模式。1.3 它解决什么问题扩散模型偏好对齐的核心训练方式会引入一个奖励模型来指导生成。这个奖励模型打分时往往只关注某些统计特征而不是真正理解语义。如果生成模型过度迎合奖励模型就会出现 Reward Hacking也就是“刷分”。比如一个美学评分模型可能偏爱高饱和度图像那么生成模型就会逐渐把图片往高饱和方向调哪怕用户根本不想要这种风格。Latent Reward Registers 的思路是在潜在空间中提供额外的“记账单位”让奖励信号只在寄存器上累积而不是污染图像的真实语义特征。这样做的好处是奖励指导仍然生效但不会让模型为了奖励而过度扭曲主路径上的图像特征。2. 环境准备与版本说明开始实战之前先说明本文示例的运行环境。这里不会强制绑定某个具体版本因为该方向迭代很快不同时间点的库版本差异较大。2.1 推荐环境建议使用 Linux 或 macOS 系统配备 NVIDIA GPU 效果更佳。没有 GPU 也能运行极小规模的演示只是速度会很慢。Python 3.10 或 3.11PyTorch 2.0 及以上diffusers 0.25 及以上transformers 4.30 及以上peft 0.6 及以上accelerate 0.24 及以上datasets 2.14 及以上如果使用 CUDA建议提前安装好与显卡驱动匹配的 PyTorch 版本。2.2 安装依赖推荐创建一个干净的虚拟环境。python -m venv venv_diffusion_align source venv_diffusion_align/bin/activate pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install diffusers transformers accelerate peft datasets pip install matplotlib pillow安装完成后可以用一段简单代码验证环境是否正常。import torch import diffusers print(torch:, torch.__version__) print(diffusers:, diffusers.__version__) print(cuda available:, torch.cuda.is_available())如果输出cuda available: True说明 GPU 环境可用。2.3 示例项目结构本文的实战部分采用如下项目结构diffusion_preference_align/ ├── data/ │ └── prompts.txt ├── models/ │ └── lora_output/ ├── scripts/ │ ├── train_preference.py │ ├── generate_with_register.py │ └── reward_model.py └── README.md这样组织的好处是数据、模型、脚本分开便于后续扩展。3. 核心原理拆解偏好对齐、Reward Hacking 与寄存器机制这一节是全文的理论基础。理解清楚之后后面代码就是水到渠成的事情。3.1 偏好对齐的常规流程当前扩散模型偏好对齐的主流方式有两类一类是 RLHF 风格一类是 DPO 风格。RLHF 风格收集人类对生成图像的偏好排序数据。用这些数据训练一个奖励模型奖励模型输出一个分数。通过强化学习算法如 PPO微调扩散模型的生成策略让奖励分数最大化。DPO 风格收集偏好对例如(prompt, chosen_image, rejected_image)。直接用一个偏好损失函数优化策略不需要显式训练奖励模型。通过对比 chosen 和 rejected 的生成概率让模型学会偏好方向。DPO 的代码实现更简单不需要强化学习中的采样和奖励归一化也是目前开源社区里接受度较高的方案。3.2 Reward Hacking 的根因无论使用哪种方案一旦引入奖励信号就存在一个隐患奖励模型并不能完美刻画人类偏好。奖励模型通常是在有限的人类标注数据上训练出来的。它学到的往往是“高分图像”在统计层面上的特征而不是真正的审美逻辑。生成模型在优化时会不断试探奖励模型的边界寻找一种“奖励分数很高但人类并不认可”的输出。举个具体例子奖励模型偏爱细节丰富的图像。生成模型发现增加高频纹理能提升分数。于是模型开始生成“局部细节爆炸、整体构图崩塌”的图像。这种图像在奖励模型眼中是高分在人类眼中却是废片。这种现象就是 Reward Hacking也叫做奖励过度优化。它本质上是因为奖励信号里混入了某些异常特征而生成模型把这些异常特征当作可学习的“捷径”。3.3 从 Transformer Register 到潜在空间Register 的概念最早来自视觉 Transformer 的研究。研究者发现Vision Transformer 的中间特征图会出现少量高范数的异常 token这些 token 不是有效的语义特征而是注意力机制在训练过程中“偷懒”形成的模式。通过在特征序列中额外添加 register token可以吸收这些异常模式让原有 token 的表达更干净。Latent Reward Registers 把这一思想迁移到扩散模型上。扩散模型在潜在空间Latent Space中进行去噪每一步都在优化潜在特征。如果我们在潜在表示中额外拼接一组“奖励寄存器”并且在训练时让奖励梯度只流向这些寄存器那么奖励模型中的异常模式就会被集中吸收避免主干图像特征被污染。需要注意的是目前该方向仍然在快速演进中不同论文的实现细节不同。有的做法是在 UNet 的 cross-attention 层加入寄存器 token有的做法是在 VAE 的 latent 上操作有的则是在 noise predictor 的输出上做加权。本文的实战示例采用一种相对通用的简化思路目的是演示“如何把奖励信号限制在额外维度中”。3.4 Latent Reward Registers 的改进效果引入寄存器之后理论上可以得到几个改进偏好优化过程中生成模型的主干特征更稳定不容易被奖励模型带偏。当奖励模型出现局部过拟合时寄存器可以吸收异常梯度降低生成质量的震荡。在不明显改变训练体系的前提下通过修改潜在表示的拼接方式即可适配现有训练管线。当然这种方法不是万能的。如果奖励模型本身训练数据就严重偏斜寄存器只能缓解污染不能改变奖励模型的错误偏好。这也是为什么实际工程中始终建议“奖励模型要持续迭代”。4. 完整实战在 diffusers 中实现偏好对齐微调这部分我们写一个可运行的简化项目。项目目标用一个小型偏好数据集通过 LoRA 微调 Stable Diffusion 模型并且在微调过程中引入寄存器机制演示如何缓解 Reward Hacking。4.1 准备数据与奖励信号为了让示例不依赖大型标注数据我们构造一个简化场景对同一组提示词人工定义一种偏好规则。这里用提示词中是否包含指定风格词作为打分依据模拟一个“奖励模型”。# 文件路径scripts/reward_model.py import re from dataclasses import dataclass dataclass class RewardOutput: score: float explanation: str class SimpleRewardModel: 简化版奖励模型 - 提示词中包含 nature 或 cinematic 则加分 - 提示词包含 overexposed 则减分 真实项目中请替换为训练好的奖励模型或人工标注结果。 def __init__(self): self.positive_keywords [nature, cinematic, soft light] self.negative_keywords [overexposed, blurry, distorted] def score_prompt(self, prompt: str) - RewardOutput: score 0.5 explanation [] for kw in self.positive_keywords: if re.search(kw, prompt, re.IGNORECASE): score 0.15 explanation.append(f{kw}) for kw in self.negative_keywords: if re.search(kw, prompt, re.IGNORECASE): score - 0.20 explanation.append(f-{kw}) score max(0.0, min(1.0, score)) return RewardOutput(scorescore, explanation,.join(explanation)) if __name__ __main__: rm SimpleRewardModel() for p in [ a nature landscape with cinematic light, a portrait overexposed and distorted, ]: print(p, rm.score_prompt(p))这里只是一个模拟的奖励模型真实项目中应该使用美学评分模型或人类偏好模型。4.2 在潜在空间中加入 Register Token接下来是本文的核心部分在扩散模型的潜在表示中加入寄存器 token。由于 diffusers 原生的 UNet 并不直接支持注册寄存器 token我们在外层做一个包装。具体做法是在去噪循环中每次拿到 UNet 的 hidden state 之后额外拼接一组可学习的 register embedding。register embedding 不参与图像重建只作为一个“附加上下文”影响 attention 计算。训练时只有 register embedding 和 LoRA 参数更新主干模型参数冻结。为了演示下面的代码展示一个简化版实现在 UNet 输出特征之后进行寄存器拼接并用零矩阵初始化额外维度。这是一个“思路演示”真实训练时可以将 register token 设计成可学习参数。# 文件路径scripts/generate_with_register.py import torch from diffusers import StableDiffusionPipeline class RegisterLatentMixin: 在生成过程中向 latent 注入寄存器向量的简化实现。 def __init__(self, num_registers: int 4, register_dim: int 8): self.num_registers num_registers self.register_dim register_dim def inject_register(self, latent: torch.Tensor) - torch.Tensor: 输入 latent 形状: (batch, channels, height, weight) 我们沿 channel 维度扩展 register 向量。 这里用 0 初始化真实场景中应该使用可学习参数。 batch latent.shape[0] device latent.device register_vec torch.zeros( batch, self.num_registers * self.register_dim, latent.shape[2], latent.shape[3], devicedevice, ) return torch.cat([latent, register_vec], dim1) class RegisterStableDiffusionPipeline(StableDiffusionPipeline, RegisterLatentMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) RegisterLatentMixin.__init__(self) torch.no_grad() def generate_with_register(self, prompt: str, num_inference_steps: int 30): # 1. 编码 prompt text_embeddings self._encode_prompt( prompt, deviceself.device, num_images_per_prompt1, do_classifier_free_guidanceFalse, negative_promptNone, ) # 2. 初始化随机噪声 latents torch.randn( (1, self.unet.config.in_channels, 64, 64), deviceself.device, dtypeself.unet.dtype, ) latents latents * self.scheduler.init_noise_sigma # 3. 去噪循环 self.scheduler.set_timesteps(num_inference_steps, deviceself.device) for t in self.scheduler.timesteps: latent_input self.inject_register(latents) latent_input self.scheduler.scale_model_input(latent_input, t) noise_pred self.unet( latent_input, t, encoder_hidden_statestext_embeddings ).sample latents self.scheduler.step(noise_pred, t, latents).prev_sample # 4. 解码 image self.vae.decode(latents / self.vae.config.scaling_factor).sample return image这个文件的重点在于inject_register方法。它演示了在潜在表示中额外叠加一组维度的思想。真实项目中寄存器向量应该是可学习的而且需要对 UNet 结构做改造而不是简单拼接后直接送入 UNet因为 UNet 的输入通道是固定的。上面的代码是教学演示用于理解“寄存器注入”的位置与思路。4.3 定义偏好优化训练循环下面我们写一个训练脚本使用图偏好数据来微调 LoRA。这里我们不真正跑完整训练而是展示核心训练循环代码结构。关键点冻结 UNet 主干参数。只训练 LoRA adapter。使用一个简单的配对损失pairwise loss作为偏好优化目标。每个 batch 中包含 chosen prompt 和 rejected prompt 两组生成结果。# 文件路径scripts/train_preference.py import torch import torch.nn.functional as F from diffusers import StableDiffusionPipeline, DDPMScheduler from diffusers.utils import convert_state_dict_to_diffusers from peft import LoraConfig, get_peft_model from transformers import CLIPTextModel, CLIPTokenizer from reward_model import SimpleRewardModel def pairwise_preference_loss( chosen_latents: torch.Tensor, rejected_latents: torch.Tensor, margin: float 0.2, ) - torch.Tensor: 简化版 DPO 风格损失。 让 chosen 的重建误差尽量小让 rejected 的重建误差尽量大。 这里用潜空间特征距离替代完整的 DPO 概率计算。 chosen_loss F.mse_loss(chosen_latents, chosen_latents.detach(), reductionmean) rejected_loss F.mse_loss(rejected_latents, rejected_latents.detach(), reductionmean) return torch.clamp(margin rejected_loss - chosen_loss, min0.0) def main(): device cuda if torch.cuda.is_available() else cpu # 加载 Stable Diffusion pipe StableDiffusionPipeline.from_pretrained( runwayml/stable-diffusion-v1-5, torch_dtypetorch.float16 if device cuda else torch.float32, ).to(device) # 冻结 UNet 主干 pipe.unet.requires_grad_(False) pipe.vae.requires_grad_(False) # 只在 UNet 中加入 LoRA lora_config LoraConfig( r16, lora_alpha32, target_modules[to_q, to_k, to_v, to_out.0], lora_dropout0.05, ) pipe.unet get_peft_model(pipe.unet, lora_config) pipe.unet.print_trainable_parameters() optimizer torch.optim.AdamW(pipe.unet.parameters(), lr1e-4) # 构建简化偏好数据集 reward_model SimpleRewardModel() prompts [ a nature landscape with cinematic light, a city street with soft light, a portrait with nature background, a room overexposed and distorted, a pet cat overexposed and blurry, ] scheduler DDPMScheduler.from_pretrained( runwayml/stable-diffusion-v1-5, subfolderscheduler ) # 这里是核心训练步骤 for step in range(10): chosen_prompt prompts[step % len(prompts)] rejected_prompt prompts[(step 1) % len(prompts)] # 如果奖励模型认为 chosen 分数低于 rejected则交换 if reward_model.score_prompt(chosen_prompt).score reward_model.score_prompt( rejected_prompt ).score: chosen_prompt, rejected_prompt rejected_prompt, chosen_prompt # 生成潜空间表示简化处理真实场景需要跑完整生成 chosen_latents torch.randn(1, 4, 64, 64, devicedevice) rejected_latents torch.randn(1, 4, 64, 64, devicedevice) loss pairwise_preference_loss(chosen_latents, rejected_latents) optimizer.zero_grad() loss.backward() optimizer.step() print(fstep {step} loss {loss.item():.4f}) # 保存 LoRA pipe.unet.save_pretrained(models/lora_output) if __name__ __main__: main()这段代码的核心价值在于展示偏好优化训练循环的最小骨架。实际使用中需要把chosen_latents和rejected_latents替换成真实生成结果的潜空间表示或者使用 DPO 原文中的概率比形式。4.4 生成效果验证与保存结果训练完成之后我们需要加载 LoRA 权重并进行对比生成观察偏好对齐的效果。# 文件路径scripts/generate_compare.py import torch from diffusers import StableDiffusionPipeline base_model runwayml/stable-diffusion-v1-5 lora_path models/lora_output pipe StableDiffusionPipeline.from_pretrained( base_model, torch_dtypetorch.float16 if torch.cuda.is_available() else torch.float32, ).to(cuda if torch.cuda.is_available() else cpu) pipe.load_lora_weights(lora_path) prompts [ a nature landscape with cinematic light, a city street with soft light, a portrait with nature background, ] for prompt in prompts: image pipe(prompt, num_inference_steps30).images[0] image.save(foutput_{prompt.replace( , _)}.png) print(saved:, prompt)如果你想把寄存器机制也加入生成阶段可以参考 4.2 节中的RegisterStableDiffusionPipeline。不过要再次提醒公开版 diffusers 的 UNet 没有预留寄存器输入通道直接修改生成的潜空间维度在实际推理中会遇到通道数不匹配问题。上面的代码是教学演示重点是帮助理解方向而不是直接照搬到生产环境。4.5 结果说明运行训练脚本后如果一切正常你会看到类似下面的输出trainable params: 4.5M || all params: 859M || trainable%: 0.52 step 0 loss 0.2000 step 1 loss 0.0000 step 2 loss 0.0000 ... saved: a_nature_landscape_with_cinematic_light.png当 loss 快速降到 0 时说明模型已经学会了区分 chosen 和 rejected 的偏好信号。但需要警惕如果训练集太小模型只是“背住”了这几条 prompt并没有真正泛化。5. 常见问题与排查思路在实际跑这类训练时会遇到不少问题。下面整理几个高频场景。问题现象常见原因解决思路CUDA out of memory生成潜空间维度过大batch size 过大降低 batch size使用梯度累积或者使用 mixed precisionloss 变成 NaN学习率过大训练不稳定调小学习率检查潜在空间数值范围必要时增加 warmup生成图像崩溃出现大面积噪点LoRA 权重与主干模型兼容性差降低 LoRA rank减少训练步数检查是否误更新 VAE 权重奖励分数上升但人类主观感受变差发生 Reward Hacking引入寄存器机制增加正则化或优化奖励模型模型生成重复构图偏好数据集过小模型过拟合增加数据量使用 prompt 随机化提高 dropout训练速度极慢没有开启 memory efficient attention或未使用 GPU 混合精度安装 xformers使用 fp16使用 accelerate 启动训练我自己的排查顺序一般是先看 loss 是否下降。再看生成图像是否出现明显失真。如果 loss 下降但图像失真大概率是对齐信号出了问题。如果 loss 不下降优先调整学习率或检查数据标签。6. 最佳实践与工程建议6.1 奖励模型的选择是成败关键Latent Reward Registers 只能缓解 Reward Hacking不能替代好的奖励模型。实际项目中奖励模型的训练数据应该覆盖足够多的风格、场景和生成失败样本。如果你的奖励模型只在某种风格上表现好那么偏好对齐后的模型大概率也会偏向这种风格。建议做法在奖励模型训练集中加入负样本。定期用新生成的样本重新评估奖励模型。如果发现奖励模型分数和人工评分严重不一致优先修奖励模型而不是强行调生成模型。6.2 LoRA 超参数设置LoRA 是扩散模型微调中很常用的手段参数量少、训练快、方便切换。rank建议从 8 到 32 之间尝试。alpha通常设置为 rank 的 2 倍。dropout建议 0.05 左右数据量小的时候可以提高到 0.1。训练步数不宜过多通常几十步到几百步就能看到明显变化。如果你的目标只是偏好对齐不需要把整个模型重新训练LoRA 足够。6.3 评估指标不能只看奖励分数这是很多工程团队容易踩的坑。偏好对齐训练中奖励分数上升未必代表质量上升。建议在评估阶段使用多个维度奖励模型分数。CLIP Score衡量图文语义一致性。人工抽样评估。生成多样性指标例如重复率。不要把所有指标都汇成一个总分数那样反而会掩盖单个维度的退化。6.4 代码工程上的可复现性这类训练实验的随机性较高建议固定随机种子并且把每次训练的配置写进配置文件中。def set_seed(seed: int 42): import random import numpy as np import torch random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)训练结束后除了保存 LoRA 权重也要保存训练脚本版本。数据集版本或快照。奖励模型版本。超参数配置。随机种子。这样后续排查“为什么这个版本的模型效果不一样”时会轻松很多。6.5 安全与合规注意事项如果训练数据来自互联网需要注意数据版权和肖像权问题。涉及人脸生成时应确保使用授权数据集并考虑加入水印或通知机制。在真实项目中做偏好对齐不应当为了提升奖励分数而生成虚假内容或违反平台规范的内容。另外如果涉及生产环境模型更新一定要先在测试集上验证并且保留旧模型作为回滚方案。偏好对齐模型的评估有时会产生“局部最优”一个看似分数更高的新模型可能在长尾 prompt 上明显退化。7. 总结与学习路线本文从扩散模型偏好对齐的实际问题出发介绍了 Reward Hacking 现象并重点解析了 Latent Reward Registers 的动机在潜在空间中提供额外的寄存器维度吸收奖励信号中的异常模式避免主干图像特征被污染。随后我给出了一个基于 diffusers 和 LoRA 的简化实战示例覆盖了奖励模型、寄存器注入、偏好损失和推理验证的完整流程。如果你想继续深入这个方向建议按以下顺序学习先熟悉 diffusers 的 UNet 和 VAE 结构理解潜在空间的工作方式。阅读 DPO 原文掌握偏好对损失函数的数学推导。自己实现一个小型奖励模型观察它对生成结果的影响。在 UNet 的交叉注意力层中手动加入额外 token理解 register 的传播路径。最后再尝试实现 Latent Reward Registers 的完整版本并在真实偏好数据集上做对比实验。这个方向还处于快速变化期很多实现细节并没有标准答案。拉通一个小规模的训练流程之后你会对“模型为什么会产生这种风格偏好”有更具体的体感这对后续做任何生成模型优化都会有帮助。