[论文解读] CWY Parametrization for Scalable Learning of Orthogonal and Stiefel Matrices

Valerii Likhosherstov, Jared Quincy Davis+2 more|arXiv (Cornell University)|Apr 18, 2020

Advanced Graph Neural Networks被引用 3

一句话总结

本文提出了CWY和T-CWY参数化方法，通过利用紧凑且可并行化的Householder反射表示，实现了在GPU和TPU上对正交矩阵和Stiefel流形的高效优化。该方法在随机梯度下降下可收敛至驻点，并在神经机器翻译和视频预测的RNN训练中表现出色。

ABSTRACT

We introduce an efficient approach for optimization over orthogonal groups on highly parallel computation units such as GPUs or TPUs. As in earlier work, we parametrize an orthogonal matrix as a product of Householder reflections. However, to overcome low parallelization capabilities of computing Householder reflections sequentially, we propose employing an accumulation scheme called the compact WY (or CWY) transform -- a compact parallelization-friendly matrix representation for the series of Householder reflections. We further develop a novel Truncated CWY (or T-CWY) approach for Stiefel manifold parametrization which has a competitive complexity and, again, yields benefits when computed on GPUs and TPUs. We prove that our CWY and T-CWY methods lead to convergence to a stationary point of the training objective when coupled with stochastic gradient descent. We apply our methods to train recurrent neural network architectures in the tasks of neural machine translation and video prediction.

研究动机与目标

• 解决正交矩阵优化中顺序Householder反射计算导致的并行化效率低下问题。
• 实现在现代并行硬件（如GPU和TPU）上对正交和Stiefel流形的可扩展、高性能优化。
• 开发一种紧凑且适合并行计算的矩阵表示方法，同时保持数值稳定性和收敛性。
• 通过一种新颖的截断CWY（T-CWY）方法将框架扩展至Stiefel流形，计算复杂度具有竞争力。
• 在真实世界深度学习任务（如神经机器翻译和视频预测）中对方法进行实证验证。

提出的方法

• 本文使用Householder反射对正交矩阵进行参数化，将顺序计算替换为紧凑的WY（CWY）表示，以支持并行执行。
• CWY变换被推导为矩阵乘积，以紧凑且数值稳定的形式累积Householder反射，适用于GPU和TPU加速。
• 提出一种新颖的截断CWY（T-CWY）方法用于Stiefel流形的参数化，降低计算成本的同时保持收敛性。
• 该方法与随机梯度下降结合，理论证明在标准假设下可收敛至目标函数的驻点。
• 通过参数化隐式维持正交性约束，优化过程中无需显式投影步骤。
• 该框架被应用于循环神经网络架构，优化直接在正交或Stiefel矩阵的流形上进行。

实验结果

研究问题

• RQ1 紧凑且可并行化的Householder反射表示能否提升GPU和TPU上的优化效率？
• RQ2 在与随机梯度下降结合时，CWY参数化是否能保持收敛性保证？
• RQ3 T-CWY方法能否以更低的计算成本有效参数化Stiefel流形？
• RQ4 在序列建模任务中，该方法与现有方法相比在训练速度和模型性能方面表现如何？
• RQ5 CWY和T-CWY方法对神经机器翻译和视频预测中RNN的训练动态有何影响？

主要发现

• CWY参数化实现了在GPU和TPU等高度并行硬件上对正交矩阵的高效可扩展优化。
• 该方法与随机梯度下降结合时，可保证收敛至训练目标的驻点。
• T-CWY方法为Stiefel流形参数化提供了计算复杂度具有竞争力的替代方案，并具备强大的并行化支持。
• 实证结果表明，该方法在神经机器翻译和视频预测任务的RNN训练中显著提升了效率和性能。
• 紧凑的WY表示显著减少了顺序Householder反射带来的计算瓶颈，使并行架构上的收敛速度更快。
• 该框架在优化过程中无需显式投影步骤即可保持数值稳定性和正交性约束。