[论文解读] projUNN: efficient method for training deep networks with unitary matrices

Bobak T. Kiani, Randall Balestriero+2 more|arXiv (Cornell University)|Mar 10, 2022

Domain Adaptation and Few-Shot Learning被引用 11

一句话总结

该论文提出 projUNN，一种通过秩-$k$ 更新实现单位矩阵高效训练的深度网络方法，通过投影或传输保持单位性。其训练复杂度接近最优的 $O(kN^2)$，与当前最优的单位矩阵 RNN 相当或更优，并支持正交卷积网络的稳定、可扩展训练。

ABSTRACT

In learning with recurrent or very deep feed-forward networks, employing unitary matrices in each layer can be very effective at maintaining long-range stability. However, restricting network parameters to be unitary typically comes at the cost of expensive parameterizations or increased training runtime. We propose instead an efficient method based on rank-$k$ updates -- or their rank-$k$ approximation -- that maintains performance at a nearly optimal training runtime. We introduce two variants of this method, named Direct (projUNN-D) and Tangent (projUNN-T) projected Unitary Neural Networks, that can parameterize full $N$-dimensional unitary or orthogonal matrices with a training runtime scaling as $O(kN^2)$. Our method either projects low-rank gradients onto the closest unitary matrix (projUNN-T) or transports unitary matrices in the direction of the low-rank gradient (projUNN-D). Even in the fastest setting ($k=1$), projUNN is able to train a model's unitary parameters to reach comparable performances against baseline implementations. In recurrent neural network settings, projUNN closely matches or exceeds benchmarked results from prior unitary neural networks. Finally, we preliminarily explore projUNN in training orthogonal convolutional neural networks, which are currently unable to outperform state of the art models but can potentially enhance stability and robustness at large depth.

研究动机与目标

• 为解决使用单位矩阵训练深度或循环网络时计算成本过高的问题，此类方法通常每步梯度计算需 $O(N^3)$ 次操作。
• 在保持低训练运行时间的前提下，实现 $N\times N$ 单位矩阵的完整参数化，尤其适用于大 $N$ 的情况。
• 通过避免昂贵的 SVD 或矩阵求逆步骤，保持单位性，从而提升深度网络的稳定性和可扩展性。
• 将单位矩阵/正交矩阵的应用扩展至卷积网络，特别是针对深度架构。

提出的方法

• projUNN 使用秩-$k$ 更新来近似梯度，将单位矩阵更新的计算成本从 $O(N^3)$ 降低至 $O(kN^2)$。
• Direct 变体（projUNN-D）通过低秩更新子空间的特征分解，将更新后的矩阵投影到最近的单位矩阵上。
• Tangent 变体（projUNN-T）使用指数映射更新，沿低秩梯度方向对单位矩阵进行传输。
• 两种变体均对完整单位矩阵流形进行参数化，并在每次梯度更新后强制保持单位性，确保数值稳定性。
• 对于卷积层，projUNN 在傅里叶域中使用 FFT 执行正交卷积，保持厄米特对称性，实现高效计算。
• 为单位矩阵/正交参数设计专用优化器，学习率通常设置为标准参数的 1/10 或 1/20。

实验结果

研究问题

• RQ1 是否可以使用低秩梯度近似在不损失性能的前提下高效训练深度网络中的单位矩阵？
• RQ2 所提方法是否在保持单位矩阵流形完整参数化的同时，实现接近最优的训练复杂度？
• RQ3 在长序列任务中，projUNN 与现有单位矩阵 RNN 相比，在准确率和训练速度方面表现如何？
• RQ4 projUNN 是否可有效扩展至正交卷积网络，以提升深度架构的稳定性？
• RQ5 与哈尔回收或单位矩阵初始化相比，哪些初始化方案能带来更好的收敛性和稳定性？

主要发现

• projUNN 每步梯度计算的复杂度达到 $O(kN^2)$，显著降低了先前基于李代数方法的 $O(N^3)$ 成本。
• 即使 $k=1$，projUNN 在基准序列学习任务上仍能匹配或超越当前最优的单位矩阵 RNN 性能。
• Tangent 变体（projUNN-T）在长时间训练中表现出优于 Direct 变体（projUNN-D）的数值稳定性。
• projUNN 实现了正交卷积网络的高效训练，为极深模型的稳定性提升提供了可行路径。
• 实验结果表明，使用 $2\times2$ 斜对称块（Henaff 风格）的初始化优于单位矩阵或哈尔回收初始化。
• 每隔 $O(N)$ 步进行一次稀疏的完整投影，可在不降低平均 $O(kN^2)$ 运行时间的前提下维持单位性。