扩散模型在图像生成任务上表现突出，但高成本的多步采样一直限制实时落地。潜在一致性模型（Latent Consistency Model，LCM）把一致性映射引入潜空间，用极少采样步数获得接近原始扩散模型的图像质量。理解LCM的关键不在网络结构，而在训练方法：一致性蒸馏通过自一致性约束，让模型学会从任意噪声点直接跳到数据端。接下来从原理、损失设计、训练代码与推理部署四个层面展开。
一、潜在一致性模型LCM要解决什么问题
常规扩散模型依赖迭代去噪。DDPM、DDIM或DPM-Solver等虽然不断降低采样步数，但为了保证图像质量，实际部署时仍然需要十几步甚至几十步推理。逐步蒸馏方法需要把原始扩散模型压缩成较少步数的学生模型，通常分阶段训练，流程较长且容易在中间步上积累误差。LCM的核心差异在于：它不再逼学生去复现少步采样的中间状态，而是直接学习一个一致性函数，使模型具备从任意时间步直接恢复初始数据的能力。
一致性函数建立在扩散过程的概率流常微分方程之上。对于同一条从数据到噪声的轨迹，任意两个不同的噪声状态理论上都会映射回同一个初始数据点。如果定义模型为f(x_t, t)，它接收带噪数据x_t和时间t，输出初始数据x_0，那么一致性性质要求f(x_t, t)对同一轨迹上的不同点保持不变。边界条件是f(x_0, 0) = x_0，即当输入已经无噪声时，模型只需要原样返回。这个性质一旦学到位，推理时从任意噪声点调用一次模型就能得到生成结果，而无需逐步迭代。
LCM进一步把一致性函数放到潜空间。标准潜在扩散模型会先使用VAE将图像编码为潜变量，再对潜变量做扩散与去噪，最后用解码器还原图像。LCM的训练和推理都基于编码后的潜变量z，这大幅减少了特征维度，降低显存和算力成本。更重要的是，LCM可以直接复用Stable Diffusion等预训练潜在扩散模型的UNet权重作为初始化，从而在较短时间内完成蒸馏。
二、一致性蒸馏的训练目标与跳步机制
一致性蒸馏无法用显式的z_0标签训练全部时间步，因为无法遍历所有可能的噪声轨迹。实际做法是：从训练数据中取出真实潜变量z_0，按照噪声调度加噪到较晚的时间步t_{n+k}，得到z_{t_{n+k}}。随后利用预训练扩散模型作为单步求解器，从z_{t_{n+k}}估计出更早的时间步z_{t_n}。训练学生模型f_θ时，要求它在两个时间点上的输出尽量一致。
损失函数通常写成最小化d(f_θ(z_{t_{n+k}}, t_{n+k}), f_{θ^-}(z_hat_t_n, t_n))。其中z_hat_t_n是预训练扩散模型执行一次去噪后得到的估计值，θ^-表示学生模型权重的指数移动平均版本。距离函数d可以选择均方误差、Huber损失或感知损失。目标模型θ^-不参与梯度更新，只用来提供稳定的回归目标，避免训练过程中的剧烈振荡。同时，还需要在t=0附近加入边界约束，让模型在输入无噪声时尽量保持恒等映射。
时间步离散化和跳步大小k对训练效果影响很大。前期如果k太大，两个时间点距离过远，一致性损失难以优化；如果k太小，学生模型只会处理接近无噪声的状态，单步生成能力不足。常见的做法是采用课程策略，训练初期使用较小的跳步距离，让模型先学会局部一致性，再逐步拉大时间间隔，增强全局映射能力。最终模型能够在四步甚至一步内完成采样，而不会出现明显的结构崩坏。
三、LCM训练的核心代码实现
下面的代码展示了一个LCM包装器和一致性损失的基础实现。实际系统中需要接入预训练UNet、时间嵌入以及VAE编码器，这里为了保持可读性做了简化。学生模型通常复用Stable Diffusion的UNet权重，目标模型使用同一结构但不求梯度，并通过EMA平滑更新。
import copy
import torch
import torch.nn as nn
class LCMWrapper(nn.Module):
def __init__(self, unet, ema_decay=0.999):
super().__init__()
self.online_model = unet
self.target_model = copy.deepcopy(unet)
for p in self.target_model.parameters():
p.requires_grad_(False)
self.ema_decay = ema_decay
def forward(self, z, t):
# 假设 unet 接受潜变量 z 和归一化时间 t
return self.online_model(z, t)
def update_target(self):
with torch.no_grad():
for online_param, target_param in zip(
self.online_model.parameters(),
self.target_model.parameters()
):
target_param.data.mul_(self.ema_decay)
target_param.data.add_(
online_param.data, alpha=1.0 - self.ema_decay
)
def consistency_loss(online, target, z_tnk, z_tn_hat, t_nk, t_n, criterion):
# 训练时对两个时间点的输出做一致性约束
pred_tnk = online(z_tnk, t_nk)
with torch.no_grad():
target_tn = target(z_tn_hat, t_n)
return criterion(pred_tnk, target_tn)
训练循环中需要从潜变量构造带噪样本，并调用预训练扩散模型执行一次单步去噪。跳步k可以根据训练轮次动态调整，以配合课程策略。损失只通过学生模型的反向传播更新，目标模型保持停止梯度。
for z0 in loader:
# z0 为编码后的潜变量，形状 [B, C, H, W]
eps = torch.randn_like(z0)
n = torch.randint(0, N, (1,)).item()
k = schedule_k(epoch) # 训练进程控制跳步大小
noise_scale_nk = (1.0 - alpha_bar[n + k]).sqrt()
z_tnk = alpha_bar[n + k].sqrt() * z0 + noise_scale_nk * eps
with torch.no_grad():
# 使用预训练扩散模型从 t_{n+k} 单步得到近似 z_{t_n}
z_tn_hat = pretrained_ode_step(z_tnk, t_nk, t_n)
loss = consistency_loss(
lcm_online,
lcm_target,
z_tnk,
z_tn_hat,
t_nk,
t_n,
huber_criterion
)
optimizer.zero_grad()
loss.backward()
optimizer.step()
lcm_online.update_target()
推理时，LCM不需要像传统扩散模型那样反复调用UNet。给定随机潜变量，模型可以直接输出对初始潜变量的预测，再交给VAE解码器生成图像。如果想使用多步采样提高质量，可以在每步预测后根据当前时间重新加噪，再进行下一轮预测。下面的代码展示了这种多步潜一致性采样流程，其中比较逻辑使用不等于判断，避免额外符号处理。
def lcm_sampling(z, lcm_model, steps=4, t_start=1.0, t_end=0.0):
# 多步潜一致性采样：每一步预测 z0 后按当前时间重新加噪
ts = torch.linspace(t_start, t_end, steps + 1)
for i in range(steps):
t = ts[i]
z0_pred = lcm_model(z, t)
if i != steps - 1:
t_next = ts[i + 1]
alpha = noise_schedule(t_next)
noise = torch.randn_like(z)
z = alpha.sqrt() * z0_pred + (1.0 - alpha).sqrt() * noise
else:
z = z0_pred
return z
四、推理加速与工程实践建议
LCM推理步数可以配置为1到8步。1步模式适合对延迟要求极高的场景，但细节可能略弱；4步模式通常能在质量和速度之间取得较好平衡。多步采样并不是简单重复模型输出，而是在每轮预测后按当前时间重新加入适量噪声，让模型有机会修正前一步的偏差。这种机制比直接多次调用模型要稳定得多，也能进一步改善边缘和纹理表现。
与无分类器引导结合时需要注意，LCM对引导强度比较敏感。蒸馏阶段如果固定了引导尺度，推理时不宜大幅偏离。一般可以尝试1.5到3.0的低引导值，否则容易出现过曝或色彩饱和失真。对于条件生成任务，可以在训练阶段把文本嵌入一并送入UNet，保持与原始Stable Diffusion一致的条件输入结构。
工程部署时，LCM常以LoRA方式对现有模型做少量参数微调。这样既能保留基础模型的生成能力，又能降低存储和传输成本。推理加速还可以结合FP16精度、TensorRT编译或INT8量化。同时不要忽略VAE解码器在总延迟中的占比，若端到端时延要求非常苛刻，可以选择更轻量的VAE或提前缓存文本和噪声调度相关张量。