Gradient Flow Through Diagram Expansions: Learning Regimes and Explicit Solutions

TL;DR

This paper introduces a mathematical framework using diagram expansions to analyze gradient flow in large learning problems, revealing different learning regimes and providing explicit solutions, especially for tensor decomposition models.

Contribution

It develops a novel diagram-based power series expansion approach to analyze gradient flow, identifying learning regimes and deriving explicit solutions for complex models.

Findings

Different learning regimes depend on parameter scaling and tensor symmetry.

Explicit solutions of gradient flow are derived in certain regimes.

Theoretical predictions align well with experimental results.

Abstract

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 1st order and solvable by the method of characteristics. We observe a very good agreement of our theoretical predictions with experiment.