This repository constains the official Pytorch implementation of ICLR 22 paper "Fast Differentiable Matrix Square Root" and the expanded T-PAMI journal "Fast Differentiable Matrix Square Root and Inverse Square Root".

You can find the presentation of our work by the slides and poster.

Check torch_utils.py for the implementation. Minimal exemplery usage is given as follows:

# Import and define function
from torch_utils import *
FastMatSqrt = MPA_Lya.apply
FastInvSqrt = MPA_Lya_Inv.apply

# For any batched matrices, compute their square root or inverse square root:
rand_matrix = torch.randn(5,32,32)
rand_cov = rand_matrix.bmm(rand_matrix.transpose(1,2))
rand_cov_sqrt = FastMatSqrt(rand_cov)
rand_inv_sqrt = FastInvSqrt(rand_cov)

All the codes for the following experiments are available:

• Decorrelated Batch Normalization (BN)
• Second-order Vision Transformer (So-ViT)
• Neural Style Transfer by Whitening and Coloring Transform (WCT)
• Temporal-Attentive Covariance Pooling (TACP) for Video Action Recognition
• Global Covariance Pooling for Large-Scale and Fine-grained Visual Recognition

Please consider citing our paper if you think the code is helpful to your research.

@inproceedings{song2022fast,
  title={Fast Differentiable Matrix Square Root},
  author={Song, Yue and Sebe, Nicu and Wang, Wei},
  booktitle={ICLR},
  year={2022}
}

@article{song2022fast2,
  title={Fast Differentiable Matrix Square Root and Inverse Square Root},
  author={Song, Yue and Sebe, Nicu and Wang, Wei},
  journal={IEEE TPAMI},
  year={2022},
  publisher={IEEE}
}

If you have any questions or suggestions, please feel free to contact me

yue.song@unitn.it