引言

想象一下，你训练了一个最先进的计算机视觉模型，用于自动驾驶汽车的行人检测。它在原本进行训练的阳光明媚的加利福尼亚运行完美。但当你把它部署到伦敦阴雨绵绵的街道上时，准确率直线下降。视觉条件——即“分布”——发生了变化。

为了解决这个问题，研究人员使用了测试时适应 (Test-Time Adaptation, TTA) 。 TTA 并不在训练后冻结模型，而是允许模型从新的、不断进入的数据 (比如伦敦的雨街) 中实时继续学习。它有效地“在线”微调模型以适应当前环境。

然而，这里有一个巨大的陷阱: 内存 。

更新神经网络需要反向传播。反向传播需要在前向传递过程中存储网络的中间输出 (即激活值) ，以便在反向传递中用于计算梯度。在大型服务器 GPU 上，这没问题。但在手机、无人机或物联网传感器等边缘设备上，这种内存需求往往是无法满足的。

今天，我们将深入探讨一篇提出巧妙解决方案的论文: SURGEON 。 该方法引入了“动态激活稀疏性 (Dynamic Activation Sparsity) ”。它无需改变模型的原始训练方式，就能智能地决定哪些部分的网络内存值得保存，哪些可以丢弃 (剪枝) 以节省空间。

瓶颈: 为什么 TTA 在边缘端很难

在理解 SURGEON 之前，我们需要明确它所解决的问题。

TTA 与 FTTA 的对比

在标准的测试时适应 (TTA) 中，你可能拥有原始训练数据，或者被允许修改训练过程，以便为未来的适应做准备。

然而，在许多现实场景中，我们面临的是全测试时适应 (Fully Test-Time Adaptation, FTTA) 。 在 FTTA 中:

  1. 你拥有一个预训练模型。
  2. 你没有原始训练数据 (通常出于隐私或存储原因) 。
  3. 你不能修改模型最初的训练方式。
  4. 你必须立即适应新数据。

内存墙

FTTA 最大的障碍是反向传播的内存成本。当神经网络处理图像时 (前向传播) ，它会在每一层生成“激活值”。为了更新权重 (反向传播) ，微积分的链式法则需要这些保存下来的激活值。

\[ \Delta W _ { i } = \frac { \partial L } { \partial W _ { i } } = \frac { \partial L } { \partial A _ { i + 1 } } \frac { \partial A _ { i + 1 } } { \partial W _ { i } } = A _ { i } ^ { T } \frac { \partial L } { \partial A _ { i + 1 } } , \]

正如上面的公式 (取自论文的预备知识部分) 所示，计算权重更新 \(\Delta W_i\) 需要 \(A_i\) (激活值) 。存储每一层的这些 \(A_i\) 张量会消耗大量内存——通常超过边缘硬件的容量。

现有解决方案及其缺陷

研究人员之前曾尝试解决这个问题，但之前的尝试都有局限性:

比较 EcoTTA、MECTA 和 SURGEON 的图 1。

如图 1 所示:

  • (a) EcoTTA: 这种方法向模型添加了轻量级的“元网络”。它冻结主干网络，只更新这些小网络。虽然内存效率高，但它不符合 FTTA 的要求，因为它需要对原始训练过程进行特定的修改来“预热”这些元模块。
  • (b) MECTA: 这种方法试图通过仅更新特定通道或层来节省内存。然而，它根据批归一化 (BN) 统计数据来决定更新什么。这使得它与不大量使用 BN 层的新架构 (如现代视觉 Transformer，ViT) 不兼容。
  • (c) SURGEON (本文方法) : 这就是我们要分析的方法。它适用于任何架构 (CNN 或 Transformer) ，并且不需要对原始训练进行任何更改。它通过剪枝激活值来解决问题——也就是字面上删除内存缓存中非绝对必要的部分。

SURGEON 方法: 动态激活稀疏性

SURGEON 的核心洞察是，并非所有的激活值都是生而平等的。有些层对学习 (准确性) 至关重要，而有些层只是内存消耗大户。此外，这种情况会随着数据的不同而变化。

SURGEON 没有为反向传播保留 100% 的激活值，而是应用了动态激活稀疏性 。 它保存激活值的“稀疏”版本——可能只保留 50% 或 30% 的数值——并将剩余部分归零。

工作流程

图 2 展示了带有剪枝的前向和反向传播过程。

图 2 优雅地展示了这个过程:

  1. 前向传播 (Forward Pass) : 数据流经网络。在每一层 \(i\)，模型计算激活值 \(A_i\)。
  2. 剪枝 (Pruning) : 在将 \(A_i\) 存储到缓存之前，SURGEON 会为该层计算一个特定的剪枝率 \(p_i\)。它移除最不重要的元素，得到一个稀疏张量 \(\dot{A}_i\)。
  3. 缓存 (Caching) : 只有 \(\dot{A}_i\) 的非零值 (及其索引) 被存储在内存中。
  4. 反向传播 (Backward Pass) : 在计算梯度时，系统重建稀疏张量并使用它来更新权重。

其中的奥妙在于 SURGEON 如何决定从每一层剪掉多少 。 它不使用固定的数字 (比如“到处都剪掉 50%”) 。相反，它针对两个相互竞争的目标进行优化: 最小化内存使用量，同时最大化准确性。

\[ \operatorname* { m i n } \quad \alpha \cdot \mathbf { M e m o r y } - \beta \cdot \mathbf { A c c u r a c y } , \]

为了实现这种平衡，作者引入了两个指标: 梯度重要性和层激活内存重要性 。

指标 1: 梯度重要性 (\(G\))

首先，系统需要知道哪些层实际上在学习。如果一层的梯度非常小，这意味着权重变化不大，因此该层的激活值对适应过程贡献不显著。

作者基于权重梯度的幅度定义了层 \(i\) 的梯度重要性 (\(G_i\)) :

\[ \Delta w _ { i } = \frac { \partial L } { \partial w _ { i } } , \quad G _ { i } = \sqrt { \frac { \sum _ { j = 1 } ^ { N _ { i } } \left( \Delta w _ { j } \right) ^ { 2 } } { N _ { i } } } , \]

这里，\(G_i\) 本质上是该层梯度的平均强度。\(G_i\) 越高意味着该层正在积极适应新数据，因此我们应该保留更多它的激活值 (少剪枝) 以确保准确性。

指标 2: 层激活内存 (\(M\))

其次，系统会考虑成本。有些层 (通常是 CNN 中早期的、高分辨率的层) 会产生巨大的激活图，吞噬 RAM。其他层 (网络深层) 则较小。

内存重要性 (\(M_i\)) 指标量化了一层在内存方面的“效率”:

\[ m _ { i } = \mathbf { s i z e } ( A _ { i } ) , \quad M _ { i } = - \log \left( \frac { m _ { i } } { \sum _ { i = 1 } ^ { l } m _ { i } } \right) , \]

该公式计算激活值大小 \(m_i\) 相对于网络总大小的比例。对数缩放有助于归一化数值。本质上，占用大量内存的层会导致其被保留的优先级降低——我们要大力剪枝这些层以节省空间。

将指标结合为剪枝率

最后，SURGEON 将这两个指标结合成一个单一的重要性指标 (\(I_i\)) 。它将两个指标归一化到 0-1 的范围，以便公平比较:

\[ I _ { i } = \mathbf { N o r m } ( M _ { i } ) \times \mathbf { N o r m } ( G _ { i } ) , \]

这个指标 \(I_i\) 告诉我们层 \(i\) 的激活值有多“宝贵”。

  • 高 \(I_i\): 该层学习量大 (高 \(G\)) 并且/或者内存效率高 (高内存得分，意味着尺寸小) 。 动作: 少剪一点。
  • 低 \(I_i\): 该层学习量不大，或者它占用了太多 RAM。 动作: 多剪一点。

最终的剪枝率 \(p_i^t\) (丢弃激活值的百分比) 是针对每个批次 \(t\) 动态计算的:

\[ p _ { i } ^ { t } = 1 - \frac { I _ { i } ^ { t } } { \operatorname* { m a x } _ { i \in \{ 1 , 2 , . . . , l \} } { ( I _ { i } ^ { t } ) } } , \]

这个公式确保网络中最重要的层被剪枝最少，其他所有层都相对于它进行缩放。

实验结果

这真的有效吗？作者将 SURGEON 与标准基线 (如 TENT 和 CoTTA) 以及内存高效基线 (如 EcoTTA 和 MECTA) 进行了测试。

为什么动态优于静态

你可能会想: “为什么要计算所有这些指标？为什么不直接制定一个全局规则，把所有地方的激活值都剪掉 50% 呢？”

作者将 SURGEON 与“静态激活稀疏性” (固定剪枝率) 进行了比较。

图 3: 动态与静态稀疏性。

图 3 展示了对比结果。X 轴代表剪枝率 (丢弃了多少数据) ，Y 轴是错误率 (越低越好) 。

  • 看红色星号 (SURGEON) 与青色线 (Static) 的对比。
  • SURGEON 在保持高稀疏度 (约 80-90%) 的同时，实现了显著更低的错误率。
  • 这证明了剪什么与剪多少同样重要。通过保留重要数据并丢弃无用数据，SURGEON 优于盲目的静态剪枝。

分析指标

为了理解算法在网络内部实际做了什么，让我们看看重要性指标在各层之间是如何变化的。

图 4: 跨层重要性指标的可视化。

图 4 可视化了 ResNet 各层的重要性得分。

  • 块 1 (第 4-16 层) : 注意这里的重要性急剧下降。这些早期层拥有巨大的激活图 (高内存成本) 。SURGEON 识别到了这一点并降低了它们的重要性得分，从而导致激进的剪枝以节省 RAM。
  • 深层: 在后面的块中，重要性较高。这些层较小 (低内存成本) ，但携带对分类至关重要的语义信息 (高梯度重要性) 。SURGEON 保留了这些层。

红线 (“Ours”) 代表结合后的指标，有效地平衡了蓝线 (仅梯度) 和绿线 (仅内存) 。

基准测试: CIFAR 和 ImageNet

论文提供了详尽的表格，但让我们关注标准基准测试中的关键结论。

CIFAR-10-C 结果 (表 1) :

表 1: CIFAR-10-C 结果。

在表 1 中，看看 CoTTA 和 SURGEON 之间的比较:

  • CoTTA: 准确率不错 (16.2% 错误率) ，但缓存大小巨大 (3697 MB) 。这根本无法放入许多边缘芯片中。
  • SURGEON: 准确率相当 (18.1% 错误率) ，但缓存大小极小 (325 MB) 。
  • 减少量: 这是内存使用量的 10 倍减少 。
  • 此外，请注意 SURGEON 在“Original” (原始) 列中标记为“X”，这意味着它在不修改训练过程的情况下工作，这与 EcoTTA 不同。

ImageNet-C 结果 (表 4) :

表 4: ImageNet-C 结果。

在更难的 ImageNet 数据集( 表 4 )上，趋势仍在继续。

  • TENT: 2714 MB 缓存。
  • SURGEON: 1125 MB 缓存。
  • SURGEON 将内存使用量减少到了 TENT 的一半以下，同时保持了几乎相同的准确率 (55.5% vs 55.2%) 。

真实世界部署: Jetson Xavier

任何“高效”算法的终极测试都是在实际边缘硬件上运行。作者将 SURGEON 部署在了 NVIDIA Jetson Xavier NX (一种流行的嵌入式 AI 计算机) 上。

表 5: Jetson Xavier NX 上的性能。

表 5 揭示了实际影响:

  • CoTTA 处理一个批次需要 522 毫秒。
  • SURGEON 仅需 117 毫秒。
  • 内存: SURGEON 使用的缓存大约是 CoTTA 的 1/10，与 MECTA 的内存相当。

至关重要的是，由于 SURGEON 使用的内存更少，它避免了设备耗尽 RAM 时发生的“交换 (swapping) ”减速，从而实现了更快的整体推理速度。

结论

AI 从研究实验室向现实世界的过渡依赖于鲁棒性。模型必须适应不断变化的环境 (雨、雾、传感器噪声) ，而不会导致运行它们的硬件崩溃。

SURGEON 代表了全测试时适应 (FTTA) 向前迈出的重要一步。通过将内存视为一种动态资源——在有助于学习的地方花费它，在无助于学习的地方节省它——它使得复杂的适应算法能够在受限设备上运行。

主要收获:

  1. 即插即用: 与 EcoTTA 不同，SURGEON 适用于标准的预训练模型。
  2. 架构无关: 与 MECTA 不同，它不依赖于批归一化 (Batch Norm) ，这使其能够适应未来的 Transformer 架构。
  3. 智能缓存: 通过利用梯度和内存重要性，它优于静态剪枝，在极小的准确率损失下节省高达 10 倍的内存。

对于希望在边缘部署鲁棒 AI 的学生和工程师来说，SURGEON 提供了一个平衡适应性和效率的实用蓝图。