Exact Learning Dynamics of a Linear Autoencoder Through Diagram Expansions

Eugene Golikov, Yaroslav Gusev, Dmitry Yarotsky

Keywords: gradient flow, analytic solutions, diagram expansion, wide networks, linear networks, autoencoders, matrix factorization

TL;DR: We apply a novel method akin to Feynman diagrams to obtain exact solutions for gradient flow dynamics of a linear autoencoder in the limit of large input and hidden sizes

Abstract:

We consider the problem of factorizing the identity matrix as a product of two matrices - that is, learning a shallow linear autoencoder - with gradient flow.
We formally expand the loss as a function of time, and show how the expansion terms could be expressed via graphs akin to Feynman diagrams.
This turns the problem of computing the loss as a function of training time into a purely combinatorial problem.
By analyzing this problem, we provide a complete classification of learning regimes in the limit of large input and hidden size.
These limit regimes can be identified with various faces of a *Pareto polygon* describing the problem.
We compute loss expansion terms exactly in the limit regimes corresponding to most of these faces.
These formal expansions turn out to be summable and lead to analytic formulas for the average loss evolution which agree very well with the experiment.
We emphasize that some of our solutions correspond to intrinsically nonlinear training dynamics; such solutions are scarce in the literature.

Submission Number: 31