One-Step Flow for Image Super-Resolution with Tunable Fidelity-Realism Trade-offs

OFTSR 输入拼接示意与 ImageNet 上的 FID-PSNR 对比气泡图
Fig. 1 — (a) 模型输入是「低清图 LR」与其「加噪版本 Augmented LR」沿通道拼接，单步模型一次前向就能输出高清图；调整插值参数 t 即可在 t=0（高保真，PSNR 25.9）和 t=1（高真实感，LPIPS 0.133）之间切换。(b) ImageNet 256×256 上各类扩散/flow 超分方法的 FID-PSNR 对比，气泡半径表示所需 NFE（函数求值次数）——OFTSR（橙色）用 1 次 NFE 就贴着帕累托前沿走完整条曲线。（论文 Fig. 1）

1. 出发点 (Motivation)

图像超分（super-resolution, SR）这几年被扩散 / flow 模型统治，因为它们生成的细节比传统回归式 CNN 更「真」。但这类方法卡在两个互相矛盾的需求之间：

这里的「保真度-真实感权衡」（perception-distortion tradeoff, Blau & Michaeli 2018）是图像复原的根本矛盾：数学上已证明不可能同时把「像 ground truth」（低失真、高 PSNR）和「看起来真实」（低 LPIPS/FID）都做到最好。医学影像、遥感、影视修复等不同场景对这两者的偏好截然不同——医学影像宁可保真也不要模型「脑补」，而影视修复要的就是细节真实。所以一个能在推理时调权衡的单步模型才有实用价值。

多步扩散其实天然能调这个权衡：调 NFE 就行（步数少→偏保真，步数多→偏真实感）。OFTSR 想问的问题是：能不能把多步模型里「沿采样轨迹滑动」的这个连续权衡能力，原封不动地压进一个单步网络，用一个超参 t 来调？

答案是两阶段方案：(1) 先训一个噪声增广的条件 rectified flow当教师；(2) 用一个特制蒸馏约束，把教师沿 ODE 轨迹的「整条权衡曲线」蒸进一步学生里。

2. 方法 (Method)

2.1 背景：rectified flow 是什么

Rectified flow（Liu et al. 2022）是一类基于 ODE 的生成模型。给定初始分布 \(p_0\) 和目标分布 \(p_1\)，它训练一个神经网络去拟合一个速度场 \(v_\theta\)：

\[\mathcal{L}_{rf}(\theta) := \mathbb{E}_{x_1\sim p_1, x_0\sim p_0}\left[\int_0^1 \|v_\theta(x_t, t) - (x_1 - x_0)\|^2\, dt\right], \quad x_t = (1-t)x_0 + t x_1.\]

—— 翻译：在 $x_0$（起点）和 $x_1$（终点）之间画一条直线，$x_t$ 是这条线上 $t$ 处的点。网络要学会在每个 $x_t$ 处预测「指向终点的方向和速度」，而这个真值就是直线的常数斜率 $x_1 - x_0$。训练好后，从 $x_0$ 出发解 ODE $\frac{dx_t}{dt} = v_\theta(x_t,t)$ 走到 $t=1$ 就生成了样本。

相比扩散模型，flow 的关键优势是：初始分布 \(p_0\) 不必是高斯。这就给超分留了空间——能不能直接学一条从「LR 图分布」到「HR 图分布」的 flow？

2.2 第一阶段：噪声增广的条件 flow（教师）

直接学 LR→HR 的 flow 会崩塌：训练时每个 LR 被强行拉向唯一的 HR，模型学成确定性映射，丢掉了「一个 LR 对应多个合理 HR」的多样性（论文 Tab. 7 验证了直接做效果差）。

OFTSR 的修法是给 LR 加噪，扩展初始分布的支撑集。用 Variance-Preserving（VP）方式加噪：

\[x_0 = \sqrt{1 - \sigma_p^2}\, x_{LR} + \sigma_p\, \epsilon, \quad \epsilon \sim \mathcal{N}(0, I).\]

—— 翻译：把 LR 图按 $\sqrt{1-\sigma_p^2}$ 缩一下，再掺进强度为 $\sigma_p$ 的高斯噪声当作 flow 的起点 $x_0$。$\sigma_p$ 越大，起点越「随机」，模型能脑补的多样性越大，但 LR 里的信息也丢得越多。VP 形式保证总方差不变。

这个 VP 公式在代码里是一行 GaussianNoise，与 Eq. (4) 一字不差：

repo/fm/measurements.py:L168-L174 — gaussian_VP 加噪算子，即论文 Eq. (4)

@register_noise(name='gaussian_VP')
class GaussianNoise(Noise):
    def __init__(self, sigma, **kwargs):
        self.sigma = sigma

    def forward(self, data):
        return np.sqrt(1 - self.sigma**2) * data + self.sigma * torch.randn_like(data, device=data.device)

加噪丢了信息怎么办？把原始 \(x_{LR}\) 沿通道维拼回去当条件输入。于是教师的训练目标是：

\[\mathcal{L}_{flow}(\theta) = \mathbb{E}_{x_1\sim p_1}\left[\int_0^1 D\big(v_\theta(x_{t,LR}, t),\, (x_1 - x_0)\big)\, dt\right], \quad x_{t,LR} = \text{concat}(x_t, x_{LR}).\]

—— 翻译：和普通 rectified flow 一样学速度场，只是把「噪声起点 $x_0$ → 高清图 $x_1$」当作要走的路径，并且每一步都把干净的 LR 图 $x_{LR}$ 拼在输入里当条件。$D$ 是 ℓ1 或 ℓ2 距离（论文发现 ℓ1 更好）。

这个 VP 公式很「统一」：\(\sigma_p = 0\) 退化成 InDI 的最小增广，\(\sigma_p = 1\) 就变成 SR3 的训练方式。代码里加噪 + 拼通道的逻辑在这里：

repo/fm/flow_matching.py:L100-L111 + L84-L99 — get_train_tuple_IR 加噪 + model_forward_wrapper 拼通道条件

def get_train_tuple_IR(self, batch, mask=None, **kwargs):
    x = batch['gt'].to(self.device)          # x1 = HR
    yn_ = batch['lq'].to(self.device)        # LR
    self.cond = yn_.detach().clone()         # 保留干净 LR 当条件
    yn = self.noiser_pertub(yn_)             # Eq.(4): VP 加噪得到 x0
    self.x1 = x
    self.x0 = yn

def model_forward_wrapper(self, model, x, t, **kwargs):
    # 把 x 与条件 LR 沿通道拼接后再喂模型
    x = torch.cat([x, self.cond], dim=1) if self.flow_config.use_cond else x
    model_output = model(x, t*999, label, augment_labels)
    return model_output[:, :3]

2.3 第二阶段：ODE 轨迹对齐蒸馏（核心创新）

OFTSR 蒸馏损失示意图
Fig. 2 — 蒸馏损失的核心思路。一步学生 $v_\phi$ 从同一个输入 $x_0$，在两个不同时刻 $t$ 和 $s$（$s>t$）各做一次预测，得到 $x_t$ 和 $x_s$。约束是：用教师 $v_\theta$ 从 $x_t$ 走 $(s-t)$ 步应该正好到达 $x_s$——也就是「学生的两个预测必须落在教师的同一条 PF-ODE 轨迹上」。（论文 Fig. 2，图中省略了 LR 条件）

为什么这样设计能保住「整条权衡曲线」？关键观察（Fig. 6 与前人工作）：沿 ODE 轨迹，从中间态 \(x_t\) 一步估计终点 \(x_{t_1}\)，估计值本身就躺在一条保真-真实感曲线上——\(t\) 越大（越接近 1）细节越丰富、LPIPS 越低（真实感好）；\(t\) 越小（越接近 0）越糊但 PSNR 越高（保真好）。所以只要让学生忠实复现教师的整条轨迹，学生就自动继承了用 \(t\) 调权衡的能力。

形式化：给定同一输入 \(x_{0,LR}\)，对 \(s > t\)，要求学生产生的 \(x_t, x_s\) 落在教师定义的同一 ODE 轨迹上：

\[x_s = x_t + (s-t)\, v_\theta(x_{t,LR}, t). \tag{6}\]

—— 翻译：教师视角下，从轨迹上的 $x_t$ 出发，沿教师速度 $v_\theta$ 走 $(s-t)$ 这么长的一步，应该到达 $x_s$。这就是一阶 Euler 离散的 PF-ODE。

而 \(x_t, x_s\) 都由一步学生算出（注意学生永远从 \(x_0\) 出发，只是把目标时刻当输入）：

\[x_t = x_0 + t\, v_\phi(x_{0,LR}, t). \tag{7}\]

—— 翻译：学生不沿轨迹爬，而是「一步到位」——从固定起点 $x_0$ 直接预测时刻 $t$ 处的状态 $x_t$。这正是单步模型的本质：把任意目标时刻 $t$ 当成可调旋钮。

把 (7) 代入 (6)，整理出对学生的约束（这是代码里实际优化的 v 空间形式）：

\[s\big(v_\phi(x_{0,LR}, s) - v_\phi(x_{0,LR}, t)\big) = (s-t)\big(v_\theta(x_{t,LR}, t) - v_\phi(x_{0,LR}, t)\big). \tag{8}\]

—— 翻译：把「学生两个时刻的速度差」和「教师速度与学生速度的差」用 $(s-t)/s$ 这个比例系数挂钩。它是 Eq.(6) 轨迹约束的等价改写，只是写成速度（$v$）而非状态（$x$）的关系，数值上更稳。

设 \(dt = s - t\)，加上 stop-gradient，得到最终蒸馏损失：

\[\mathcal{L}_{distill}(\phi) = \mathbb{E}_{x_1, t}\left[\left\| v_\phi(x_{0,LR}, s) - \text{SG}\Big[v_\phi(x_{0,LR}, t) + \tfrac{dt}{s}\big(v_\theta(x_{t,LR}, t) - v_\phi(x_{0,LR}, t)\big)\Big]\right\|^2_2\right]. \tag{9}\]

—— 翻译：让学生在时刻 $s$ 的速度，去匹配一个「目标速度」——这个目标速度是「学生在 $t$ 的速度」朝「教师速度」方向修正 $dt/s$ 比例后的结果。SG[·] 是停梯度，只让 $s$ 这一支回传梯度，防止训练塌掉。$t>0$ 保证 $dt/s$ 不会除零。

代码里这正是默认的 v_boot 分支，与 Eq. (9) 完全对应（delta_lambda = dt/s）：

repo/distill_fm.py:L390-L423 — 学生双时刻预测 + 教师 RK2 求速度 + v_boot 蒸馏损失（Eq. 8/9）

v_pred_t = flow.model_forward_wrapper(flow.model, x_0, t)   # 学生在 t 的速度
v_pred_s = flow.model_forward_wrapper(flow.model, x_0, s)   # 学生在 s 的速度
xt = x_0 + torch.einsum('b,bijk->bijk', t, v_pred_t.detach())  # Eq.(7): 学生算出 x_t

with torch.no_grad():                                       # 教师不回传梯度
    v_teacher_t = flow.model_forward_wrapper(flow.teacher, xt, t)
    if boot_config.teacher_solver == 'rk2':
        # r=0.5 -> midpoint；r=1 -> Heun；r=2/3 -> Ralston
        pred = v_teacher_t.clone()
        x_2 = xt + boot_config.rk2_r * torch.einsum('b,bijk->bijk', t_step, pred)
        t_2 = boot_config.rk2_r * s + (1 - boot_config.rk2_r) * t
        pred_2 = flow.model_forward_wrapper(flow.teacher, x_2, t_2)
        v_teacher_t = (pred + 1/(2*boot_config.rk2_r) * (pred_2 - pred))
# v_boot: Eq.(9)，delta_lambda = dt/s
delta_lambda = (t_step / s)[..., None, None, None]
loss_distil = F.mse_loss(v_pred_s,
                         (v_pred_t + delta_lambda*(v_teacher_t - v_pred_t)).detach(),
                         reduction='mean')

2.4 对齐损失与边界损失

光有蒸馏损失还不够稳。OFTSR 再加两项（借鉴 BOOT 的边界条件思想）：

\[\mathcal{L}_{align}(\phi) = \mathbb{E}_{x_1, t}\left[(1-t)^2\,\big\| v_\phi(x_{0,LR}, t) - v_\theta(x_{t,LR}, t)\big\|^2_2\right]. \tag{10}\]

—— 翻译：学生（从起点看 $t$）和教师（从轨迹上的 $x_t$ 看 $t$）都该指向同一个终点 $x_1$；把两者速度拉近，权重 $(1-t)^2$ 是因为换算到终点空间时会乘这个系数。$t=0$ 时它退化成 BOOT 的边界损失。

总目标：\(\mathcal{L}(\phi) = \mathcal{L}_{distill} + \lambda_{align}\mathcal{L}_{align} + \lambda_{BC}\mathcal{L}_{BC}\)，默认 \(\lambda_{align}=0.01\)、\(\lambda_{BC}=0.1\)。

2.5 推理：一次前向 + 一个旋钮

训练完，单步学生 \(v_\phi\) 一次前向就出 HR，权衡由 \(t\) 控制：

\[x_{t_1} = x_0 + v_\phi(x_0, x_{LR}, t). \tag{13}\]

—— 翻译：从加噪起点 $x_0$ 出发，加上学生预测的「残差速度」，直接得到高清图。想保真就把 $t$ 调小（输出偏向 LR 的均值，PSNR 高），想真实感就把 $t$ 调大（细节丰富，LPIPS 低）。整个推理 NFE = 1。

repo/fm/flow_matching.py:L232-L244 — 单步推理，对应 Eq. (13)

def image_restoration_t_(self, x_0, yn, t):
    samples = []
    for i in range(int(np.ceil(x_0.shape[0] / self.psnr_batch_size))):
        batch_x0 = x_0[i*self.psnr_batch_size:(i+1)*self.psnr_batch_size]
        vec_t = torch.ones(batch_x0.shape[0],).to(self.device) * t   # 旋钮 t
        self.cond = yn.detach()[...]
        with torch.no_grad():
            v_pred = self.model_forward_wrapper(self.model, batch_x0, vec_t)
            sample = batch_x0 + (self.T - self.eps) * v_pred         # x_0 + (T-t)*v
        samples.append(sample.cpu())
    return torch.cat(samples, dim=0)
OFTSR 在 t 从 1 到 0 时输出的连续过渡
Fig. 3 — 同一个单步模型，仅调 t：从 t=1（真实感最高，LPIPS 0.055 / PSNR 27.66）连续滑到 t=0（保真最高，LPIPS 0.160 / PSNR 30.03）。注意 PSNR 单调升、LPIPS 单调降——这正是把多步扩散「调 NFE 换权衡」的能力压进了单步推理。（论文 Fig. 3）

3. 结论 (Key findings)

在 FFHQ 256×256、DIV2K、ImageNet 256×256 和多个真实世界 SR 数据集上（4× 超分）：

4. 实现细节 (Implementation notes)

5. 批判性总结 (Critical assessment)

Strengths

Limitations / open questions

When to use / not use

Further reading

讨论 / Comments

评论托管在本仓库的 GitHub Discussions, 需 GitHub 账号。