本文是一篇关于提高零阶优化的扩展性的研究,代码已开源,论文已被 ICLR 2024 接收。
今天介绍一篇密歇根州立大学 (Michigan State University) 和劳伦斯・利弗莫尔国家实验室(Lawrence Livermore National Laboratory)的一篇关于零阶优化深度进修框架的文章 “DeepZero: Scaling up Zeroth-Order Optimization for Deep Model Training”,本文被 ICLR 2024 接收,代码已开源。
论文地址:https://arxiv.org/abs/2310.02025
项目地址:https://www.optml-group.com/posts/deepzero_iclr24
1. 背景
零阶(Zeroth-Order, ZO)优化已成为解决机器进修(Machine Learning)问题的热门技术,特别是在一阶(First-Order, FO)信息难以或无法获得的情况下:
物理学和化学等学科:机器进修模型可能与复杂的仿照器或实验相互作用,其中底层系统是不可求导的。
黑盒进修场景:当深度进修(Deep Learning)模型与第三方 API 集成时,如针对黑盒深度进修模型的对抗性攻击和防卫,以及语言模型服务的黑盒提示进修。
硬件限制:用于计算一门路度的原理性反向传播(backpropagation)机制在硬件系统上实现深度进修模型时可能不受支持。
然而,目前零阶优化的可扩展性仍然是一个未解决的问题:其运用主要限于相对较小规模的机器进修问题,如样本级的对抗性攻击生成。随着问题维度的增加,传统零阶办法的准确性和效率会下降。这是因为基于零阶有限差分的梯度估计是一门路度的有偏预算,且在高维空间中偏差更加明显。这些挑拨激发了本文讨论的核心问题:如何扩展零阶优化使其可以训练深度进修模型?
2. 零门路度预算:RGE 还是 CGE?
零阶优化器仅通过提交输入和接收相应的函数值与目标函数进行交互。主要有两种梯度预算办法:坐标梯度预算(Coordinate Gradient Estimation, CGE)和随机梯度预算(Random Gradient Estimation, RGE),如下所示:
其中表示对优化变量(例如,神经网络的模型参数)的一门路度的预算。
在(RGE)中,表示随机扰动向量,例如,从标准高斯分布中抽取;是扰动大小(又称平滑参数);q 是用于获得有限差分的随机方向数。
在(CGE)中,表示标准基向量,提供了在对应坐标的偏导数的有限差分估计。
与 CGE 相比,RGE 具有可以减少函数评估次数的灵活性。尽管查询效率高,但 RGE 在从头开始训练深度模型时是否能提供令人满意的准确性仍不确定。为此,我们进行了调查,其中我们运用 RGE 和 CGE 对不同大小的小型卷积神经网络(CNN)在 CIFAR-10 上进行了训练。如下图所示,CGE 可以实现与一阶优化训练相当的测试精度,并显著优于 RGE,同时也比 RGE 具有更高的时间效率。
基于 CGE 在准确性和计算效率方面相对于 RGE 的优势,我们选择 CGE 作为首选的零门路度估计器。然而,CGE 的查询复杂性仍然是一个瓶颈,因为它随模型大小增加而扩大。
3. 零阶深度进修框架:DeepZero
据我们所知,之前的工作没有展示出 ZO 优化在训练深度神经网络(DNN)时不会显著降低性能的有效性。为了克服这一障碍,我们开发了 DeepZero,一种原理性零阶优化深度进修框架,可以将零阶优化扩展到从头开始的神经网络训练。
a) 零阶模型修剪(ZO-GraSP):一个随机初始化的密集神经网络往往包含一个高质量的稠密子网络。然而,大多数有效的修剪办法都包含模型训练作为中间步骤。因此,它们不适合通过零阶优化找到稠密性。为了解决上述挑拨,我们受到了无需训练的修剪办法的启发,称为初始化修剪。在这类办法中,梯度信号保留(GraSP)被选用,它是一种通过随机初始化网络的梯度流识别神经网络的稠密性先验的办法。
b) 稠密梯度:为了保留训练密集模型的准确性优势,在 CGE 中我们结合了梯度稠密性而不是权重稠密性。这确保了我们在权重空间中训练一个密集模型,而不是训练一个稠密模型。具体而言,我们利用 ZO-GraSP 确定可以捕获 DNN 可压缩性的逐层修剪比率(Layer-wise Pruning Ratios, LPRs),然后零阶优化可以通过不断迭代更新部分模型参数权重来训练密集模型,其中稠密梯度比率由 LPRs 确定。
c) 特征重用:由于 CGE 逐元素扰动每个参数,它可以重用紧接扰动层之前的特征,并执行剩余的前向传播操作,而不是从输入层开始。从经验上看,带有特征重用的 CGE 在训练时间上可以实现 2 倍以上的减少。
d) 前传并行化:CGE 支持模型训练的并行化。这种解耦特性使得通过分布式机器扩展前向传播成为可能,从而显著提高零阶训练速度。
4. 实验分析
a) 图象分类
在 CIFAR-10 数据集上,我们将 DeepZero 训练的 ResNet-20 与两种通过一阶优化训练的变体进行比较:
(1)通过一阶优化训练获得的密集 ResNet-20
(2)通过一阶优化训练通过 FO-GraSP 获得的稠密 ResNet-20
如下图所示,尽管在 80% 至 99% 的稠密区间中,与(1)相比,运用 DeepZero 训练的模型仍存在准确度差距。这突出了 ZO 优化用于深度模型训练的挑拨,其中高稠密度的实现是被期望的。值得注意的是,在 90% 至 99% 的稠密区间中,DeepZero 优于(2),展示了 DeepZero 中梯度稠密性相对于权重稠密性的优越性。
b) 黑箱防卫
当模型的所有者不愿意与防卫者共享模型细节时,会出现黑盒防卫问题。这对于运用一阶优化训练直接增强白盒模型的现有鲁棒性增强算法构成了挑拨。为了克服这一挑拨,ZO-AE-DS 被提出,在白盒去噪平滑(Denoised Smoothing, DS)防卫操作和黑盒图象分类器之间引入了自动编码器(AutoEncoder, AE),以解决 ZO 训练的维度挑拨。ZO-AE-DS 的缺点是难以扩展到高分辨率数据集(例如,ImageNet),因为运用 AE 会损害输入到黑盒图象分类器的图象的保真度,并导致较差的防卫性能。相比之下,DeepZero 可以直接进修与黑盒分类器集成的防卫操作,无需自动编码器。如下表所示,就认证准确率(Certified Accuracy, CA)而言 DeepZero 在所有输入扰动半径上始终优于 ZO-AE-DS。
c) 与仿真耦合的深度进修
数值办法在提供物理信息仿照方面不可或缺,但它们自身存在挑拨:离散化不可避免地产生数值误差。通过与迭代偏微分方程(Partial Differential Equation, PDE)求解器的循环交互训练纠正神经网络的可行性,被称为” 求解器环路”(Solver-in-the-Loop, SOL)。虽然现有工作专注于运用或开发可微仿照器进行模型训练,我们通过利用 DeepZero 扩展了 SOL,使其能够与不可微或黑盒仿照器一起运用。下表比较了 ZO-SOL(通过 DeepZero 实现)与三种不同的可微办法的测试误差纠正性能:
(1) SRC(低保真仿照无误差纠正);
(2) NON(非交互式训练,运用预生成的低和高保真仿照数据在仿照循环外进行);
(3) FO-SOL(给定可微仿照器时,用于 SOL 的一阶训练)。
每个测试仿照的误差计算为与高保真仿照相比的纠正仿照的平均绝对误(MAE)。结果表明,通过 DeepZero 实现的 ZO-SOL 在只有基于查询的仿照器访问权限的情况下依然优于 SRC 和 NON,并缩小了与 FO-SOL 的性能差距。与 NON 相比,ZO-SOL 的表现突显了在有黑盒仿照器集成时的 ZO-SOL 前景。
5. 总结与讨论
这篇论文介绍了一个深度网络训练中零阶优化深度进修框架 (DeepZero)。具体来说,DeepZero 将坐标梯度估计、零阶模型修剪带来的梯度稠密性、特征重用以及前传并行化整合到统一的训练流程中。利用这些创新,DeepZero 在包括图象分类任务和各种实际黑箱深度进修场景中表现出了效率和有效性。此外,还探索了 DeepZero 在其他领域的适用性,如涉及不可微物理实体的应用,以及在计算图和反向传播的计算不被支持的设备上进行训练。
作者介绍
张益萌,密歇根州⽴⼤学 OPTML 实验室, 计算机博士在读, 研究兴趣⽅向包括 Generative AI, Multi-Modality, Computer Vision, Safe AI, Efficient AI。