Analyzing the Training Dynamics of Image Restoration Transformers: A Revisit to Layer Normalization

该论文揭示了传统层归一化在图像恢复 Transformer 训练中会导致特征幅值发散和通道熵坍塌的问题，并提出了一种名为 i-LN 的定制化层归一化方法，通过整体归一化和自适应重缩放有效解决了这一矛盾，显著提升了训练动态与恢复性能。

要点

这篇论文就像是在给图像修复（比如把模糊照片变清晰、把下雨天照片变晴天）的“大脑”（Transformer 模型）做了一次深度体检，发现了一个被大家忽略的“隐形杀手”，并发明了一个简单的“急救包”来解决问题。

我们可以用**“装修房子”和“指挥交通”**的比喻来理解这篇论文的核心内容：

1. 发现了什么问题？（“失控的装修队”）

想象一下，你雇佣了一个装修队（AI 模型）来修复一张破损的老照片。

• 传统的做法（LayerNorm）： 装修队里有个“监工”（LayerNorm，层归一化）。他的规矩是：不管这面墙原本是什么颜色，也不管隔壁墙是什么颜色，他都要强行把每一块砖（图像中的每一个像素点）的颜色都调整成“标准灰度”，并且把砖块的大小强行统一。
• 后果： 这个监工太死板了！
  • 破坏邻里关系： 照片里的细节（比如猫的眼睛和胡须）是紧密相连的。强行把每一块砖单独标准化，就像把猫的眼睛和胡须强行拆开，导致它们失去了原本的联系（空间相关性被破坏）。
  • 数值爆炸： 为了对抗这个死板的监工，装修队里的工人们（网络特征）开始“发疯”。他们为了绕过监工的规则，把砖块做得巨大无比，数值甚至膨胀到了一百万那么大！这就好比为了把墙刷白，工人们把油漆桶都炸开了，满屋子都是飞溅的油漆（特征发散）。
  • 信息丢失： 原本照片里丰富的色彩层次（熵），因为被强行统一，变得像白开水一样平淡无奇（通道熵崩塌）。

2. 为什么以前没发现？（“只盯着局部看”）

以前的研究只关注模型能不能把图修好，没去观察装修过程中工人们的状态。就像只盯着房子最后有没有修好，没发现装修队其实一直在“带病工作”，随时可能因为数值爆炸而崩溃。

3. 他们提出了什么新方法？（"i-LN：懂变通的智能监工”）

作者发明了一个新工具，叫 i-LN。它不是把旧监工开除，而是给他换了一套更聪明的“工作手册”。

• 策略一：不再“各自为政”，而是“统筹全局” (Spatial Holisticness)

  • 旧做法： 监工盯着每一块砖单独看，不管它旁边是谁。
  • 新做法 (i-LN)： 监工退后一步，把整面墙（甚至整个房间）看作一个整体。他计算整面墙的平均颜色和整体大小，然后统一调整。
  • 比喻： 就像指挥交通，以前是每辆车自己看红绿灯，现在变成了交警看着整个路口的车流，统一指挥。这样，原本紧密相连的“砖块”（像素）就能保持它们原本的空间关系，猫的眼睛和胡须依然紧紧挨着。

• 策略二：不再“一刀切”，而是“看人下菜碟” (Input-Adaptive Rescaling)

  • 旧做法： 不管输入的是晴天还是暴雨，监工都按同一套标准调整。
  • 新做法 (i-LN)： 监工会看一眼现在的天气（输入图像的特征）。如果是暴雨天，他就允许颜色深一点；如果是晴天，就允许亮一点。
  • 比喻： 就像裁缝量体裁衣。以前是大家都穿均码（统一归一化），现在是根据每个人的身材（输入统计信息）重新调整衣服的大小。这样既保留了衣服原本的版型，又穿得舒服。

4. 效果怎么样？（“从混乱到井井有条”）

用了这个新工具后，奇迹发生了：

• 数值稳定了： 那些疯狂膨胀到“一百万”的数值，现在乖乖地回到了正常范围（像 1 左右），不再炸锅。
• 细节更清晰： 因为保留了像素间的“邻里关系”，修复出来的照片边缘更锐利，纹理更清晰，不再有那种模糊或奇怪的伪影。
• 更省电（低精度推理）： 因为数值不再乱跳，这个模型甚至可以在手机或低端芯片上运行（使用低精度计算），而不会像以前那样因为数值太大导致手机死机或画面变黑。

总结

这篇论文告诉我们：在修复图像时，不要强行把每个像素都变成“标准件”。

以前的 AI 模型像是在用**“流水线作业”修图，虽然快但容易把细节修坏，甚至把自己逼疯（数值爆炸）。
作者提出的 i-LN 就像是给模型装上了“全局视野”和“灵活应变”**的能力。它告诉模型：“请尊重照片原本的结构，根据具体情况灵活调整。”

结果就是：模型修图修得更快、更稳、更漂亮，而且连手机都能跑得动了。这是一个简单却极其有效的“四两拨千斤”的改进。

技术摘要

这是一篇发表于 ICLR 2026 的会议论文，题为《分析图像恢复 Transformer 的训练动态：重新审视层归一化》（ANALYZING THE TRAINING DYNAMICS OF IMAGE RESTORATION TRANSFORMERS: A REVISIT TO LAYER NORMALIZATION）。

以下是对该论文的详细技术总结：

1. 研究背景与核心问题 (Problem)

背景：
图像恢复（Image Restoration, IR）任务（如超分辨率、去噪、去雨等）近年来广泛采用基于 Vision Transformer (ViT) 的架构，并通常结合层归一化（LayerNorm, LN）作为标准组件。

核心问题：
作者通过深入分析 IR Transformer 的训练动态，发现了一个关键但被忽视的现象：

• 特征发散（Feature Divergence）：在使用传统 LayerNorm 时，网络内部特征的量级（magnitude）会剧烈发散，达到百万级（$10^6$）的规模。
• 通道熵坍塌（Channel-wise Entropy Collapse）：特征在通道维度上的熵急剧下降，意味着激活值高度集中在少数几个通道中，导致特征分布极度不平衡。

原因分析：
作者认为，这种现象源于传统 LayerNorm 的约束与图像恢复任务的独特需求之间存在错位（Misalignment）：

1. 逐 Token 归一化破坏空间相关性：传统 LN 对每个 Token（像素块）独立进行归一化，忽略了 Token 之间的空间关系（inter-pixel relationships），破坏了图像恢复至关重要的空间结构。
2. 输入无关的缩放丢弃统计信息：传统 LN 将特征映射到统一的归一化空间，限制了内部表示的范围灵活性，忽略了 IR 任务中固有的输入依赖统计特性（input-dependent statistics）。
3. 网络的对抗行为：为了绕过 LN 施加的与其任务目标冲突的约束，网络被迫学习生成极大值的特征来“欺骗”归一化层，从而导致特征发散。

2. 方法论：i-LN (Methodology)

为了解决上述问题，作者提出了图像恢复 Transformer 专用层归一化（Image Restoration Transformer Tailored Layer Normalization, i-LN）。这是一种即插即用（drop-in replacement）的改进方案，包含两个核心组件：

A. 空间整体归一化 (Spatial Holistic Normalization, LN*)

• 机制：不同于传统 LN 对每个 Token 单独计算均值和方差，i-LN 在整个空间 - 通道维度（spatio-channel dimension）上计算全局统计量（均值 $\mu$ 和方差 $\sigma^2$）。
• 数学原理：
  $$LN^*(x) = \gamma \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$
  其中期望 $E_{\ell,c}$ 是同时对空间位置 $\ell$ 和通道 $c$ 进行的。
• 理论优势：证明了这种操作是一个位似变换（Homothety），能够保持 Token 之间的相对差异（即空间结构），直到一个全局缩放因子。这解决了传统 LN 破坏空间相关性的问题。

B. 输入自适应重缩放 (Input-Adaptive Rescaling)

• 机制：在注意力（Attention）和前馈网络（FFN）层之后，利用前一步归一化过程中计算出的标准差 $\sigma$ 对输出进行显式重缩放。
• 公式：
  $$B(x; f, i\text{-}LN) = x + \sqrt{\sigma^2 + \epsilon} \cdot f(LN^*(x))$$
• 目的：传统归一化会丢失全局缩放信息（即特征的范围灵活性）。通过引入 $\sqrt{\sigma^2 + \epsilon}$ 项，i-LN 显式地恢复了丢失的全局缩放项，允许网络根据输入的具体统计特性灵活调整特征范围，从而保留输入依赖的低层特征统计信息。

3. 主要贡献 (Key Contributions)

1. 现象发现与理论解释：首次系统性地揭示了 IR Transformer 中特征量级发散和通道熵坍塌的现象，并从理论上证明了传统逐 Token LayerNorm 无法保持图像的空间结构，导致网络产生病态行为。
2. 提出 i-LN：设计了一种简单但有效的归一化方案，结合了“空间整体归一化”和“输入自适应重缩放”，完美契合 IR 任务对空间相关性和特征范围灵活性的需求。
3. 广泛的实验验证：在多种 IR 任务（超分辨率 SR、去噪 DN、去雨 DR、JPEG 伪影去除 CAR）和多种骨干网络（SwinIR, HAT, DRCT, SRFormer）上进行了验证。
4. 鲁棒性分析：证明了 i-LN 不仅提升了性能，还显著增强了模型在低精度推理（如 INT8/FP16）下的稳定性，防止了因特征过大导致的数值溢出（NaN/Inf）。

4. 实验结果 (Results)

• 训练稳定性：
  • 使用传统 LN 时，特征量级随训练迭代迅速发散至 $10^6$ 甚至 $10^7$，且通道熵在训练早期急剧下降。
  • 使用 i-LN 后，特征量级被稳定控制在合理范围（接近 $N(0,1)$ 分布），通道熵保持高位且分布均匀，训练过程极其稳定。
• 性能提升：
  • **超分辨率 **(SR)：在 Set5, Set14, Urban100 等基准上，PSNR 和 SSIM 均有显著提升（例如 HAT1 + i-LN 在 Urban100 ×4 SR 上 PSNR 提升约 0.62 dB）。
  • 其他任务：在去雨、去噪和 JPEG 去块效应任务中，i-LN 均优于传统 LN 及其他归一化变体（如 LayerScale, RMSNorm, InstanceNorm 等）。
  • 消融实验：证明了“空间整体性”和“重缩放”两个组件缺一不可，共同作用才能达到最佳效果。
• 低精度推理鲁棒性：
  • 在 FP16 和 INT8 量化设置下，传统 LN 模型会出现大量黑点（数值溢出/Inf），性能崩溃。
  • i-LN 模型在低精度下几乎无损，保持了与全精度相当的恢复质量，证明了其数值稳定性。
• 空间相关性建模：
  • 可视化相对位置嵌入（RPE）显示，使用 i-LN 的模型学习到了结构清晰、类似卷积核模式的 RPE，而传统 LN 模型的 RPE 则充满噪声，表明 i-LN 更好地捕捉了像素间的空间关系。

5. 意义与影响 (Significance)

• 重新定义 IR 中的归一化：该工作挑战了 Transformer 中“逐 Token 归一化”作为默认标准的假设，指出在图像恢复等需要保留空间结构和低层统计信息的任务中，这种设计是有害的。
• 简单高效的改进：i-LN 不需要改变网络架构，只需替换归一化层即可带来显著收益，具有极高的实用价值和推广性。
• 解决数值稳定性难题：为 IR Transformer 在边缘设备上的低精度部署扫清了障碍，解决了因特征发散导致的量化失效问题。
• 理论指导实践：通过数学推导（位似变换）和实证分析，为理解 Transformer 在低级视觉任务中的训练动态提供了新的视角，即网络可能会为了适应不合适的归一化约束而“病态”地调整特征分布。

总结：
这篇论文通过深入分析训练动态，发现并解决了图像恢复 Transformer 中由传统 LayerNorm 引起的特征发散问题。提出的 i-LN 方法通过空间整体归一化和输入自适应重缩放，不仅显著提升了恢复性能，还极大地增强了训练的稳定性及低精度推理的鲁棒性，为未来 IR 模型的设计提供了重要的理论依据和工程指导。