# Gradient Flow Through Diagram Expansions: Learning Regimes and Explicit Solutions

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 61 references that are in the map.
- First posted on arXiv 4 Feb 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Towards Understanding Grokking: An Effective Theory of Representation Learning](https://mlsanity.com/paper/2205.10343/) (2022)
- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks](https://mlsanity.com/paper/2306.17844/) (2023)
- [Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learning](https://mlsanity.com/paper/2012.09839/) (2020)
- [A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations](https://mlsanity.com/paper/2302.03025/) (2023)
- [Tensor Programs II: Neural Tangent Kernel for Any Architecture](https://mlsanity.com/paper/2006.14548/) (2020)
- [Dynamics of Deep Neural Networks and Neural Tangent Hierarchy](https://mlsanity.com/paper/1909.08156/) (2019)
- [The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon](https://mlsanity.com/paper/2206.04817/) (2022)
- [Saddle-to-Saddle Dynamics in Deep Linear Networks: Small Initialization Training, Symmetry, and Sparsity](https://mlsanity.com/paper/2106.15933/) (2021)
- [Grokking modular arithmetic](https://mlsanity.com/paper/2301.02679/) (2023)
- [Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking](https://mlsanity.com/paper/2311.18817/) (2023)
- [Asymptotics of Wide Networks from Feynman Diagrams](https://mlsanity.com/paper/1909.11304/) (2019)
- [Simplicity Bias in Transformers and their Ability to Learn Sparse Boolean Functions](https://mlsanity.com/paper/2211.12316/) (2022)
- [A Rigorous Framework for the Mean Field Limit of Multilayer Neural Networks](https://mlsanity.com/paper/2001.11443/) (2020)
- [Tensor Programs IIb: Architectural Universality of Neural Tangent Kernel Training Dynamics](https://mlsanity.com/paper/2105.03703/) (2021)
- [From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks](https://mlsanity.com/paper/2409.14623/) (2024)
- [Mean Field Analysis of Deep Neural Networks](https://mlsanity.com/paper/1903.04440/) (2019)
- [Why Do You Grok? A Theoretical Analysis of Grokking Modular Addition](https://mlsanity.com/paper/2407.12332/) (2024)
- [Learning the greatest common divisor: explaining transformer predictions](https://mlsanity.com/paper/2308.15594/) (2023)
- [Feature emergence via margin maximization: case studies in algebraic tasks](https://mlsanity.com/paper/2311.07568/) (2023)
- [Emergence in non-neural models: grokking modular arithmetic via average gradient outer product](https://mlsanity.com/paper/2407.20199/) (2024)
- [Mean Field Limit of the Learning Dynamics of Multilayer Neural Networks](https://mlsanity.com/paper/1902.02880/) (2019)
- [Mechanism of feature learning in deep fully connected networks and kernel machines that recursively learn features](https://mlsanity.com/paper/2212.13881/) (2022)
- [On the asymptotics of wide networks with polynomial activations](https://mlsanity.com/paper/2006.06687/) (2020)

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

- [A prism hierarchy of learning regimes in large linear autoencoders](https://mlsanity.com/paper/2606.05335/) (2026)

## Abstract (arXiv)

We develop a general mathematical framework to analyze scaling regimes and derive explicit analytic solutions for gradient flow (GF) in large learning problems. Our key innovation is a formal power series expansion of the loss evolution, with coefficients encoded by diagrams akin to Feynman diagrams. We show that this expansion has a well-defined large-size limit that can be used to reveal different learning phases and, in some cases, to obtain explicit solutions of the nonlinear GF. We focus on learning Canonical Polyadic (CP) decompositions of high-order tensors, and show that this model has several distinct extreme lazy and rich GF regimes such as free evolution, NTK and under- and over-parameterized mean-field. We show that these regimes depend on the parameter scaling, tensor order, and symmetry of the model in a specific and subtle way. Moreover, we propose a general approach to summing the formal loss expansion by reducing it to a PDE; in a wide range of scenarios, it turns out to be first-order and solvable by the method of characteristics. We observe a very good agreement of our theoretical predictions with experimental results.