[论文解读] Identifying Equivalent Training Dynamics

William T. Redman, Juan M. Bello‐Rivas+4 more|arXiv (Cornell University)|Feb 17, 2023

Model Reduction and Neural Networks被引用 6

一句话总结

本文提出一种基于数据的框架，利用Koopman算子理论通过分析优化轨迹的谱相似性，识别机器学习中的等价训练动态。结果表明，共轭动态（即等价优化）可通过重叠的Koopman谱检测到，揭示了学习率与批量大小之比、网络宽度、数据集类型及激活函数在前馈神经网络中对共轭性的控制作用。

ABSTRACT

Study of the nonlinear evolution deep neural network (DNN) parameters undergo during training has uncovered regimes of distinct dynamical behavior. While a detailed understanding of these phenomena has the potential to advance improvements in training efficiency and robustness, the lack of methods for identifying when DNN models have equivalent dynamics limits the insight that can be gained from prior work. Topological conjugacy, a notion from dynamical systems theory, provides a precise definition of dynamical equivalence, offering a possible route to address this need. However, topological conjugacies have historically been challenging to compute. By leveraging advances in Koopman operator theory, we develop a framework for identifying conjugate and non-conjugate training dynamics. To validate our approach, we demonstrate that comparing Koopman eigenvalues can correctly identify a known equivalence between online mirror descent and online gradient descent. We then utilize our approach to: (a) identify non-conjugate training dynamics between shallow and wide fully connected neural networks; (b) characterize the early phase of training dynamics in convolutional neural networks; (c) uncover non-conjugate training dynamics in Transformers that do and do not undergo grokking. Our results, across a range of DNN architectures, illustrate the flexibility of our framework and highlight its potential for shedding new light on training dynamics.

研究动机与目标

• 为解决当前缺乏通用方法来判断机器学习训练中不同超参数选择是否导致等价优化轨迹的问题。
• 实现基于其底层动力学行为（特别是共轭性）对机器学习优化算法的分类。
• 提供一种数据驱动且可泛化的通用方法，用于在无需访问内部算法方程的情况下，比较多种机器学习方法的训练动态。
• 将Koopman算子理论扩展至迭代式机器学习优化过程的分析，尤其适用于深度神经网络。
• 揭示哪些超参数（如学习率、批量大小、宽度、激活函数和数据集）在训练过程中引发共轭或非共轭动态。

提出的方法

• 将迭代式机器学习优化建模为离散时间动力系统，其中每个更新步骤对应一个状态转移映射。
• 应用Koopman模态分解，从观测到的训练轨迹中提取谱对象（特征值与模态）。
• 以谱相似性（特别是重叠的Koopman特征值）作为标准，识别共轭动力系统，从而表明优化动态等价。
• 采用数据驱动的Koopman分析方法，比较在线镜像下降与梯度下降，以数据驱动方式验证其分析等价性。
• 在全连接网络中系统性地改变超参数（学习率、批量大小、宽度、激活函数、数据集），以评估其对Koopman谱的影响。
• 使用多个随机种子之间Koopman谱的Wasserstein距离，量化谱相似性的统计鲁棒性。

Figure 1: Data-driven identification of the conjugacy between online mirror and gradient descent. (A) Example trajectory for $x_{t}$ (red) and $q(u_{t})$ (black), when $f(x)=\sum_{i=1}^{d}x_{i}^{4}$ . Rhombus denotes initialization and star denotes the end of $500$ iterations. (B) Koopman spectra co

实验结果

研究问题

• RQ1 在何种条件下，深度神经网络训练中的不同超参数配置会导致等价的优化动态？
• RQ2 Koopman谱相似性能否可靠检测出不同优化算法（如在线镜像下降与梯度下降）之间的共轭性？
• RQ3 哪些超参数——学习率、批量大小、层宽、数据集、激活函数——最显著地影响训练动态的共轭结构？
• RQ4 在无法访问底层算法方程的情况下，数据驱动的Koopman分析在多大程度上可识别等价的训练轨迹？
• RQ5 不同数据集（如MNIST与合成数据）和激活函数在多大程度上影响优化过程的谱结构与共轭性？

主要发现

• 在特定问题上，在线镜像下降与梯度下降之间观察到高度重叠的Koopman谱，证实了其分析等价性。
• 学习率与批量大小之比是决定优化动态是否共轭的关键因素。
• 网络宽度和数据集的性质（如手写数据与合成数据）显著影响训练动态的共轭结构。
• 激活函数的选择会影响Koopman算子的谱特性，从而影响优化轨迹是否等价。
• 通过Koopman模态分解的谱分析表明，即使训练性能相似，不同的超参数组合也可能导致非共轭动态。
• 该框架在多种场景中成功识别出共轭关系，包括求根算法与神经网络训练，展现出广泛的适用性。

Figure 2: Learning rate to batch size ratio controls the conjugacy of training. (A) Training loss, as a function of number of images shown, for different choices of batch size. Loss was first smoothed using a square filter. Width of curves denotes $25^{\text{th}}$ and $75^{\text{th}}$ percentile of