TL;DR. Tensor programs are a mathematical language for describing
computations in wide neural networks. Their master theorem says that, as the
width grows, averages over the neurons become predictable: they converge to
expectations in a much simpler scalar random process. This turns the analysis
of a high-dimensional random network into a tractable probability calculation
and provides a rigorous foundation for studying its infinite-width behavior,
including Gaussian-process limits
Lee et al. (2018) and neural tangent
kernels Jacot et al. (2018).
This note explains the proof for computations that may reuse a random weight
matrix and its transpose, as happens with weight sharing and backpropagation.
Such reuse creates correlations, so each multiplication cannot simply be
replaced by independent Gaussian noise. The proof conditions on all earlier
uses of the matrix and separates the next output into a correction determined
by those uses and a genuinely fresh Gaussian part. An induction then shows
that the scalar process tracks the full computation, including cases in which
some intermediate vectors are linearly dependent.
Why matrix reuse matters
Let and let
,
independently of . Conditioned on , each coordinate of is Gaussian
with variance . This may suggest replacing every
matrix product by fresh Gaussian noise, but the next use of the same matrix
disproves that idea:
The first term converges to . If , the second term,
conditioned on , is centered Gaussian with variance
and is independent of . Hence
Thus in fixed-coordinate distribution, not in vector
norm. The -term is precisely the dependence that a fresh-noise
approximation misses.
The master theorem and the -rules
The program language, scalar -calculus, and master theorem in this section
are from Section 2 of Yang's Tensor Programs III
[Yang 2020]. A
program starts
from random vectors and random matrices
, and repeatedly applies
Here and every nonlinearity are fixed as . The matrices
are independent, with
and the coordinate slices are iid copies of a
fixed Gaussian vector .
The limiting scalar attached to each program vector is defined
recursively.
Initial vectors and nonlinearities
For , use the corresponding component of
, and set
For a coordinatewise operation,
Matrix products
Write
For a fixed oriented matrix , the variables are
jointly centered Gaussian, with
Fresh Gaussian families belonging to different matrices or orientations
(in particular, and ) are independent. Reuse through the opposite
orientation is restored by
Here the sum is over earlier products on whose fresh Gaussian
variables depends. The derivative is symbolic: first express as a
deterministic function of all preceding 's, then differentiate
that expression.
For nonsmooth functions, the expected derivatives can equivalently be
defined by Gaussian regression. If are the
relevant earlier products, set
Then define the vector of
to be
.
For differentiable functions this agrees with Stein's identity
and its multivariate version.
Master theorem
A function is polynomially bounded if
for some .
If all program nonlinearities are polynomially bounded, then for any program
vectors and any polynomially bounded test function ,
The theorem concerns empirical coordinate laws; it does not assert that the
finite-width coordinates are independent.
The central conditional calculation
This is the conditioning step used in the proof of the master theorem
(Appendices K and L of
Yang 2020).
Order the initial and matrix-generated vectors (the G-vars) as
. Suppose the next one is
where .
Collect all earlier uses of in both orientations:
and form
Empty matrices are allowed when or .
Let . Conditioned on ,
all the displayed vectors are fixed and is subject to
Gaussian conditioning under these linear constraints gives the exact law
where is an independent copy of ,
, , and
Multiplying by yields
with independent of and
This separates the part forced by previous uses of from the remaining
Gaussian randomness.
Limits of the conditional parameters
Introduce the finite Gram data
Moment convergence for the preceding G-vars gives deterministic limits,
denoted with bars; for example,
Rank stability, proved in the next section, implies
and
almost surely. Therefore
Indeed, the expression in parentheses is the squared -distance from
to .
Using Equation (5), define
Then
Consequently, with
we have
Identification with the -calculus
Set .
By Equation (1),
Gaussian regression therefore gives, conditionally on all earlier
-variables,
The mean and variance here are exactly
and Equation (8).
Applying the ZDot rule, Equation (2), with Stein's
identity or its pseudoinverse form, gives the matching drift identity
Adding Equation (11) and Equation (12) proves the key scalar conditional law
conditionally on the earlier -variables. For , this reduces
to , recovering the opening example.
Why the Gram-matrix limits are stable
The core-set and rank-stability mechanism below follows Appendix L of
Yang (2020). The proof simultaneously
maintains two induction statements.
Moment invariant
For every polynomially bounded ,
Core-set invariant
There is such that:
the vectors form a basis for all preceding G-vars, with
unique coefficients independent of ;
the law of and Lebesgue measure have the same null
sets;
for every prescribed triangular array of Lebesgue-null sets ,
almost surely all sufficiently large coordinate slices
avoid .
The last property upgrades almost-everywhere limiting identities to exact
finite-width identities at every coordinate.
To see the consequence, let be polynomially bounded
functions of the preceding G-vars and let
If , then
Thus the linear combination vanishes almost everywhere under the core law.
The density and null-avoidance properties make it vanish at every finite-width
coordinate for all large , so . Ordinary lower
semicontinuity of rank gives the reverse inclusion. Therefore, eventually,
The pseudoinverse is continuous on a fixed-rank stratum, proving the
pseudoinverse convergence used above.
Updating the core set
Return to . By Equation (8),
There are two cases.
Zero residual variance
If , fixed coefficients satisfy
Expressing and each as coordinatewise functions of the old core
turns this into an almost-everywhere functional identity. Core density and
null avoidance make it exact at finite width:
almost surely for all sufficiently large . Hence is already in the old
span and the core need not change.
Positive residual variance
If , Equation (13) has the form
Conditioned on the old core, has an everywhere-positive Gaussian
density. Fubini's theorem then shows that adjoining preserves the two-way
null-set property, so the new core is .
At finite width, . Moreover, rank stability and
old-core null avoidance imply
Thus every , conditionally on the past, retains positive Gaussian
variance. Applying Fubini to each prescribed null set and then old-core null
avoidance proves the new null-avoidance property. The same fresh Gaussian
component prevents from lying in the old span. This completes the
core-set update in both cases.
Moment convergence
The final concentration argument combines the projection estimates and
weakly correlated Gaussian strong law from Appendices K and L of
Yang (2020). It remains to prove
. When , the exact relation above
makes a fixed polynomially bounded function of earlier G-vars, so
applies directly. Assume henceforth that
.
Conditioned on , the second term inside braces is the expectation
of the first. Although the coordinates of are correlated, the
dependence has fixed rank. Indeed, if , then
the normalized off-diagonal correlations of satisfy
A strong law for such weakly correlated Gaussian arrays applies. Its required
high moments follow from polynomial boundedness,
, the convergence of , and
. Therefore
Term : replace finite parameters by their limits
The projection diagonal satisfies
Hence at most coordinates have ; their contribution
vanishes by Hölder's inequality and the same polynomial moment bounds.
On the remaining coordinates, both variances stay bounded away from zero
because . Gaussian convolution smooths even a
merely measurable polynomially bounded , with
These derivatives are bounded by polynomial weights in the preceding
coordinates and . Meanwhile, Equation (10), Equation (8), and
Equation (17) imply
The mean-value bound followed by Cauchy--Schwarz now gives
Term : use the induction hypothesis
Every and is a polynomially bounded coordinatewise function of
earlier G-vars. Consequently
is polynomially bounded. Applying , then using
the scalar identification in Equation (13), gives
Thus , proving .
Closing the induction
For the initial Gaussian vectors, the ordinary strong law proves the moment
invariant. Choose a maximal linearly independent subfamily of their limiting
Gaussian variables as the first core. Its covariance is positive definite, so
its law has an everywhere-positive density. Coordinatewise linear relations
and countable null-set avoidance prove the other core properties.
The preceding sections establish the simultaneous step
Induction therefore proves the moment statement for all G-vars.
Finally, every program vector is a fixed polynomially bounded
coordinatewise function of the complete list of G-vars:
The composition of polynomially bounded functions is polynomially bounded.
Applying the G-var moment result to gives
Equation (3), completing the proof.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S. Schoenholz, Jeffrey
Pennington, and Jascha Sohl-Dickstein. "Deep neural networks as Gaussian
processes." International Conference on Learning Representations, 2018.
https://arxiv.org/abs/1711.00165.
Arthur Jacot, Franck Gabriel, and Clément Hongler. "Neural tangent
kernel: Convergence and generalization in neural networks." Advances in
Neural Information Processing Systems 31, 2018.
https://arxiv.org/abs/1806.07572.
TL;DR. Tensor programs are a mathematical language for describing computations in wide neural networks. Their master theorem says that, as the width grows, averages over the neurons become predictable: they converge to expectations in a much simpler scalar random process. This turns the analysis of a high-dimensional random network into a tractable probability calculation and provides a rigorous foundation for studying its infinite-width behavior, including Gaussian-process limits Lee et al. (2018) and neural tangent kernels Jacot et al. (2018).
This note explains the proof for computations that may reuse a random weight matrix and its transpose, as happens with weight sharing and backpropagation. Such reuse creates correlations, so each multiplication cannot simply be replaced by independent Gaussian noise. The proof conditions on all earlier uses of the matrix and separates the next output into a correction determined by those uses and a genuinely fresh Gaussian part. An induction then shows that the scalar process tracks the full computation, including cases in which some intermediate vectors are linearly dependent.
Why matrix reuse matters
Let
The first term converges to
and is independent of
Thus
The master theorem and the
The program language, scalar
Here
and the coordinate slices
The limiting scalar
Initial vectors and nonlinearities
For
For a coordinatewise operation,
Matrix products
Write
For a fixed oriented matrix
Fresh Gaussian families belonging to different matrices or orientations (in particular,
Here the sum is over earlier products
For nonsmooth functions, the expected derivatives can equivalently be defined by Gaussian regression. If
Then define the vector of
and its multivariate version.
Master theorem
A function is polynomially bounded if
The theorem concerns empirical coordinate laws; it does not assert that the finite-width coordinates are independent.
The central conditional calculation
This is the conditioning step used in the proof of the master theorem (Appendices K and L of Yang 2020). Order the initial and matrix-generated vectors (the G-vars) as
where
and form
Empty matrices are allowed when
Let
Gaussian conditioning under these linear constraints gives the exact law
where
Multiplying by
with
This separates the part forced by previous uses of
Limits of the conditional parameters
Introduce the finite Gram data
Moment convergence for the preceding G-vars gives deterministic limits, denoted with bars; for example,
Rank stability, proved in the next section, implies
Indeed, the expression in parentheses is the squared
Using Equation (5), define
Then
Consequently, with
we have
Identification with the
Set
Gaussian regression therefore gives, conditionally on all earlier
The mean and variance here are exactly
Applying the ZDot rule, Equation (2), with Stein's identity or its pseudoinverse form, gives the matching drift identity
Adding Equation (11) and Equation (12) proves the key scalar conditional law
conditionally on the earlier
Why the Gram-matrix limits are stable
The core-set and rank-stability mechanism below follows Appendix L of Yang (2020). The proof simultaneously maintains two induction statements.
Moment invariant
For every polynomially bounded
Core-set invariant
There is
The last property upgrades almost-everywhere limiting identities to exact finite-width identities at every coordinate.
To see the consequence, let
If
Thus the linear combination vanishes almost everywhere under the core law. The density and null-avoidance properties make it vanish at every finite-width coordinate for all large
The pseudoinverse is continuous on a fixed-rank stratum, proving the pseudoinverse convergence used above.
Updating the core set
Return to
There are two cases.
Zero residual variance
If
Expressing
almost surely for all sufficiently large
Positive residual variance
If
Conditioned on the old core,
At finite width,
Thus every
Moment convergence
The final concentration argument combines the projection estimates and weakly correlated Gaussian strong law from Appendices K and L of Yang (2020). It remains to prove
From Equation (6),
Let
and define the one-dimensional Gaussian average
The difference in the moment invariant is bounded by
Term
Conditioned on
A strong law for such weakly correlated Gaussian arrays applies. Its required high moments follow from polynomial boundedness,
Term
The projection diagonal satisfies
Hence at most
On the remaining coordinates, both variances stay bounded away from zero because
These derivatives are bounded by polynomial weights in the preceding coordinates and
The mean-value bound followed by Cauchy--Schwarz now gives
Term
Every
is polynomially bounded. Applying
Thus
Closing the induction
For the initial Gaussian vectors, the ordinary strong law proves the moment invariant. Choose a maximal linearly independent subfamily of their limiting Gaussian variables as the first core. Its covariance is positive definite, so its law has an everywhere-positive density. Coordinatewise linear relations and countable null-set avoidance prove the other core properties.
The preceding sections establish the simultaneous step
Induction therefore proves the moment statement for all G-vars.
Finally, every program vector
The composition of polynomially bounded functions is polynomially bounded. Applying the G-var moment result to
References