[论文解读] Fast Differentiable Matrix Square Root

Song Yue, Nicu Sebe+1 more|arXiv (Cornell University)|Jan 21, 2022

Advanced Vision and Imaging被引用 5

一句话总结

本文提出两种快速可微分矩阵平方根方法——矩阵泰勒多项式（MTP）与矩阵有理逼近（MPA），用于高效前向计算，结合基于李雅普诺夫方程的迭代求解器实现反向传播。该方法在保持视觉任务中竞争力表现的同时，相较SVD与牛顿-舒尔茨迭代实现了显著的速度提升。

ABSTRACT

Computing the matrix square root or its inverse in a differentiable manner is important in a variety of computer vision tasks. Previous methods either adopt the Singular Value Decomposition (SVD) to explicitly factorize the matrix or use the Newton-Schulz iteration (NS iteration) to derive the approximate solution. However, both methods are not computationally efficient enough in either the forward pass or in the backward pass. In this paper, we propose two more efficient variants to compute the differentiable matrix square root. For the forward propagation, one method is to use Matrix Taylor Polynomial (MTP), and the other method is to use Matrix Padé Approximants (MPA). The backward gradient is computed by iteratively solving the continuous-time Lyapunov equation using the matrix sign function. Both methods yield considerable speed-up compared with the SVD or the Newton-Schulz iteration. Experimental results on the de-correlated batch normalization and second-order vision transformer demonstrate that our methods can also achieve competitive and even slightly better performances. The code is available at \href{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}.

研究动机与目标

• 解决SVD与牛顿-舒尔茨迭代在深度学习框架中，尤其是批量矩阵平方根运算时的计算低效问题。
• 克服因特征值差异导致的特征分解在反向传播中不稳定且高成本的问题。
• 开发适用于GPU、可微分且兼具高速与数值稳定性的矩阵平方根方法。
• 实现矩阵平方根在视觉应用（如协方差池化与归一化层）中的高效使用。

提出的方法

• 在前向传播中使用矩阵泰勒多项式（MTP）与矩阵有理逼近（MPA）近似矩阵平方根，替代SVD或迭代方法。
• 基于矩阵符号函数设计迭代李雅普诺夫求解器，高效计算反向传播中的梯度。
• 将反向传播建模为求解连续时间李雅普诺夫方程，避免显式特征分解。
• 在深度学习框架中集成该方法，计算开销极小，充分利用GPU上高度优化的矩阵乘法运算。
• 对输入矩阵进行预归一化，以提升迭代组件的收敛稳定性。
• 支持混合精度训练，尤其在大模型中进一步提升效率。

实验结果

研究问题

• RQ1 MTP与MPA能否为深度学习中的矩阵平方根计算提供一种比SVD更快且可微分的替代方案？
• RQ2 所提出的基于李雅普诺夫的反向求解器在梯度计算速度上是否优于传统的SVD与牛顿-舒尔茨迭代？
• RQ3 在ZCA白化与二阶视觉Transformer等真实视觉基准测试中，所提方法在准确率与速度上的表现如何？
• RQ4 在大模型中，该方法能否在减少训练时间与显存占用的同时保持或提升性能？

主要发现

• 在So-ViT-14模型上，我们的MPA-Lya方法相较牛顿-舒尔茨迭代快3.19倍，相较SVD-Padé快25.63倍。
• MTP-Lya方法在前向传播中速度最快，单次矩阵平方根计算仅需2.36毫秒，优于SVD与NS迭代。
• 在CIFAR-100数据集上，MPA-Lya实现19.55% ± 0.20的验证误差，推理时间仅2.52毫秒，速度与准确率均优于SVD与NS迭代。
• 在二阶视觉Transformer中，MPA-Lya仅需254个训练周期即可收敛，少于基于SVD的方法，且在top-1准确率上保持相当或略优。
• MTP-Lya在ResNet-18上性能与SVD和NS相当，但在ResNet-50上略有下降，可能源于近似误差的累积。
• 基于李雅普诺夫的反向求解器相较NS迭代减少了矩阵乘法次数，从而实现更快的反向传播。