[论文解读] Distributed Adversarial Training to Robustify Deep Neural Networks at Scale

Gaoyuan Zhang, Songtao Lu+8 more|arXiv (Cornell University)|Jun 13, 2022

Adversarial Robustness in Machine Learning被引用 5

一句话总结

本文提出分布式对抗训练（DAT），一种可扩展的、通用的多机大规模批量对抗训练框架。在非凸设置下，其收敛至一阶驻点的速率为 $O(1/\sqrt{T})$，同时在 ImageNet 上保持了鲁棒准确率，并支持梯度压缩与半监督学习，实现高效训练。

ABSTRACT

Current deep neural networks (DNNs) are vulnerable to adversarial attacks, where adversarial perturbations to the inputs can change or manipulate classification. To defend against such attacks, an effective and popular approach, known as adversarial training (AT), has been shown to mitigate the negative impact of adversarial attacks by virtue of a min-max robust training method. While effective, it remains unclear whether it can successfully be adapted to the distributed learning context. The power of distributed optimization over multiple machines enables us to scale up robust training over large models and datasets. Spurred by that, we propose distributed adversarial training (DAT), a large-batch adversarial training framework implemented over multiple machines. We show that DAT is general, which supports training over labeled and unlabeled data, multiple types of attack generation methods, and gradient compression operations favored for distributed optimization. Theoretically, we provide, under standard conditions in the optimization theory, the convergence rate of DAT to the first-order stationary points in general non-convex settings. Empirically, we demonstrate that DAT either matches or outperforms state-of-the-art robust accuracies and achieves a graceful training speedup (e.g., on ResNet-50 under ImageNet). Codes are available at https://github.com/dat-2022/dat.

研究动机与目标

• 解决由于计算成本高和数据规模大，导致对抗训练（AT）在大规模深度学习中可扩展性受限的问题。
• 在数据或计算资源被划分到多台机器的分布式环境中，实现有效的对抗训练。
• 开发一个通用框架，支持分布式环境中监督、半监督以及压缩梯度形式的对抗训练。
• 在标准非凸优化条件下，理论上建立 DAT 的收敛性保证。
• 通过实证验证 DAT 在多种设置下的鲁棒性与效率：ImageNet、随机平滑、无标签数据以及不同硬件配置。

提出的方法

• 设计一种分布式极小极大优化框架，通过大规模批量更新实现跨多台机器的对抗训练可扩展性。
• 集成梯度压缩与自适应学习率（如 LAMB 优化器），以减少分布式训练中的通信开销。
• 将算法设计为统一管道，支持多种攻击生成方法（如 PGD、FGSM）和多种数据类型（有标签、无标签）。
• 理论上分析梯度估计、量化、自适应学习率以及内部最大化查询器带来的下降误差。
• 推导出在一般非凸设置下，收敛至一阶驻点的速率为 $O(1/\sqrt{T})$。
• 使用 NCCL 通信后端实现框架，并支持基于 InfiniBand 和 NVLink 互连的 HPC 集群。

实验结果

研究问题

• RQ1 能否在不牺牲鲁棒性的前提下，通过分布式计算将对抗训练有效扩展到大规模批量设置？
• RQ2 分布式优化组件（如梯度压缩与自适应学习率）如何影响对抗训练中的收敛性与鲁棒准确率？
• RQ3 在非凸、随机且分布式设置下，分布式对抗训练的理论收敛速率是多少？
• RQ4 所提出的框架在使用混合数据类型（有标签/无标签）的大规模数据集（如 ImageNet）上训练时，是否能保持鲁棒性能？
• RQ5 该框架在不同硬件配置（包括标准集群与高性能计算（HPC）系统）下的表现如何？

主要发现

• DAT 在使用 ResNet-50 的 ImageNet 上实现了与最先进方法相当或更优的鲁棒准确率，即使在大批次设置下亦如此。
• 该框架在提升鲁棒性的同时保持了较高的标准准确率，避免了传统分布式 AT 中常见的性能下降问题。
• 理论分析证实，其收敛至一阶驻点的速率为 $O(1/\sqrt{T})$，与非凸设置下的标准 SGD 一致。
• 实证结果表明，随着计算节点数量增加（测试了 1、3、6 个节点），训练时间实现平稳加速，准确率下降可忽略。
• DAT 支持半监督鲁棒训练以及鲁棒预训练-微调，展示了在多样化训练范式中的泛化能力。
• 在采用 InfiniBand 和 NVLink 的 HPC 集群上，通信成本更低，可扩展性优于标准分布式系统。