[论文解读] Regularizing activations in neural networks via distribution matching with the Wasserstein metric

Taejong Joo, Donggu Kang+1 more|arXiv (Cornell University)|Feb 13, 2020

Advanced Neural Network Applications参考文献 56被引用 8

一句话总结

该论文提出了一种名为投影误差函数正则化（Projected Error Function Regularization, PER）的新方法，通过最小化神经网络激活分布与其标准正态分布之间1- Wasserstein距离的上界来正则化激活。通过将激活投影到随机的一维方向上，并应用类似Pseudo-Huber的损失函数，PER能够捕捉高阶统计量和单元间相关性，从而在不增加参数或改变网络架构的前提下，提升训练稳定性和泛化性能。

ABSTRACT

Regularization and normalization have become indispensable components in training deep neural networks, resulting in faster training and improved generalization performance. We propose the projected error function regularization loss (PER) that encourages activations to follow the standard normal distribution. PER randomly projects activations onto one-dimensional space and computes the regularization loss in the projected space. PER is similar to the Pseudo-Huber loss in the projected space, thus taking advantage of both $L^1$ and $L^2$ regularization losses. Besides, PER can capture the interaction between hidden units by projection vector drawn from a unit sphere. By doing so, PER minimizes the upper bound of the Wasserstein distance of order one between an empirical distribution of activations and the standard normal distribution. To the best of the authors' knowledge, this is the first work to regularize activations via distribution matching in the probability distribution space. We evaluate the proposed method on the image classification task and the word-level language modeling task.

研究动机与目标

• 为解决深层网络训练过程中激活分布不稳定的难题，该问题可能阻碍收敛和泛化。
• 克服现有方法仅正则化一阶和二阶统计量（如均值和方差）的局限性，通过捕捉高阶矩和单元间相关性来改进。
• 探索利用Wasserstein度量这一理论坚实且鲁棒的距离度量，在概率分布空间中进行正则化的潜力。
• 开发一种无需参数、与网络架构无关的正则化方法，以在多种任务中提升训练稳定性和模型性能。

提出的方法

• 将每一层的激活投影到从单位球面上随机采样的单维方向上，以实现在低维空间中的分布分析。
• 在投影空间中应用Pseudo-Huber损失，该损失结合了L1和L2正则化的优点，既平滑又对异常值具有鲁棒性。
• 最小化激活经验分布与标准正态分布N(0, I)之间1-Wasserstein距离的上界。
• 通过蒙特卡洛采样随机投影，近似计算切片Wasserstein距离，以高效估计Wasserstein距离。
• 将正则化损失表述为投影激活及其与目标高斯分布偏差的函数。
• 通过避免可学习参数并保持可微性，确保方法计算高效且与标准训练流程兼容。

实验结果

研究问题

• RQ1 对激活的完整分布进行正则化（而非仅其矩）是否能提升深层神经网络的训练稳定性和泛化性能？
• RQ2 与传统的归一化和正则化技术相比，基于Wasserstein度量的分布匹配在性能和鲁棒性方面表现如何？
• RQ3 PER在无需显式归一化层或批量统计信息的情况下，能在多大程度上稳定训练过程中的激活分布？
• RQ4 使用随机投影和Pseudo-Huber损失是否能有效捕捉隐藏表示中的单元间相关性和高阶统计量？
• RQ5 PER是否可在无需架构修改的前提下，有效应用于图像分类和语言建模等多种任务？

主要发现

• PER在训练过程中显著稳定了激活的分布，有效防止了均值和方差漂移，尤其在训练初期表现突出。
• 与原始模型和批归一化（BN）基线相比，PER下激活分布与标准正态分布之间的1-Wasserstein距离始终更低。
• 在图像分类（CIFAR-10、CIFAR-100）和词级别语言建模（WikiText-2）任务中，PER实现了微小但一致的测试准确率提升。
• 尽管训练时间略有增加（每轮迭代0.093秒 vs. 原始模型的0.071秒），PER不引入额外参数，且与标准训练流程完全兼容。
• 分析表明，PER有效控制了高阶统计量和单元间相关性，表现为分布漂移减少和收敛速度加快。
• 即使未使用批归一化或显式矩控制，PER在保持激活分布接近目标分布N(0, I)方面，仍优于BN和VCL方法。