# Tensor Normal Training for Deep Learning Models

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

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 58 references that are in the map.
- First posted on arXiv 5 Jun 2021.

## Builds on (most cited first, 19 of 19)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Practical Gauss-Newton Optimisation for Deep Learning](https://mlsanity.com/paper/1706.03662/) (2017)
- [A Kronecker-factored approximate Fisher matrix for convolution layers](https://mlsanity.com/paper/1602.01407/) (2016)
- [Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers](https://mlsanity.com/paper/2007.01547/) (2020)
- [Three Mechanisms of Weight Decay Regularization](https://mlsanity.com/paper/1810.12281/) (2018)
- [On Empirical Comparisons of Optimizers for Deep Learning](https://mlsanity.com/paper/1910.05446/) (2019)
- [A Stochastic Quasi-Newton Method for Large-Scale Optimization](https://mlsanity.com/paper/1401.7020/) (2014)
- [Practical Quasi-Newton Methods for Training Deep Neural Networks](https://mlsanity.com/paper/2006.08877/) (2020)
- [Stochastic Quasi-Newton Methods for Nonconvex Stochastic Optimization](https://mlsanity.com/paper/1607.01231/) (2016)
- [Stochastic Block BFGS: Squeezing More Curvature out of Data](https://mlsanity.com/paper/1603.09649/) (2016)
- [RES: Regularized Stochastic BFGS Algorithm](https://mlsanity.com/paper/1401.7625/) (2014)
- [A Statistically Identifiable Model for Tensor-Valued Gaussian Random Variables](https://mlsanity.com/paper/1911.02915/) (2019)

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

- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [A New Perspective on Shampoo's Preconditioner](https://mlsanity.com/paper/2406.17748/) (2024)
- [ASGO: Adaptive Structured Gradient Optimization](https://mlsanity.com/paper/2503.20762/) (2025)
- [Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures](https://mlsanity.com/paper/2311.00636/) (2023)
- [Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective](https://mlsanity.com/paper/2402.03496/) (2024)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [ASDL: A Unified Interface for Gradient Preconditioning in PyTorch](https://mlsanity.com/paper/2305.04684/) (2023)
- [An Improved Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/2406.06420/) (2024)
- [On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer](https://mlsanity.com/paper/2603.09952/) (2026)
- [Stochastic Approximation Beyond Gradient for Signal Processing and Machine Learning](https://mlsanity.com/paper/2302.11147/) (2023)
- [Trade-Offs of Diagonal Fisher Information Matrix Estimators](https://mlsanity.com/paper/2402.05379/) (2024)
- [A Mini-Block Fisher Method for Deep Neural Networks](https://mlsanity.com/paper/2202.04124/) (2022)
- [Correlations Are Ruining Your Gradient Descent](https://mlsanity.com/paper/2407.10780/) (2024)
- [On the Convergence Analysis of Muon](https://mlsanity.com/paper/2505.23737/) (2025)
- [Optimizer-Model Consistency: Full Finetuning with the Same Optimizer as Pretraining Forgets Less](https://mlsanity.com/paper/2605.06654/) (2026)
- [KrADagrad: Kronecker Approximation-Domination Gradient Preconditioned Stochastic Optimization](https://mlsanity.com/paper/2305.19416/) (2023)
- [Studying K-FAC Heuristics by Viewing Adam through a Second-Order Lens](https://mlsanity.com/paper/2310.14963/) (2023)
- [Rank-1 Approximation of Inverse Fisher for Natural Policy Gradients in Deep Reinforcement Learning](https://mlsanity.com/paper/2601.18626/) (2026)
- [Tight Complexity of Muon under Heavy-Tailed Noise](https://mlsanity.com/paper/2605.18528/) (2026)
- [Eva: A General Vectorized Approximation Framework for Second-order Optimization](https://mlsanity.com/paper/2308.02123/) (2023)
- [Gathering and Exploiting Higher-Order Information when Training Large Structured Models](https://mlsanity.com/paper/2312.03885/) (2023)
- [Layer-wise Adaptive Step-Sizes for Stochastic First-Order Methods for Deep Learning](https://mlsanity.com/paper/2305.13664/) (2023)
- [Spectral-factorized Positive-definite Curvature Learning for NN Training](https://mlsanity.com/paper/2502.06268/) (2025)

## Abstract (arXiv)

Despite the predominant use of first-order methods for training deep learning models, second-order methods, and in particular, natural gradient methods, remain of interest because of their potential for accelerating training through the use of curvature information. Several methods with non-diagonal preconditioning matrices, including KFAC, Shampoo, and K-BFGS, have been proposed and shown to be effective. Based on the so-called tensor normal (TN) distribution, we propose and analyze a brand new approximate natural gradient method, Tensor Normal Training (TNT), which like Shampoo, only requires knowledge of the shape of the training parameters. By approximating the probabilistically based Fisher matrix, as opposed to the empirical Fisher matrix, our method uses the block-wise covariance of the sampling based gradient as the pre-conditioning matrix. Moreover, the assumption that the sampling-based (tensor) gradient follows a TN distribution, ensures that its covariance has a Kronecker separable structure, which leads to a tractable approximation to the Fisher matrix. Consequently, TNT's memory requirements and per-iteration computational costs are only slightly higher than those for first-order methods. In our experiments, TNT exhibited superior optimization performance to state-of-the-art first-order methods, and comparable optimization performance to the state-of-the-art second-order methods KFAC and Shampoo. Moreover, TNT demonstrated its ability to generalize as well as first-order methods, while using fewer epochs.