在扩散模型中，RectifiedFlowScheduler.step() 方法的作用远不止简单的噪声相减，而是通过 数学优化路径 和 动态时间步控制 来实现高效、稳定的去噪过程。以下从原理、设计逻辑和代码实现三个层面详细解析：

一、核心问题：为什么不能直接相减噪声？

1. 传统扩散模型的局限性

• 在 DDPM 等传统扩散模型中，去噪公式为：
  $$x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}t}} \epsilon\theta \right) + \sigma_t z$$
  其中需依赖预设的噪声表（$\alpha_t, \beta_t$）和随机噪声 $z$ 。
• 问题：
  • 计算复杂，需维护噪声系数表；
  • 步数多（通常需 1000 步），效率低；
  • 随机性控制困难（如 DDIM 需切换确定性/随机模式）。

2. Rectified Flow 的革新思路

Rectified Flow 通过 直线路径优化 解决上述问题：

• 路径公式：$z_t = (1-t) \cdot \text{data} + t \cdot \text{noise}$，用直线连接数据与噪声分布。
• 去噪目标：预测速度场（velocity field）$v_\Theta$，而非直接预测噪声。
  去噪步骤简化为：
  $$z_{t-1} = z_t - \Delta t \cdot v_\Theta$$
  其中 $\Delta t$ 是动态计算的时间间隔。

二、RectifiedFlowScheduler.step() 的核心作用

1. 动态时间步计算

1
2
3
4

# 计算当前时间步与下一时间步的间隔 Δt
lower_mask = timesteps_padded < timestep - t_eps
lower_timestep = timesteps_padded[lower_mask][0]  # 找到最近的下一时间步
dt = timestep - lower_timestep  # 计算 Δt

• 意义：
  • 避免固定步长导致的误差累积；
  • 自适应选择最优时间间隔，提升数值稳定性。

2. 去噪执行（核心公式）

1

prev_sample = sample - dt * model_output  # 即 z_{t-1} = z_t - Δt * v_Θ

• 与直接减噪声的本质区别：
  • model_output 是速度场 $v_\Theta$，而非噪声 $\epsilon$；
  • 通过 $v_\Theta$ 沿直线路径反向求解 ODE，实现高效传输。

3. 随机性控制

1
2
3
4

if stochastic_sampling:  # 随机采样模式
    x0 = sample - timestep[..., None] * model_output  # 预测原始数据 x0
    next_timestep = timestep[..., None] - dt
    prev_sample = self.add_noise(x0, torch.randn_like(sample), next_timestep)

• 设计意图：
  • 默认关闭随机性（stochastic_sampling=False），确保结果可复现；
  • 开启后模拟传统扩散的随机过程，增加生成多样性。

三、条件掩码（Conditioning Mask）的作用

1
2

tokens_to_denoise_mask = (t - t_eps < (1.0 - conditioning_mask)).unsqueeze(-1)
return torch.where(tokens_to_denoise_mask, denoised_latents, latents)

• 功能解析：
  • conditioning_mask：标识每个潜空间位置的“条件强度”（1.0=完全条件，0.0=无条件）。
  • 动态去噪：仅对满足 $t < (1.0 - \text{mask})$ 的 token 执行去噪（即时间步早于预设条件步）。
• 应用场景：
  • 视频生成中，局部区域/帧 按不同时间步去噪（如首帧强条件，后续帧弱条件）；
  • 避免全局同步去噪导致的时序错位。

四、Rectified Flow 的架构优势

五、总结：调度器的不可替代性

1. 路径优化
   Rectified Flow 的直线路径设计取代了传统扩散的弯曲路径，显著减少误差累积，实现 “一步直达”式去噪。
2. 动态时间控制
   通过 $\Delta t$ 的自适应计算，避免固定步长的数值不稳定问题，尤其适合 高分辨率生成（如 4K 视频）。
3. 条件生成兼容性
   与条件掩码协同，支持 细粒度时空控制（如视频中特定帧的局部编辑），这是直接减噪声无法实现的。

因此，RectifiedFlowScheduler.step() 不仅是数学优化的核心，也是平衡 质量、速度与控制力 的关键模块，直接相减噪声无法达到同等效果。