[论文解读] Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves

Luke Metz, Niru Maheswaranathan+3 more|arXiv (Cornell University)|Sep 23, 2020

Machine Learning and Data Classification参考文献 67被引用 14

一句话总结

本论文提出了一种可扩展的、通用的元优化器，其在数以千计的多样化机器学习任务上使用大规模计算资源进行训练。该优化器采用分层神经网络架构，可访问验证损失及其他任务特征，从而实现稳健的泛化能力和自训练，其在稳定性、适应性和分布外性能方面优于先前方法。

ABSTRACT

Much as replacing hand-designed features with learned functions has revolutionized how we solve perceptual tasks, we believe learned algorithms will transform how we train models. In this work we focus on general-purpose learned optimizers capable of training a wide variety of problems with no user-specified hyperparameters. We introduce a new, neural network parameterized, hierarchical optimizer with access to additional features such as validation loss to enable automatic regularization. Most learned optimizers have been trained on only a single task, or a small number of tasks. We train our optimizers on thousands of tasks, making use of orders of magnitude more compute, resulting in optimizers that generalize better to unseen tasks. The learned optimizers not only perform well, but learn behaviors that are distinct from existing first order optimizers. For instance, they generate update steps that have implicit regularization and adapt as the problem hyperparameters (e.g. batch size) or architecture (e.g. neural network width) change. Finally, these learned optimizers show evidence of being useful for out of distribution tasks such as training themselves from scratch.

研究动机与目标

• 通过将训练规模扩展至数千个多样化任务，克服先前元优化器的脆弱性与泛化能力差的问题。
• 通过引入验证损失和架构元数据等额外任务信号，提升优化器的鲁棒性。
• 设计一种具有强归纳偏置的分层神经网络架构，以实现更优的性能与泛化能力。
• 使元优化器能够从零开始训练新的优化器，展示类似于自托管编译器的自宿主能力。
• 建立一个可扩展的、高算力的流水线，用于训练与评估元优化器，实现低方差与高保真度。

提出的方法

• 通过双层优化训练优化器：外层循环基于多个内层任务的性能来更新优化器参数。
• 采用分层神经网络架构，处理梯度、验证损失和模型超参数，以生成自适应的参数更新。
• 训练数据集包含超过1,000个多样化的机器学习任务，涵盖批量大小、网络宽度和数据分布的变化。
• 使用约60,000个CPU核心进行大规模训练，持续一个月（约5,000 CPU年），以支持长序列展开与稳定的外层损失估计。
• 通过分布式评估系统（包含评估主控与工作节点）在保留任务上运行完整的训练轨迹，以衡量泛化性能。
• 通过RPC进行摘要聚合，结合批处理与随机日志记录，以最小化开销并支持事后分析。

实验结果

研究问题

• RQ1 元优化器是否能在数千个多样化机器学习任务上实现有效泛化，而不仅限于少数任务？
• RQ2 引入验证损失与架构元数据是否能提升元优化器的稳定性与性能？
• RQ3 元优化器是否能从零开始训练新的优化器，从而展示自宿主能力？
• RQ4 算力与任务多样性的扩展如何影响元优化器的泛化能力与鲁棒性？
• RQ5 此类系统的计算与环境成本如何？是否可进一步提升效率？

主要发现

• 由于在超过1,000项多样化任务的分布上进行训练，该元优化器能有效泛化到未见过的任务，包括不同批量大小与网络宽度的任务。
• 该优化器通过其更新规则隐式学习正则化，能够动态适应超参数与架构的变化，这与标准的一阶优化器不同。
• 该模型成功实现了从零开始训练新的元优化器，展示了自宿主能力，这是类似于自托管编译器的关键里程碑。
• 在分布外任务上的评估表现出色，证实了优化器在元训练分布之外的鲁棒性。
• 训练过程约消耗60,000个CPU核心运行一个月，耗电量约为200兆瓦时，凸显了其高昂的计算成本。
• 尽管计算成本高昂，作者认为，一旦训练完成，此类优化器可消除下游训练中的超参数调优需求，从而实现长期成本节约。