论文阅读笔记：《TSD-SR: One-Step Diffusion with Target Score Distillation for Real-World Image Super-Resolution》

## 论文核心思想

本文提出了一种名为 TSD-SR 的新颖框架，旨在将强大的预训练文本到图像（T2I）扩散模型蒸馏成一个高效且有效的单步真实世界图像超分辨率（Real-ISR）模型。其核心思想是通过新颖的 目标分数蒸馏（TSD） 和 分布感知采样模块（DASM），解决现有蒸馏方法在应用于超分辨率任务时的性能瓶颈，从而在实现极快推理速度的同时，达到甚至超越许多复杂多步模型的图像恢复质量。

模型解决了什么问题？

该模型主要针对现有基于扩散模型的超分辨率方法中存在的两大痛点：

多步模型的效率低下：大多数利用扩散模型先验的Real-ISR方法，由于其固有的迭代去噪特性，需要20-50步的推理过程，导致计算成本高昂，速度远不及GAN等单步方法。
现有单步蒸馏方法的性能局限：虽然已有工作尝试将扩散模型蒸馏为单步或少步，但它们在图像恢复和细节生成方面的表现并不尽如人意。本文作者发现，这主要是因为直接套用变分分数蒸馏（VSD） 到Real-ISR任务中存在两个关键缺陷：
- 不可靠的梯度方向：VSD依赖教师模型提供优化方向，但在训练初期，当学生模型生成的图像质量较差时，教师模型提供的梯度方向并不可靠，这会导致生成结果出现视觉伪影。
- 细节恢复能力不足：VSD损失在不同时间步（timestep）上的变化很大。传统的VSD框架采用均匀采样策略来选择时间步，这会导致对细节恢复至关重要的早期时间步的梯度被稀释，使得模型难以学习恢复精细纹理。

## 方法 (Methodology)

为了解决上述问题，TSD-SR引入了两个核心组件：TSD Loss 和 DASM。

1. TSD Loss (目标分数蒸馏损失)

TSD Loss是一种为学生模型提供稳定可靠梯度的组合损失函数。其目的是在VSD的基础上，引入一个更可靠的监督信号来指导优化。

其梯度公式如下：

这个损失由两部分构成：

目标分数匹配 (TSM):ϵ_ψ(ẑ_t; t, c_y) − ϵ_ψ(z_t; t, c_y)
- 作用：该项强制让教师模型 (ϵ_ψ) 对“学生生成的含噪潜变量” ẑ_t 和“真实高清图像的含噪潜变量” z_t 的预测尽可能保持一致。这相当于提供了一个来自真实数据分布的“锚点”，为优化提供了更可靠的方向，有效避免了伪影和过平滑问题。
变分分数蒸馏 (VSD): λ(ϵ_ψ(z_t; t, c_y) − ϵ_ϕ(ẑ_t; t, c_y))
- 作用：该项是经典VSD的变体，旨在将教师模型 ϵ_ψ 的知识蒸馏给可训练的LoRA模型 ϵ_ϕ ，促使学生模型的分布向教师模型看齐。

2. DASM (分布感知采样模块)

DASM是一个为解决细节恢复不足问题而设计的采样模块。

动机：它旨在克服传统均匀采样策略带来的梯度稀释问题。

结构与流程：
- DASM首先像常规方法一样，通过加噪得到一个初始的含噪样本 ẑ_t。
- 接着，它不再止步于此，而是利用LoRA模型 ϵ_ϕ 进行多步迭代式的“微型去噪”，从而生成一系列沿着真实扩散采样轨迹的含噪样本 ẑ_t − 1, ẑ_t − 2, … 。
- 在一次训练迭代中，所有这些由DASM生成的样本都会被用来计算TSD损失，并将它们的梯度累积起来，共同更新学生模型。
优势：通过这种方式，DASM在一次迭代中就强化了对细节恢复至关重要的早期时间步的学习，使得梯度能够更集中地用于优化细节生成，有效提升了图像的清晰度和真实感。

训练与推理 (Training and Inference)

推理策略 (Inference Strategy)

该模型的推理策略非常高效和直接，其核心是实现“单步”超分辨率。

单步完成: TSD-SR 经过蒸馏训练后，成为了一个单步推理模型。这意味着它不像传统扩散模型那样需要20-50次迭代去噪，而是仅需一步即可完成从低质量（LQ）图像到高质量（HQ）图像的转换。
流程:
1. 输入一张低质量（LQ）图像。
2. 训练好的学生模型 (Gθ) 对该图像进行处理。
3. 模型在潜空间中执行一次性的去噪和超分辨率操作。
4. 输出最终的高质量（HQ）图像。
高效性: 这种单步策略使得其推理速度极快。例如，它的速度比SeeSR快40倍以上，比SUPIR快120倍以上。其速度优势源于直接从低质量图像数据去噪并采用固定的文本提示（prompt）。

训练策略 (Training Strategy)

该模型的训练策略相对复杂，其目标是将一个强大的预训练多步扩散模型（教师模型）的知识蒸馏到一个高效的单步学生模型中。整个过程可以参考论文第17页的算法，主要包含以下几个关键部分：

1. 初始化

学生模型 (Gθ): 学生模型的主体结构和权重初始化自预训练的教师模型（论文中为SD3）。其中，VAE编码器（带有可训练的LoRA层）和去噪网络是可训练的，而VAE解码器则被冻结，以保留其强大的图像生成先验知识。
LoRA模型 (ϵϕ): 这是一个可训练的教师模型副本（通过LoRA实现），它的权重同样初始化自教师模型的去噪网络。

2. 核心训练循环

在每次训练迭代中，模型会交替更新“学生模型”和“LoRA模型”。

A. 学生模型 (Gθ) 的更新

学生模型的更新依赖于一个组合损失函数，该函数包含两大部分：

重建损失 (L_Rec): 用于保证生成图像在内容上与高质量原图相似。
- 它由像素空间的LPIPS损失和潜空间的MSE损失（均方误差损失）组成。
- 训练初期会使用MSE损失以稳定训练，后期则会移除它以避免结果过于平滑。
正则化损失 (L_Reg): 即本文核心的 TSD损失，用于提升图像的真实感和细节。
1. 初始采样与TSD计算: 首先，对学生模型生成的潜变量ẑ和真实高质量潜变量z添加噪声，得到ẑ_t和z_t 。然后，利用这两个含噪潜变量计算一次TSD损失（如论文公式(5)所示）。
2. DASM梯度累积: 接下来，启动DASM模块。它会沿着采样轨迹进行N次迭代去噪，每一步都生成新的含噪样本对（例如ẑ_cur, z_cur）。对每一对新样本都计算一次TSD损失，并将其加权累加到总的正则化损失中。
3. 这一步是DASM的核心，通过累积多个相关样本的梯度，来强化对细节恢复的学习。
最终更新: 将重建损失和正则化损失加权求和，得到学生模型的总损失ℒ_G 。使用这个总损失通过反向传播来更新学生模型的参数θ 。

B. LoRA模型 (ϵϕ) 的更新

根据VSD框架的要求，可训练的LoRA模型也需要被更新。

扩散损失 (L_loss): 计算一个扩散损失（如论文公式(9)所示），目标是让LoRA模型ϵ_ϕ的去噪能力与教师模型对齐。这个损失衡量了LoRA模型预测的噪声与目标噪声之间的差距。
更新: 使用该扩散损失来更新LoRA模型的参数ϕ 。

通过在训练中不断重复这两个更新步骤，模型最终学会了在单步内生成高质量、细节丰富的超分辨率图像。

总结与个人看法

读完有几个疑问亟待解决： - 在训练初期教师模型对于ẑ₀的预测一开始其实是不准确的，从那一张余弦相似度的图片其实就能看出来，那么为什么采用了TSD-Loss之后就可以？ >核心思想: 该部分的核心思想是，对于来自同一分布的样本，教师模型 (ϵψ) 预测出的真实分数应该是相近的。它通过比对教师模型对“合成的含噪潜变量”ẑ_t和“高质量的含噪潜变量”z_t 的预测，来鼓励两者之间的预测结果保持一致。这提供了一个更可靠的优化方向，特别是当初始生成的图像质量不高时，可以有效避免模型陷入视觉伪影或生成过度平滑的结果。

其实这一点或许就是为了避免生成artifacts，或者理解，教师模型认不出来复原模型生成的结果。

但是为什么TSD的效果比OSEDiff更好，按照对于原文的理解，或许是因为他采用了一个更加牛逼的VSD-Loss，因为他用了DASM模块，计算了多个时间步的VSD对Loss进行贡献。究其原因是因为在不同的时间步，扩散模型解决的问题不一样，在越小的时间步也就是越接近ẑ₀时越会专注于细节的生成，然而OSEDiff在所有的时间步是均匀采样的，这就导致一个问题，当采用均匀采样时，模型在大多数情况下会采到属于Stage 1的“简单”时间步，只有少数机会采到属于Stage 2的“困难但重要”的时间步。这就导致了真正能指导模型学习细节的有效梯度，在大量“简单”梯度中被平均化，其作用被大大削弱，即“被稀释了”。

难道说这个就能用来解释OSEDiff结果的平滑性质了么？