面向现代神经网络架构的Kronecker因子近似曲率
机器学习
2024-01-12 v2 机器学习
摘要
许多现代神经网络架构（如Transformer、卷积神经网络或图神经网络）的核心组件可表示为具有权重共享的线性层。Kronecker因子近似曲率（K-FAC）作为一种二阶优化方法，已显示出加速神经网络训练从而降低计算成本的潜力。然而，目前尚无将其应用于通用架构（特别是具有线性权重共享层的架构）的框架。在本工作中，我们识别出线性权重共享层的两种不同设置，由此引出K-FAC的两种变体——expand与reduce。我们证明它们在各自设置下对带权重共享的深度线性网络是精确的。值得注意的是，K-FAC-reduce通常比K-FAC-expand更快，我们利用这一点通过优化Wide ResNet的边缘似然来加速自动超参数选择。最后，我们观察到在用于训练图神经网络和视觉Transformer时，这两种K-FAC变体之间差异很小。然而，两种变体均能以一阶参考运行步数的50%-75%达到固定的验证指标目标，这转化为可比较的挂钟时间改进。这凸显了将K-FAC应用于现代神经网络架构的潜力。
引用
@article{arxiv.2311.00636,
title = {Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures},
author = {Runa Eschenhagen and Alexander Immer and Richard E. Turner and Frank Schneider and Philipp Hennig},
journal= {arXiv preprint arXiv:2311.00636},
year = {2024}
}
备注
NeurIPS 2023