实时 4K 图像修复：TSFormer 以 3.38M 参数达 40fps 高性能，推动技术落地

超高分辨率（UHD）图像修复对于需要出色视觉保真度的应用至关重要，然而现有方法往往在修复质量和效率之间存在权衡，限制了其实际部署。
在这篇论文中，作者提出了一种名为TSFormer的全功能框架，该框架将可信学习与S化简相结合，以提高 UHD 图像修复的泛化能力和计算效率。关键在于，模型中允许的 Token 移动量非常有限。
为了有效地过滤 Token ，作者使用基于随机矩阵理论的 Min-p 方法来量化 Token 的不确定性，从而提高模型的鲁棒性。作者的模型可以在实时（40fps）运行 4K 图像（3.38 M 参数）。
大量实验表明，TSFormer 在实现最先进的修复质量的同时，提高了泛化能力和降低了计算需求。
此外，作者的 Token 过滤方法可以应用于其他图像修复模型，以有效地加速推理并保持性能。

1 Introduction

超高分辨率（UHD）图像修复对于需要高分辨率图像质量的各种应用至关重要，包括医学影像、视频流媒体和数字监控。由于 UHD 图像拥有数百万像素，在资源有限的情况下进行处理是一项巨大的挑战。

目前，有一些深度学习框架[34, 39]可以在消费级 GPU 上处理 UHD 图像。然而，这些方法直接或间接地降低输入图像的分辨率，可能导致图像修复过程中丢失一些重要信息。因此，这会导致图像细节丢失，这对 UHD 图像尤为重要。

为了解决这个问题，作者提出了一种名为 TSFormer 的轻量级和可信框架，它结合稀疏化和随机矩阵理论。TSFormer 旨在保留最具信息量的特征，从而在降低计算开销的同时保持高修复质量和鲁棒性。

确实，TSFormer 的第一个关键组成部分是 Min-p 采样，这是一种基于概率稀疏化理论的稀疏表示技术，灵感来源于[1]中最近的发展。与传统的 Top-k 过滤方法[18]不同，Min-p采样根据概率阈值选择性地保留高置信度的特征，允许动态特征选择。这种"少就是多"的方法认识到，并非所有特征对最终输出做出相同程度的贡献，尤其是在高维数据（如 UHD 图像）中。然而，尽管 Min-p 采样有效降低了模型的计算需求，但在大量 Token 中可能引入不稳定性，因为存在噪声。为了解决这个问题，TSFormer 引入了一个基于随机矩阵理论的可信机制[7, 23]。可信过滤涉及分析特征矩阵的特征值，以确保仅保留稳健、高置信度的特征。如图 1(a)所示，Min-p 采样应用基于概率的阈值保留高置信度的特征，同时丢弃较低重要性的特征，实现有效的稀疏化。在图 1(b)中，基于随机矩阵理论的可信过滤进一步细化剩余特征，通过排除位于可信阈值（虚心圆）之外的点。这种可信驱动的特征选择改进了泛化能力，使模型在处理各种退化图像时能够可靠运行。

到目前为止，作者通过强制使用 Min-p 并引入可信机制构建了一个 Min-p 稀疏注意力（MSA），它能够高效、可靠地生成注意力图。MSA 被集成到 TSFormer 的每个块中，每个块还包含频域学习和多尺度学习组件。TSFormer 在多个 UHD 图像任务上表现出令人鼓舞的性能，并且可以在资源受限的设备（单个 3090 GPU 着色器，24G RAM）上实时运行 UHD 图像。

总之，作者的贡献如下：

作者开发了一种名为 TSFormer 的 UHD 图像恢复模型，它可以在单个 GPU 上实时运行 4K 分辨率图像，并具有强大的泛化能力。
作者设计了一种基于可信机制的 Token 过滤方法，将其集成到 TSFormer 中，以生成高质量的关注图。
基于可信机制的 Token 过滤方法可以在任何基于 Transformer 的图像恢复框架中使用，以提高模型的效率。大量实验结果证明了我国方法的有效性。

2 Related Work

UHD Image Restoration

超高分辨率（UHD）图像修复对于医学影像、视频 Stream 和数字监控等应用至关重要 [31, 36]。最近，深度学习方法通过采样和并行化技术有效地重构了 UHD 图像的细节和颜色。郑等人在[40]中提出了一个多引导双边上采样模型用于 UHD 图像去雾，通过多个引导输入增强清晰度。邓等人[6]开发了一个可分块整合网络用于 UHD 视频去模糊，采用多尺度整合方案来减轻运动和模糊伪影。王等人[26]提出了 LLFormer，这是一个基于 Transformer 的低光增强方法，利用轴向多头自注意力以及跨层注意力融合块来提高照明和对比度。此外，有一些方法[34, 39]可以通过构建轻量级模型和查找表实时重构清晰 UHD 图像。

尽管这些方法可以在实时中增强 UHD 图像，但取样和表格查找方法并未受到可信机制的监督，这限制了模型的泛化能力。相比之下，作者提出的 TSFormer 采用了 Min-p 采样进行自适应稀疏化，并基于随机矩阵理论采用可信过滤[2, 9]，同时提高了特征可靠性和恢复质量，保持了较低的计算成本。

Token Sampling Technology

当前，大型语言模型（LLM）使用一些 Token 采样技术来加快推理。传统方法，如 Top-k 过滤[8, 18]，根据大小优先权特征，但使用固定阈值，这些阈值可能无法适应不同的数据分布，可能丢弃有价值的信息。概率稀疏化方法，如 Min-p 采样[1, 42]，引入了基于概率的动态阈值，更好地适应数据分布，允许灵活地保留特征。受此启发，作者引入了 Token 技术来构建高效模型。在此基础上，确保实时性能的同时，引入了一种可信机制（随机矩阵理论）来准确采样 Token [4]。

Random Matrix Theory

信任过滤 Token （trusted filtering of tokens）的难度较大，这可能导致模型推理速度显著降低。与其他可信建模方法（如贝叶斯、变量推理和 Token 分布）相比，随机矩阵理论（RMT）在速度和准确性上进行权衡。随机矩阵理论（RMT）提供了一个分析并增强高维数据特征稳定性的框架 [7, 23]。通过分析特征值分布，基于 RMT 的方法可以识别和保留最稳定且显著的特征，从而提高鲁棒性和泛化能力 [10, 35]。然而，将 RMT 集成到深度学习以进行 UHD 修复的研究仍然不充分，这为在不显著增加计算成本的情况下增强特征可靠性提供了机会。

3 Method

在本节中，作者介绍 TSFormer，一种用于超高清图像恢复的高效 Transformer 网络。它是一种对称编码器和解码器结构，如图 2 所示。

Preliminaries

该方法通过作用于注意力图的稀疏化来过滤掉不重要的特征和噪声。

为了进一步优化作者的注意力机制的计算效率，作者在频域中使用快速傅里叶变换（FFT）进行注意力计算。FFT 显著降低了计算注意力分数的复杂性，特别是对于高分辨率图像。

1D x(t) 信号的前向傅里叶变换为：

逆傅里叶变换的公式为：

通过利用 FFT，作者可以更高效地执行注意力操作，尤其是在处理高分辨率图像时，通过在频域处理数据。

Overall Pipeline

解码器是解码器对高分辨率图像的细节和颜色进行重建的镜像操作。在相应的编码器和解码器 Level 之间引入特征融合块（FFB），以合并不同分辨率的特征。每个 FFB 应用了一个轻量级卷积块，带有 PReLU 激活，有助于在跨尺度上平滑地转换和细化特征，最终提高模型性能。

最后，作者在网络周围使用残差连接来获取恢复后的图像，如下所示：

Trusted Sparse Block (TSB)

自注意力机制不仅计算成本高，而且容易受到噪声影响，特别是在高频、细节丰富的 UHD 图像中。为了解决这个问题，作者设计了一个可信稀疏块（TSB）作为特征提取单元，它集成了 Min-p 采样和可信学习。

基于 Min-p 稀疏关注的 MSA (Min- Sparse Attention)。 作者的 MSA 利用 Min- 采样动态地保留只有高概率特征，使用傅里叶变换有效地在频域计算注意力。

为了捕捉局部特征交互，作者将和划分为 patches，并对每个 patches 应用快速傅里叶变换（FFT），将其转换为频域：

最小 -p 操作保留超过此阈值的元素在 M 中，通过设置较低的分数为零来稀疏化注意力图：

随机矩阵理论下的可信学习

Feature Fusion Blocks (FFB)

特征融合块（FFB）旨在将编码器和解码器不同层级的多种尺度特征进行集成。通过动态调整每个特征图的贡献，FFB 使模型能够有效地捕捉到细粒度细节和全局上下文。

给定两个来自网络不同阶段的输入特征图和，FFB 执行以下操作：

4 Experiments and Analysis

作者在五个超高清图像恢复任务上（包括低光增强、去雾、去模糊、除雪和脱雨）与最先进的算法进行了性能比较。

Experimental settings.

对于 UHD 低光图像增强，作者使用 UHD-LL 数据集[13]和 UHD-LOL4K [29]。为了评估去模糊能力，作者使用 UHD-Blur 数据集[25]。对于去雾评估，作者采用 UHD-Haze 数据集[25]。这些选择与前期研究[13, 25]所采用的方法一致。此外，为了评估作者的 UHD 图像除雪和去雨能力，作者利用了王等人[31]提出的 UHD-Snow 和 UHD-Rain 数据集。作者将 PSNR [11]和 SSIM [32]作为上述基准测试的评估指标。

对比方法。在作者研究中，作者将作者的方法与八种通用图像恢复（IR）技术进行了比较：SwinIR [16]，Uformer [33]，Restormer [38]，DehazeFormer [22]，Stripformer [24]，FFTformer [12]，和 SFNet [30]。此外，作者还包括了四种超高清图像恢复（UHDIR）方法：LLFormer [26]，UHD-Four [15]，UHD [41]，和 UHDformer [27]，以及 UHDDIP [28]。为了进行公平比较，作者使用这些模型的官方实现进行再训练，并使用与作者的提出方法相同数量的迭代进行评估。

训练细节 在作者的模型中，初始通道 C 为 32，扩展比设置为 2，FFN 中通道扩展因子 r 设置为 2.0。在训练过程中，作者使用 AdamW 优化器，批量大小为 6， patch 尺寸为 512，总共进行 300000 次迭代。初始学习率固定为。对于数据增强，随机应用水平和垂直翻转。整个框架在 PyTorch 上运行，使用 2 个 NVIDIA RTX 3090 GPU。

Main Results

低光图像增强结果。作者在 UHD 低光图像增强上评估了两个训练数据集，包括 UHD-LOL4k [26]和 UHD-LL [14]。在表 1 中，TSFormer 在 UHD-LOL4k 和 UHD-LL 数据集上都取得了最先进的低光图像增强结果，同时在这两个数据集上实现了最高的 PSNR 和 SSIM 值。尽管其性能强大，但 TSFormer 保持了轻量级的架构，只有 3.38M 参数，比其他高性能模型如 Restormer 和 UHDFour 更高效。这种准确性和效率的平衡表明 TSFormer 适合实时应用。图 3 说明了视觉改进，TSFormer 在低光条件下有效地减少了噪声并增强了细节，超越了先前的方法。

图像去模糊结果。 作者在 UHD-Blur 数据集上评估图像去模糊。如表 2 所示，UHD-Blur 数据集上的图像去模糊定量结果。TSFormer 在关键指标上取得了显著改进，性能优异。具体而言，TSFormer 实现了 29.52 dB 的 PSNR，超过了 UHDformer 的 28.82 dB，并优于其他通用模型，如 Restormer、Uformer 和 Stripformer，它们在 25-25.4 dB 之间。这突显了 TSFormer 在 UHD 图像中恢复高分辨率细节的能力。作者还评估了 GoPro [19]数据集的性能，并报告了表 3 中的结果。

图像去雾结果。表 4 展示了在 UHD-Haze 数据集上使用训练好的模型得到的 UHD-Haze 的定量结果。TSFormer 在 PSNR 和 SSIM 指标上取得了最高分，相较于 UHDformer 在 UHD-Haze 数据集上实现了 0.77 dB 的提升，同时在各视觉质量方面始终优于其他方法。TSFormer 在保持与先驱模型相当的可比参数数量的同时，显著降低了 LPIPS，展示了其在效率和性能方面的良好平衡。图 5 说明了定性结果，其中 TSFormer 在所有方法中产生的输出最为清晰，展示了其有效去除雾气并恢复其他方法通常被雾气掩盖的细微细节的能力。

图像去雨和除雪结果。 作者使用构建的 UHD-Rain 数据集评估 UHD 图像去雨。结果报告在表 5 中。TSFormer 实现了最先进的性能，显著优于所有现有方法的关键指标。与先前的模型如 UHDformer、Restormer 和 UHD-DIP 相比，TSFormer 在实现显著更高的保真度和感知质量的同时，还能有效处理 UHD 分辨率。图 6 说明了视觉比较，表明 TSFormer 在去除雨条纹的同时保留更细微的细节方面更有效。

作者在 UHD-Snow 数据集上实现了超高清降雪实验，结果汇总于表 8，TSFormer 取得了优越性能。

5 Ablation studies

在本节中，作者评估了不同采样策略对 UHD 图像去雾的影响。具体来说，作者比较了三种采样技术：Top-k 采样，无信任机制的 Min-p 采样，以及带有信任机制的 Min-p 采样。

Sampling Method Comparison

为了理解每种采样方法对注意力分布和特征保留的影响，作者在表 6 中提供了定量结果，并在图 7 中给出了累积概率分布的视觉化。

表 6 显示，Min-p 采样与可信机制在去雾质量方面优于其他采样技术，实现了较高的 PSNR 和 SSIM 分数，这些分数反映了更优秀的结构保真度。图 9 显示了每种方法如何捕捉细节，Min-p 采样与可信机制主要关注高置信区域，同时有效过滤噪声。

表格 2：在 UHD-Blur 数据集上的图像去模糊结果。TSFormer 在 PSNR，SSIM 和 LPIPS 指标上实现了最先进的性能，这证明了其在 UHD 去模糊任务中的有效性和效率。

表格 2 展示了在 UHD-Blur 数据集上的图像去模糊结果。TSFormer 在 PSNR，SSIM 和 LPIPS 指标上实现了最先进的性能，这证明了其在 UHD 去模糊任务中的有效性和效率。

Trusted Strategies Comparison

作者评估了将可信学习集成到作者的模型中以提高稳定性和鲁棒性的不同策略。

完整特征值分解（FED）。这种策略涉及对每个注意力块进行完整的特征值分解，以识别并丢弃不稳定的块。给定一个稳定性阈值 τ，只有满足最大特征值 λ_{max} < τ 的块被保留。尽管这种方法在隔离稳定特征方面有效，但计算成本高，导致推理时间较长。

迭代稳定性调整（ISA）。在这种方法中，稳定性阈值是根据在迭代过程中稳定区域（即，特征值方差较小的区域）的方差动态调整的：

然而，这种迭代调整虽然具有自适应性，但由于重复计算，会显著增加运行时间，且没有带来显著的性能提升。如表 9 所示，采用可信机制的 Min-p 采样在保证最低峰值信噪比和 SSIM 值的同时，大幅降低了运行时间，相较于其他方法具有明显优势。

6 Running Time and Application

在不考虑 I/O 操作的情况下，作者推理 4K 图像需要 40 fps。作者的方法高效的原因有两个：一是 Token 过滤机制确保了模型的稀疏性，二是特征图在可信机制下进行下采样。

图 10 展示了在 DarkFace 数据集[37]上，使用各种图像增强方法作为预处理步骤的目标检测结果的定性比较。

7 MSA of Potential

表 10 展示了在各种 UHD 模型中，通过集成 MSA 所实现的效率提升。通过动态关注高置信度特征，MSA 在 LLFormer、UHDFormer 和 UHDFour 模型中显著减少了 FLOPs 和运行时间。值得注意的是，MSA 在 UHDFormer 中的集成将 FLOPs 降低 20%，运行时间降低 25%，同时对性能的影响很小。这些改进证明了 MSA 在实时 UHD 图像修复应用中的潜力。图 11 说明了 MSA 在增强多样场景下低光图像质量的影响。通过选择性保留高置信度特征，MSA 使 LLFormer、UHDFormer 和 UHDFour 模型在应对挑战性低光条件时能更好地泛化。这种能力对于照明条件较差的实际应用环境至关重要。视觉结果表明，使用 MSA 可以提高图像的清晰度和细节保留，突显了其在实际场景中的有效性。

8 Conclusion

作者提出了一种高效且鲁棒性强的模型，称为 TSFormer，它专注于处理超高清图像。

其优势源于一种可信 Token 过滤机制，通过动态阈值和随机矩阵理论实现。

TSFormer 提出了一种可信 Token 过滤机制，可应用于其他 Transformer 框架以提高鲁棒性。

实验结果表明，与最先进的多种 UHD 恢复任务相比，TSFormer 在多个 UHD 恢复任务中实现了优越性能。

参考文献

[0]. TSFormer: A Robust Framework for Efficient UHD Image Restoration.

END

作者：小书童
来源：集智书童

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。