[论文解读] Riemannian approach to batch normalization

Minhyung Cho, Jaehyeong Lee|arXiv (Cornell University)|Sep 27, 2017

Domain Adaptation and Few-Shot Learning参考文献 20被引用 14

一句话总结

该论文通过将尺度不变的权重向量建模为Grassmann流形𝒢(1,n)上的点，提出了一种用于批量归一化（BN）的黎曼优化框架，实现了几何感知的梯度下降。该方法通过消除因权重缩放引起的梯度模糊性，在多个模型架构和数据集上提升了训练稳定性和性能，相较于使用BN的标准SGD和Adam方法，在CIFAR-100和SVHN上实现了最高达1.5%的误差降低。

ABSTRACT

Batch Normalization (BN) has proven to be an effective algorithm for deep neural network training by normalizing the input to each neuron and reducing the internal covariate shift. The space of weight vectors in the BN layer can be naturally interpreted as a Riemannian manifold, which is invariant to linear scaling of weights. Following the intrinsic geometry of this manifold provides a new learning rule that is more efficient and easier to analyze. We also propose intuitive and effective gradient clipping and regularization methods for the proposed algorithm by utilizing the geometry of the manifold. The resulting algorithm consistently outperforms the original BN on various types of network architectures and datasets.

研究动机与目标

• 解决批量归一化（BN）中因权重缩放导致的梯度更新不稳定与模糊性问题，该问题影响优化收敛性及对权重衰减的敏感性。
• 将BN权重向量的空间重新表述为黎曼流形（特别是Grassmann流形𝒢(1,n)），以内在方式强制实现尺度不变性。
• 开发尊重流形内在几何结构的SGD带动量和Adam的黎曼版本，确保优化过程的稳定与高效。
• 引入一种基于流形结构的几何梯度裁剪方法，以及一种基于变分推理的正则化技术。
• 在多种架构和数据集（包括在CIFAR-10、CIFAR-100和SVHN上的VGG与Wide ResNet）上，持续展示优于标准BN训练的性能提升。

提出的方法

• 将BN权重向量的空间建模为Grassmann流形𝒢(1,n)，其中每个点代表权重空间中的一个方向，对标量缩放保持不变。
• 在𝒢(1,n)上定义一个黎曼度量张量，以保持BN变换在权重缩放下的不变性，从而支持内在的梯度下降。
• 通过使用指数映射和并行传输，将欧氏空间的更新规则适配到流形上，推导出黎曼优化算法——带动量的黎曼SGD与黎曼Adam。
• 提出一种基于流形切空间的几何梯度裁剪方法，可在保持不变性的同时约束梯度更新。
• 提出一种基于变分推理的正则化方法，将权重向量建模为𝒢(1,n)上的随机元素，以提升泛化能力。
• 实现算法时仅引入极小的计算开销：前向与反向传播保持不变，仅权重更新步骤通过黎曼运算进行修改。

实验结果

研究问题

• RQ1 如何利用批量归一化中尺度不变权重向量的内在几何结构来提升优化稳定性？
• RQ2 在Grassmann流形𝒢(1,n)上进行黎曼优化，能否消除BN中因权重缩放导致的梯度更新模糊性？
• RQ3 当应用于多种网络架构的BN层时，几何感知优化方法是否优于标准SGD和Adam？
• RQ4 在流形上应用几何梯度裁剪与正则化，能否提升泛化能力并增强对学习率衰减等超参数的鲁棒性？
• RQ5 所提出方法在多大程度上降低了对权重衰减的敏感性，同时保持或提升测试准确率？

主要发现

• 所提出的黎曼优化方法（SGD-G与Adam-G）在所有评估的数据集和架构上，均持续优于标准SGD与Adam在BN层上的表现。
• 在WRN-40-10架构上，CIFAR-100数据集上，Adam-G变体实现了17.59%的测试误差，优于使用标准BN和Adam的最优报告结果18.30%。
• 在SVHN数据集上，WRN-22-8架构下，Adam-G方法将误差降低至1.49%，优于此前报告的Wide ResNet最优结果1.54%。
• 该方法消除了权重衰减与梯度缩放之间的有害交互，解决了BN训练中已知的不稳定性问题。
• 计算成本增加极小——每次更新仅增加2.5至3.5倍，使其在大规模深度学习中具有实际可行性。
• 该框架可推广至其他具有尺度不变权重的归一化技术，如权重归一化与归一化传播。