[论文解读] Understanding Approximate Fisher Information for Fast Convergence of Natural Gradient Descent in Wide Neural Networks

Ryo Karakida, Kazuki Osawa|arXiv (Cornell University)|Oct 2, 2020

Stochastic Gradient Optimization Techniques参考文献 30被引用 15

一句话总结

该论文建立了理论基础，表明在宽的全连接神经网络中，使用各种近似Fisher信息矩阵（如分块对角、分块三对角、K-FAC和单位级近似）的自然梯度下降（NGD）方法，其收敛速度与精确NGD达到全局最小值的速度相同。关键洞见在于，这些近似方法在函数空间中产生各向同性梯度，从而实现快速收敛，且该收敛速度与神经正切核（NTK）无关，即使在分层或单位级近似下，只要满足特定的网络宽度和学习率条件，也能实现快速收敛。

ABSTRACT

Natural Gradient Descent (NGD) helps to accelerate the convergence of gradient descent dynamics, but it requires approximations in large-scale deep neural networks because of its high computational cost. Empirical studies have confirmed that some NGD methods with approximate Fisher information converge sufficiently fast in practice. Nevertheless, it remains unclear from the theoretical perspective why and under what conditions such heuristic approximations work well. In this work, we reveal that, under specific conditions, NGD with approximate Fisher information achieves the same fast convergence to global minima as exact NGD. We consider deep neural networks in the infinite-width limit, and analyze the asymptotic training dynamics of NGD in function space via the neural tangent kernel. In the function space, the training dynamics with the approximate Fisher information are identical to those with the exact Fisher information, and they converge quickly. The fast convergence holds in layer-wise approximations; for instance, in block diagonal approximation where each block corresponds to a layer as well as in block tri-diagonal and K-FAC approximations. We also find that a unit-wise approximation achieves the same fast convergence under some assumptions. All of these different approximations have an isotropic gradient in the function space, and this plays a fundamental role in achieving the same convergence properties in training. Thus, the current study gives a novel and unified theoretical foundation with which to understand NGD methods in deep learning.

研究动机与目标

• 为解决在自然梯度下降（NGD）中，尽管计算成本高，为何启发式Fisher信息矩阵（FIM）近似在实践中表现良好这一理论空白。
• 分析在深度神经网络的无限宽极限下，NGD与近似FIM的渐近训练动态。
• 明确不同FIM近似（如分块对角、K-FAC、单位级）在何种精确条件下能保持与精确NGD相同的快速收敛性。
• 阐明函数空间梯度的各向同性作为各类NGD近似中实现快速收敛的统一原理。

提出的方法

• 该研究采用神经正切核（NTK）框架，分析无限宽极限下宽的全连接神经网络在函数空间中的训练动态。
• 推导了使用近似FIM的NGD的渐近动态，并证明在特定条件下其与精确NGD在函数空间中等价。
• 分析聚焦于分层近似（分块对角、分块三对角、K-FAC）和单位级近似，证明当学习率按网络宽度或样本量适当缩放时，其收敛行为与精确NGD完全相同。
• 论文引入一个阻尼参数 ρ > 0 以稳定FIM的逆，并推导了线性化动态与真实动态之间的偏差界限，证明在 M → ∞ 时收敛。
• 确立了实现快速收敛的关键机制是函数空间中梯度的各向同性，该特性源于近似FIM的结构。
• 数值实验通过将理论收敛速率与实际训练动态对比，验证了理论预测，尤其针对单位级NGD。

实验结果

研究问题

• RQ1 在何种条件下，使用分块或单位级FIM近似的NGD方法在宽的神经网络中能实现与精确NGD相同的快速收敛？
• RQ2 为何实践中FIM近似（如K-FAC或分块对角）虽存在理论不确定性，却能表现良好？
• RQ3 FIM近似中何种结构特性可确保函数空间中的快速收敛？其与NTK有何关联？
• RQ4 不同FIM近似如何在函数空间中诱导梯度的各向同性？为何该特性对快速收敛至关重要？

主要发现

• 在深度神经网络的无限宽极限下，使用分块对角、分块三对角、K-FAC和单位级FIM近似的近似NGD，其收敛速度与精确NGD达到全局最小值的速度相同。
• 只要学习率按网络宽度或样本量适当缩放，所有这些近似方法与精确NGD在函数空间中的收敛动态完全一致。
• 实现快速收敛的关键机制是函数空间中梯度的各向同性，该特性由近似FIM的结构所诱导，且与NTK无关。
• 单位级NGD在阻尼参数 ρ > 0 为小但非零时可实现快速收敛，其收敛速率受 A³ρ⁻⁶/√M 限制，且当 M → ∞ 时该速率趋于零。
• 数值实验表明，各向同性条件在分层和单位级近似中成立，但在逐元素对角近似中不成立，从而解释了后者的性能不佳。
• 尽管函数空间动态相同，不同近似方法在参数空间中的训练动态不同，导致收敛到不同的全局最小值和测试预测。