auto-fpt: Automating Free Probability Theory Calculations for Machine Learning Theory
Arjun Subramonian, Elvis Dohmatob
TL;DR
The paper tackles the challenge of computing the high-dimensional limiting value $r$ of the expected trace of a rational expression in large random matrices, a central task in ML theory. It introduces auto-fpt, a lightweight Python/SymPy tool that symbolically reduces the problem to a reduced system of fixed-point equations using operator-valued free probability. The authors demonstrate auto-fpt on several problems—the MP law, classical ridge regression generalization, matrix subordination, and high-dimensional random features—showing how a pencil-based input yields tractable FP systems and known theoretical results. By providing an open-source, reproducible workflow, auto-fpt aims to streamline high-dimensional ML analyses and facilitate discovery of new phenomena in scalable learning settings.
Abstract
A large part of modern machine learning theory often involves computing the high-dimensional expected trace of a rational expression of large rectangular random matrices. To symbolically compute such quantities using free probability theory, we introduce auto-fpt, a lightweight Python and SymPy-based tool that can automatically produce a reduced system of fixed-point equations which can be solved for the quantities of interest, and effectively constitutes a theory. We overview the algorithmic ideas underlying auto-fpt and its applications to various interesting problems, such as the high-dimensional error of linearized feed-forward neural networks, recovering well-known results. We hope that auto-fpt streamlines the majority of calculations involved in high-dimensional analysis, while helping the machine learning community reproduce known and uncover new phenomena.