一文了解Arm神经超级采样 (Arm Neural Super Sampling, Arm NSS) 深入探索架构、训练和推理

作者：Arm 工程部计算机视觉高级架构师 Liam O’Neil

本文将从训练、网络架构到后处理和推理等方面，深入探讨 Arm 神经超级采样 (Arm Neural Super Sampling, Arm NSS) 的工作原理，希望为机器学习 (ML) 工程师和移动端图形开发者来详细解释 Arm NSS 的运行机制，及其如何在移动端硬件上进行部署。

用神经超级采样取代启发式方法的原因

时域超级采样 (TSS)，也称为时域抗锯齿 (TAA)，已经成为过去十年中抗锯齿技术的行业标准解决方案。它具备诸多优势，包括能够解决各类锯齿问题，对于延迟渲染的计算效率高，并且可扩展至图像的优化升级。然而，该解决方案也面临一些挑战。例如，TSS 中目前常用的人工设计的启发式方法难以扩展，并且需要在不同内容中不断调整，像鬼影、去遮挡伪影和时域不稳定性等问题依然存在，并且在与图像优化升级技术结合时变得更加棘手。

Arm NSS 通过一套从数据中学习且经过训练的神经模型，而非静态规则来克服这些限制。它可以跨条件和内容类型进行泛化，更有效地适应运动动力和识别锯齿模式。这些能力使其能够更可靠地处理那些对 AMD 超级分辨率锐画技术 2 (FSR2) 和 Arm 精锐超级分辨率技术 (Arm Accuracy Super Resolution, Arm ASR) 等方法颇具挑战的边缘侧用例。

训练 Arm NSS 网络：带反馈的递归学习

Arm NSS 使用以每像素一个样本渲染的 540p 帧序列进行训练。每帧与以每像素 16 个样本渲染的 1080p 真实标签 (Ground Truth) 图像配对。序列大约包含 100 帧，以帮助模型理解图像内容如何随时间变化。

解决方案的输入包括渲染图像的颜色、运动向量和深度，以及引擎元数据，如抖动向量和摄像头矩阵。该模型采用递归式训练，在每次执行反向传播之前，会先对多帧序列进行前向传播。这使得网络能够随时间传播梯度，并学习如何累积信息。

该网络采用时空损失函数进行训练，能对空域保真度和时域一致性方面的误差进行惩罚。空域保真度用于确保单帧图像清晰、细节丰富且视觉准确，有助于保留边缘、纹理和精细结构；而时域稳定性则抑制了连续帧之间可能出现的闪烁、抖动或其他形式的时域噪声。

此项训练是在 PyTorch 中采用成熟的技术实践进行，其中包括 Adam 优化器、余弦退火学习率调度，以及标准数据增强策略。预处理和后处理流程使用 Slang 语言编写，以确保灵活性和性能，而量化感知训练则采用 ExecuTorch 完成。

网络架构和输出设计

Arm NSS 网络采用具有跳跃连接的四层级 UNet 主干架构，以保持空间结构。它分别通过三个编码器和解码器模块对输入数据进行下采样和上采样处理。

Arm 评估了多种方法：

图像预测‌：易于实现，但在量化条件下表现不佳，并会产生视觉伪影。

内核预测：泛化能力良好且能有效量化，但由于需要许多大型内核映射，导致带宽开销较高。

参数预测（所选方法）‌：为每个像素输出少量参数。这些参数驱动后处理步骤（如滤波和样本累积）。该方法对量化友好且带宽效率高。

该网络生成三种逐像素输出：

4 x 4 滤波器内核；

用于累积和校正的时域系数‌；

隐藏状态张量‌作为时域反馈传递至下一帧。

该网络输出服务于两条路径：

滤波器内核与时域系数‌被后处理阶段用于计算经优化升级的最终图像；

隐藏状态被前向传递至下一帧推理。与 Arm ASR 此类技术依赖手动调优的启发式方法不同，Arm NSS 等 ML 方法具有三重优势：

Arm NSS 能估计动态内核滤波器和参数，以逐像素粒度解决锯齿问题；

Arm NSS 通过利用采集多帧历史状态的时域反馈，实现更出色的时域稳定性；

Arm NSS 可针对新游戏内容进行调优，使开发者能进一步针对其特定游戏优化图像质量。

通过时域反馈提升帧间一致性

Arm NSS 引入了两种关键反馈机制来解决时域不稳定性问题：

将先前帧的隐藏特征前向传递，使网络能够学习变化与持续的部分；

通过计算亮度导数来检测闪烁的细微特征，从而突显指示不稳定的时域差异。

这些输入帮助模型在不依赖人工设计规则的情况下保持时域稳定性。

预处理阶段：输入准备

在推理开始之前，基于 GPU 的预处理阶段会准备 Arm NSS 所需的输入数据，其中包括收集每个像素的属性（如颜色、运动向量和深度）。此外，该阶段还会计算亮度导数，这是一个用于标记细微特征闪烁的时域信号，以及去遮挡掩码，以突出显示陈旧的历史记录，并重投影历史隐藏特征。

这些数据将被整合为用于神经网络的单一输入张量。该阶段作为计算着色器运行，在推理调用前执行，即在使用面向 Vulkan 的 ML 扩展的 GPU 上运行。

后处理：从原始输出到最终帧

在推理完成后，作为计算着色器运行的后处理阶段将构建最终输出颜色。所有步骤均集成至渲染图中，并专为移动端高效运行设计，具体包括以下步骤：

运动向量扩张，减少重投影历史时的锯齿问题；

历史重投影‌，采用 Catmull-Rom 滤波器降低重投影的模糊度；

滤波，应用 4 x 4 内核对当前颜色输入执行抗锯齿处理；

稀疏优化升级，将抖动的低分辨率样本映射到高分辨率网格上，缺失像素都以零值填充，然后用 4 x 4 内核稀疏滤波，执行插值和抗锯齿处理，类似于去马赛克；

校正，使用预测的 theta 参数剔除陈旧历史数据；

样本累积‌，通过预测的 alpha 参数将新数据与历史缓冲混合，在色调映射域执行以避免“萤火虫”伪影。

验证质量

Arm 使用了峰值信噪比 (Peak Signal-to-Noise Ratio, PSNR)、结构相似性指数 (Structural Similarity Index, SSIM) 和渲染聚焦感知误差指标 FLIP 等指标来评估 Arm NSS。虽然这些指标未必总是匹配人类感知，但它们有助于发现问题案例，并通过跟踪各种指标增强信心。

通过持续集成 (CI) 工作流回放测试序列并记录 Arm NSS、Arm ASR 和其他基准的性能。有关视觉比较和感知评估的详细内容，请参阅技术白皮书。

白皮书：https://developer.arm.com/documentation/111019

在 540p 到 1080p 的比较中，Arm NSS 提升了稳定性和细节保留性能，并且在快速运动、部分遮挡物体和细微几何体场景中表现出色。与 Arm ASR 或 AMD FSR2 等非神经方法不同，Arm NSS 无需反应掩码，即可处理粒子效果。

Arm NSS 能否实时运行？

虽然搭载神经加速器的芯片尚未发布，但我们可以通过一些最低性能假设和执行网络推理所需的乘积累加运算 (MAC) 数量来估计 Arm NSS 是否足够快。该分析适用于任何在吞吐量、功率和利用率方面满足这些相同假设的加速器。我们假定在可持续的 GPU 时钟频率下，目标每瓦 10 TOP/s 的神经加速是可实现的。

在持续的性能条件下，我们对每帧的优化升级处理设定了‌不超过四毫秒‌的目标。在低时钟频率 GPU 上，推理前后运行的着色器阶段共耗时约 ‌1.4 毫秒‌。基于此预算，Arm NSS 需控制在约 ‌27 GOPs ‌以下。而参数预测的网络使用约为 ‌10 GOPs‌，即使神经网络加速器的效率仅为 40%，这在该范围内仍绰绰有余。

早期模拟数据显示，Arm NSS 在 1.5 倍优化升级（平衡模式）下的运行时间约为 Arm ASR的 75%，而在两倍优化升级（平衡模式）下则预计性能会优于 Arm ASR。这一效率提升得益于用简化的推理过程取代了复杂的启发式方法。

即刻利用 Arm NSS 进行构建

Arm NSS 引入了一种由 ML 驱动的时域超级采样的实用方法。它用经过学习的滤波器和稳定性提示取代了人工调优的启发式方法，并能够在移动端硬件的实时性限制内运行。

其训练方法、紧凑架构和利用面向 Vulkan 的 ML 扩展使其兼具高性能和适应性。对于构建神经渲染解决方案的 ML 工程师来说，Arm NSS 是一个可部署、结构良好，且运行于图形管线内的推理示例。

想要探索 Arm 神经图形开发套件，查看示例代码，并探索网络结构，请访问 Arm 开发者中心的 Arm NSS 页面。欢迎开发者使用开发套件或为自己的内容重新训练 NSS，并提供反馈，你的见解将有助于塑造移动设备上神经渲染的未来。