# Fast Differentiable Matrix Square Root and Inverse Square Root

arXiv: https://arxiv.org/abs/2201.12543 · HTML: https://mlsanity.com/paper/2201.12543/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 64 references that are in the map.
- First posted on arXiv 29 Jan 2022.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Fine-Grained Visual Classification of Aircraft](https://mlsanity.com/paper/1306.5151/) (2013)
- [Universal Style Transfer via Feature Transforms](https://mlsanity.com/paper/1705.08086/) (2017)
- [Whitening for Self-Supervised Representation Learning](https://mlsanity.com/paper/2007.06346/) (2020)
- [Bilinear CNNs for Fine-grained Visual Recognition](https://mlsanity.com/paper/1504.07889/) (2015)
- [Improved Texture Networks: Maximizing Quality and Diversity in Feed-forward Stylization and Texture Synthesis](https://mlsanity.com/paper/1701.02096/) (2017)
- [RobustNet: Improving Domain Generalization in Urban-Scene Segmentation via Instance Selective Whitening](https://mlsanity.com/paper/2103.15597/) (2021)
- [TEA: Temporal Excitation and Aggregation for Action Recognition](https://mlsanity.com/paper/2004.01398/) (2020)
- [A Closed-form Solution to Photorealistic Image Stylization](https://mlsanity.com/paper/1802.06474/) (2018)
- [Switchable Whitening for Deep Representation Learning](https://mlsanity.com/paper/1904.09739/) (2019)
- [Iterative Normalization: Beyond Standardization towards Efficient Whitening](https://mlsanity.com/paper/1904.03441/) (2019)
- [Decorrelated Batch Normalization](https://mlsanity.com/paper/1804.08450/) (2018)
- [Towards Faster Training of Global Covariance Pooling Networks by Iterative Matrix Square Root Normalization](https://mlsanity.com/paper/1712.01034/) (2017)
- [Image-to-Image Translation via Group-wise Deep Whitening-and-Coloring Transformation](https://mlsanity.com/paper/1812.09912/) (2018)
- [Is Second-order Information Helpful for Large-scale Visual Recognition?](https://mlsanity.com/paper/1703.08050/) (2017)
- [Deep CNNs Meet Global Covariance Pooling: Better Representation and Generalization](https://mlsanity.com/paper/1904.06836/) (2019)
- [Diversified Arbitrary Style Transfer via Deep Feature Perturbation](https://mlsanity.com/paper/1909.08223/) (2019)
- [Training Deep Networks with Structured Layers by Matrix Backpropagation](https://mlsanity.com/paper/1509.07838/) (2015)
- [Improved Bilinear Pooling with CNNs](https://mlsanity.com/paper/1707.06772/) (2017)
- [Whitening and Coloring batch transform for GANs](https://mlsanity.com/paper/1806.00420/) (2018)
- [Backpropagation-Friendly Eigendecomposition](https://mlsanity.com/paper/1906.09023/) (2019)
- [Why Approximate Matrix Square Root Outperforms Accurate SVD in Global Covariance Pooling?](https://mlsanity.com/paper/2105.02498/) (2021)
- [Robust Differentiable SVD](https://mlsanity.com/paper/2104.03821/) (2021)
- [Keep it Simple: Image Statistics Matching for Domain Adaptation](https://mlsanity.com/paper/2005.12551/) (2020)
- [Eigendecomposition-free Training of Deep Networks with Zero Eigenvalue-based Losses](https://mlsanity.com/paper/1803.08071/) (2018)

## Built on by (most cited first, 16 of 16)

- [Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Space](https://mlsanity.com/paper/2304.05398/) (2023)
- [SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training](https://mlsanity.com/paper/2412.13148/) (2024)
- [Adaptive Log-Euclidean Metrics for SPD Matrix Learning](https://mlsanity.com/paper/2303.15477/) (2023)
- [Riemannian Multinomial Logistics Regression for SPD Neural Networks](https://mlsanity.com/paper/2305.11288/) (2023)
- [Understanding Matrix Function Normalizations in Covariance Pooling through the Lens of Riemannian Geometry](https://mlsanity.com/paper/2407.10484/) (2024)
- [Gradient Multi-Normalization for Stateless and Scalable LLM Training](https://mlsanity.com/paper/2502.06742/) (2025)
- [SPD Matrix Learning for Neuroimaging Analysis: Perspectives, Methods, and Challenges](https://mlsanity.com/paper/2504.18882/) (2025)
- [Orthogonal SVD Covariance Conditioning and Latent Disentanglement](https://mlsanity.com/paper/2212.05599/) (2022)
- [Improving Covariance Conditioning of the SVD Meta-layer by Orthogonality](https://mlsanity.com/paper/2207.02119/) (2022)
- [Fast and Stable Riemannian Metrics on SPD Manifolds via Cholesky Product Geometry](https://mlsanity.com/paper/2407.02607/) (2024)
- [PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training](https://mlsanity.com/paper/2601.22137/) (2026)
- [Batch-efficient EigenDecomposition for Small and Medium Matrices](https://mlsanity.com/paper/2207.04228/) (2022)
- [Adversarial Defence without Adversarial Defence: Enhancing Language Model Robustness via Instance-level Principal Component Removal](https://mlsanity.com/paper/2507.21750/) (2025)
- [Cross-Level Multi-Instance Distillation for Self-Supervised Fine-Grained Visual Categorization](https://mlsanity.com/paper/2401.08860/) (2024)
- [Learning to Discover Iterative Spectral Algorithms](https://mlsanity.com/paper/2602.09530/) (2026)
- [A Short Note on Batch-efficient Divide-and-Conquer Algorithm for EigenDecomposition](https://mlsanity.com/paper/2604.27325/) (2026)

## Abstract (arXiv)

Computing the matrix square root and 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 the backward pass. In this paper, we propose two more efficient variants to compute the differentiable matrix square root and the inverse square root. For the forward propagation, one method is to use Matrix Taylor Polynomial (MTP), and the other method is to use Matrix Pad'e Approximants (MPA). The backward gradient is computed by iteratively solving the continuous-time Lyapunov equation using the matrix sign function. A series of numerical tests show that both methods yield considerable speed-up compared with the SVD or the NS iteration. Moreover, we validate the effectiveness of our methods in several real-world applications, including de-correlated batch normalization, second-order vision transformer, global covariance pooling for large-scale and fine-grained recognition, attentive covariance pooling for video recognition, and neural style transfer. The experimental results demonstrate that our methods can also achieve competitive and even slightly better performances. The Pytorch implementation is available at https://github.com/KingJamesSong/FastDifferentiableMatSqrt