基于短预热阶段的快速认证鲁棒训练

近年来，基于边界传播（bound propagation）的认证鲁棒训练方法被提出用于训练具有可证明鲁棒性保证的神经网络。尽管包括区间边界传播（IBP）和 CROWN-IBP 在内的最先进（SOTA）方法每批次训练复杂度与标准神经网络训练相近，它们通常使用长达数百或数千轮次的预热（warmup）调度才能达到 SOTA 性能，因而仍然代价高昂。在本文中，我们指出了现有方法中的两个重要问题，即初始化时边界爆炸，以及 ReLU 激活状态的不平衡，并改进了 IBP 训练。这两个问题使得认证训练困难且不稳定，因此先前工作需要使用长预热调度。为缓解这些问题并以更短预热进行更快的认证训练，我们基于 IBP 训练提出三项改进：1) 我们推导了一种用于 IBP 训练的新权重初始化方法；2) 我们提出在模型每一层完全加入批归一化（BN），因为我们发现 BN 能减少 ReLU 激活状态的不平衡；3) 我们还设计了正则化，在预热期间显式收紧认证边界并平衡 ReLU 激活状态。我们能够在极短训练调度（分别为总共 160 和 80 轮次）下，在 CIFAR-10 ($\epsilon=\frac{8}{255}$) 上取得 65.03% 的验证误差，在 TinyImageNet ($\epsilon=\frac{1}{255}$) 上取得 82.36% 的验证误差，优于在相同网络架构下使用数百或数千轮次训练的文献 SOTA。代码见 https://github.com/shizhouxing/Fast-Certified-Robust-Training。