[论文解读] Test-time Batch Normalization

Tao Yang, Shenglong Zhou+3 more|arXiv (Cornell University)|May 20, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文提出GpreBN，一种新型的测试时批量归一化层，其在保持训练时梯度反向传播形式的同时，利用数据集级别的统计信息以提升对数据分布偏移的鲁棒性。通过在推理过程中通过熵最小化优化仿射参数，GpreBN在域泛化和鲁棒性基准测试中均达到最先进性能，优于Tent和T3A等先前方法。

ABSTRACT

Deep neural networks often suffer the data distribution shift between training and testing, and the batch statistics are observed to reflect the shift. In this paper, targeting of alleviating distribution shift in test time, we revisit the batch normalization (BN) in the training process and reveals two key insights benefiting test-time optimization: $(i)$ preserving the same gradient backpropagation form as training, and $(ii)$ using dataset-level statistics for robust optimization and inference. Based on the two insights, we propose a novel test-time BN layer design, GpreBN, which is optimized during testing by minimizing Entropy loss. We verify the effectiveness of our method on two typical settings with distribution shift, i.e., domain generalization and robustness tasks. Our GpreBN significantly improves the test-time performance and achieves the state of the art results.

研究动机与目标

• 解决深度学习模型在训练与测试之间数据分布偏移的问题。
• 识别有效测试时批量归一化层适应的关键设计原则。
• 开发一种在推理过程中保持梯度反向传播保真度并利用稳定数据集级统计信息的方法。
• 在包括域泛化和鲁棒性在内的多种分布偏移场景中实现最先进性能。

提出的方法

• 提出GpreBN，一种新型批量归一化层，将梯度反向传播形式与归一化统计选择解耦。
• 通过在测试时优化过程中动态计算批量统计量，保留训练时所见的跨样本梯度依赖关系。
• 使用数据集级别的运行统计量（如移动平均）进行归一化，提升推理过程中的稳定性和鲁棒性。
• 通过在测试批量数据上对GpreBN的仿射参数（γ, β）进行熵最小化优化。
• 在不改变计算图或梯度流动的前提下，灵活使用不同统计量（如源域、目标域或目标域平均统计量）。
• 引入rT3A，一种互补的分类器自适应技术，选择性地将目标域统计量应用于浅层网络，验证了在鲁棒性任务中分布偏移源于特征提取器的假设。

实验结果

研究问题

• RQ1 保持训练时梯度反向传播形式对测试时自适应性能有何影响？
• RQ2 使用数据集级别的统计信息是否能提升测试时自适应在分布偏移下的鲁棒性？
• RQ3 为何现有测试时方法如Tent在域泛化任务中失败，尽管其在鲁棒性任务中表现成功？
• RQ4 分布偏移的位置（如浅层与深层）在多大程度上影响测试时自适应的有效性？
• RQ5 是否可以设计一种统一的测试时自适应方法，在域泛化和鲁棒性任务中均有效？

主要发现

• GpreBN在域泛化任务中达到最先进性能，在CIFAR-10-C数据集上将平均损坏误差（mCE）降低至8.9，优于Tent（13.2）和T3A（14.4）。
• 在VLCS数据集上，使用目标域统计量的GpreBN实现13.1%的top-1错误率，显著优于Tent（13.8%）和使用目标域统计量的BN（13.8%）。
• GpreBN在域泛化和鲁棒性任务中均保持强劲性能，而Tent在域泛化任务中失败，T3A在鲁棒性任务中表现欠佳。
• 消融实验确认，保持跨样本梯度反向传播和使用数据集级别统计信息对性能至关重要。
• rT3A变体表明，仅对前50%的层应用目标域统计量导致性能下降极小（14.6 vs. 14.4），支持了鲁棒性偏移源于早期特征提取器的假设。
• 即使在模型于源域上微调的情况下，该方法依然有效，而目标统计量在此类场景中失效，表明GpreBN对微调引起的分布偏移具有鲁棒性。