Source-linked AI summary
Towards Faster Training of Global Covariance Pooling Networks by Iterative Matrix Square Root Normalization
Peihua Li, Jiangtao Xie, Qilong Wang, Zilin Gao
TL;DR
Existing global covariance pooling methods rely on EIG or SVD, which hinder efficient GPU training. This paper introduces iSQRT-COV, a Newton-Schulz-based meta-layer with pre-normalization and post-compensation for end-to-end training. The method runs faster through GPU-suitable matrix multiplications while achieving competitive or superior benchmark performance, including state-of-the-art results on three fine-grained benchmarks.
Problem
Existing matrix square root normalization methods depend on EIG or SVD, whose limited GPU support makes global covariance pooling training inefficient.
Method
iSQRT-COV uses a three-layer meta-layer combining pre-normalization, coupled Newton-Schulz iteration, and post-compensation for end-to-end covariance-pooling training.
Results
iSQRT-COV runs much faster using GPU-suitable matrix multiplications and achieves competitive ImageNet performance, superior results to existing counterparts, and state-of-the-art results on three fine-grained benchmarks.
Takeaways & Limitations
The method supports efficient end-to-end global covariance pooling and extends its reported evaluation across ImageNet and three fine-grained benchmarks.
Takeaways & Limitations
The paper conjectures that approximate matrix square root may be a better robust covariance estimator than the exact square root, leaving this question for future research.
Abstract
from arXiv · showhide
Global covariance pooling in convolutional neural networks has achieved impressive improvement over the classical first-order pooling. Recent works have shown matrix square root normalization plays a central role in achieving state-of-the-art performance. However, existing methods depend heavily on eigendecomposition (EIG) or singular value decomposition (SVD), suffering from inefficient training due to limited support of EIG and SVD on GPU. Towards addressing this problem, we propose an iterative matrix square root normalization method for fast end-to-end training of global covariance pooling networks. At the core of our method is a meta-layer designed with loop-embedded directed graph structure. The meta-layer consists of three consecutive nonlinear structured layers, which perform pre-normalization, coupled matrix iteration and post-compensation, respectively. Our method is much faster than EIG or SVD based ones, since it involves only matrix multiplications, suitable for parallel implementation on GPU. Moreover, the proposed network with ResNet architecture can converge in much less epochs, further accelerating network training. On large-scale ImageNet, we achieve competitive performance superior to existing counterparts. By finetuning our models pre-trained on ImageNet, we establish state-of-the-art results on three challenging fine-grained benchmarks. The source code and network models will be available at http://www.peihuali.org/iSQRT-COV
1. Introduction
Global covariance pooling improves on first-order pooling, but existing matrix square root normalization relies on GPU-unfriendly EIG or SVD. iSQRT-COV addresses this with Newton-Schulz-based end-to-end training and a structured meta-layer designed for convergence and deep-network performance.
- Global covariance pooling has improved recognition over first-order pooling across visual tasks, with matrix square root normalization central to recent state-of-the-art results.
- EIG and SVD limit efficient training because GPU support is weak, forcing CPU computation that restricts multi-GPU concurrency and throughput.
- iSQRT-COV uses Newton-Schulz iteration in both forward and backward propagation, avoiding GPU-unfriendly Schur decomposition or EIG during training.
- Its meta-layer combines pre-normalization, coupled matrix iteration, and post-compensation, with pre-normalization supporting convergence and post-compensation supporting deep ConvNet performance.
2. Related Work
Prior covariance-pooling networks use matrix square root normalization but remain expensive when forward or backward computation depends on SVD, Schur decomposition, EIG, or Lyapunov-equation solutions. iSQRT-COV instead performs both propagation directions with Newton-Schulz iteration and sandwiches it between pre-normalization and post-compensation.
- Improved B-CNN combines matrix square root normalization with element-wise square root and ℓ2-normalization, but its backward propagation can require Lyapunov-equation solutions or decomposition-based gradients.
- Existing improved B-CNN training remains expensive because SVD, Schur decomposition, or EIG are GPU-unfriendly.
- iSQRT-COV uses Newton-Schulz iteration for both forward and backward propagation, so training involves GPU-friendly matrix multiplications.
- The method differs by sandwiching Newton-Schulz iteration between pre-normalization and post-compensation, and by evaluating ImageNet plus three fine-grained benchmarks.
- Global Gaussian-distribution networks and matrix power-normalized covariance pooling also use matrix square root normalization as part of their second-order representations.
3. Proposed iSQRT-COV Network
iSQRT-COV replaces GPU-unfriendly eigendecomposition and SVD in covariance-pooling networks with Newton-Schulz iteration embedded in a three-layer meta-layer. Pre-normalization ensures convergence, post-compensation restores data magnitudes, and matrix backpropagation supplies end-to-end gradients.
- Meta-Layer Design: The meta-layer applies pre-normalization, coupled matrix iteration, and post-compensation as three consecutive nonlinear structured layers.Pre-normalization divides the covariance matrix by its trace or Frobenius norm, while post-compensation counteracts the resulting magnitude change.
- Meta-Layer Design: Pre-normalization makes the locally convergent Newton-Schulz iteration satisfy its convergence condition for covariance matrices.The trace- and Frobenius-norm schemes produce norms below 1, satisfying the stated condition.
- Matrix Square Root Computation: Newton-Schulz iteration computes an approximate matrix square root using only matrix products, enabling parallel GPU implementation.The method uses a fixed iteration count selected by cross-validation; experiments report matching or better performance with no more than 5 iterations.
- Post-Compensation: Post-compensation is necessary for deep-network training: without it, prevalent ResNet fails to converge, while the proposed scheme outperforms Batch Normalization by about 1%.The comparison concerns alternatives for counterbalancing the magnitude changes introduced by pre-normalization.
- Backward Propagation: Gradients for the structured layers are derived with matrix backpropagation methodology, including the pre-normalization, Newton-Schulz iteration, and post-compensation stages.The derivation applies matrix-function chain rules and propagates gradients through the nonlinear layers.
4. Experiments
Experiments evaluate iSQRT-COV on ImageNet and fine-grained benchmarks, examining iteration, efficiency, normalization, convergence, accuracy, and representation compactness. The method delivers faster training, competitive ImageNet results, and strong fine-grained performance.
- Experimental setup: The experiments use ImageNet classification and three fine-grained benchmarks, implemented with CUDA/cuBLAS components callable from Matlab.The fine-grained datasets are Birds, Aircrafts, and Cars.
- Iteration analysis: 1.3% lower error than Plain-COV is achieved with one Newton-Schulz iteration, while performance peaks between 5 and 7 iterations before worsening after N = 7.At three iterations, iSQRT-COV is within 0.3% of MPN-COV; between five and seven iterations it is marginally better.
- Efficiency: iSQRT-COV scales better with batch size and multiple GPUs, whereas MPN-COV tends to saturate above batch size 512.The comparison measures network training speed using forward and backward propagation.
- Design choices: Trace pre-normalization produces 0.3% lower error than Frobenius-norm pre-normalization with similar runtime.Trace-based pre-normalization is used in the remaining experiments.
- Design choices: Without post-compensation, iSQRT-COV fails to converge; Batch Normalization converges but has about 1% higher top-1 error than the proposed scheme.These results establish post-compensation as essential in the reported ResNet-50 experiments.
- ImageNet results: With ResNet-50, iSQRT-COV converges within 60 epochs at 22.14% top-1 error, about 0.6% lower than MPN-COV.Using 90 epochs, it reaches 6.12% top-5 error versus 6.22% with 60 epochs.
- ImageNet results: iSQRT-COV outperforms MPN-COV by 0.6% in top-1 error, and deeper iSQRT-COV models outperform the corresponding original ResNets.The reported improvements are 2.4% over ResNet-101 and 1.8% over ResNet-152 for the 101-layer model.
5. Conclusion
iSQRT-COV enables end-to-end training of covariance-pooling networks using coupled Newton-Schulz iteration instead of GPU-unfriendly EIG or SVD. Its matrix-multiplication-based design runs faster and shows promising performance across ImageNet and fine-grained benchmarks.
- iSQRT-COV is an end-to-end trainable iterative matrix square root normalization network for covariance pooling.
- The method uses coupled Newton-Schulz iteration and only matrix multiplications, making it suitable for parallel GPU implementation.
- Compared with EIG- or SVD-dependent methods, iSQRT-COV runs much faster during training.
- The authors validate iSQRT-COV on ImageNet and challenging fine-grained benchmarks, reporting efficiency and promising performance.