Tensor Programs (Yang’s Framework)

What a tensor program actually is

A tensor program is a straight-line program over three kinds of object: vectors in R^n (width n, the thing going to infinity), matrices in R^(n×n) sampled iid Gaussian at initialization, and scalars. Only two operations are allowed. MatMul: from a matrix W and a vector x, form Wx or W⊤x. Nonlin: from vectors x¹…xᵏ and a coordinatewise function ψ, form the vector whose i-th entry is ψ(x¹_i, …, xᵏ_i). A third rule extracts scalars as coordinate averages, (1/n)Σ_i ψ(x¹_i,…).

That is deliberately spartan, and the payoff is that training fits inside it too: forward pass, backward pass, and any fixed number of SGD steps are all matmuls and coordinatewise maps. A one-hidden-layer net, both directions:

u   := input vector          # n coords, iid at init
h   := W u                   # MatMul
a   := φ(h)                   # Nonlin
f   := (1/n) Σ_i v_i a_i     # scalar readout

da  := v                     # backward pass starts
dh  := φ’(h) · da            # Nonlin (coordinatewise product)
du  := W⊤ dh                  # MatMul — the SAME W, transposed

Note the last line: W appears twice, once transposed. That reuse is the hinge everything else turns on.

The master theorem, and the term the Gaussian misses

The master theorem says: as n → ∞, for every vector x in the program there is a scalar random variable Z^x such that coordinate averages converge almost surely to expectations — (1/n)Σ_i ψ(x¹_i,…,xᵏ_i) → E[ψ(Z^x¹,…,Z^xᵏ)]. The coordinates are not independent Gaussians; what converges is the joint law across all program vectors, and the theorem hands you rules to compute it.

Nonlin is easy: Z^ψ(x,y) = ψ(Z^x, Z^y). MatMul is where the content lives, and it splits in two — write Zhat for the Gaussian part and Zdot for the correction:

Z^(Wx) = Zhat^(Wx)  +  Zdot^(Wx)

  Zhat^(Wx)   jointly Gaussian across all Wy in the program, with
              Cov(Zhat^(Wx), Zhat^(Wy)) = σ_W² · E[Z^x Z^y]

  Zdot^(Wx)   a linear combination of the earlier vectors y whose W⊤y was
              used, weighted by E[∂Z^x / ∂Zhat^(W⊤y)]

At initialization, x does not depend on W, the correction vanishes, and you recover the classical Gaussian-process picture. After one gradient step it does not vanish, because the backward pass used W⊤ and the weights are now correlated with the activations they multiply. Feature learning lives entirely in that correction term.

A numerical check on the first line

Take the simplest case: h = Wu with W_ij ~ N(0, 1/n) and u normalized so (1/n)Σ_i u_i² = 1. The rules say Z^h ~ N(0,1), so a ReLU second moment should converge to E[relu(Z)²] = 1/2. How close is a real, finite width?

target:  E[relu(Z)²] = ½ · E[Z²] = 0.5
Var:     E[relu(Z)⁴] − (½)² = ½·3 − 0.25 = 1.25
sd of the empirical average at n = 4096:
         √(1.25 / 4096) ≈ 0.0175   →  ~3.5% of the target

simulation, 200 draws at n = 4096:  mean 0.5015, sd 0.0174

Two lessons. The limit is real — the mean is dead on. And the approach is only O(1/√n): at width 4096 a single average still wobbles by several percent, and quadrupling the width merely halves that. Every statement below inherits that rate, which is why the theory guides scaling well and any particular number badly.

Why this subsumes the earlier width limits

Before Tensor Programs, each infinite-width result was an architecture-specific theorem. Neal’s Gaussian-process limit for one-layer nets, its deep extensions (the NNGP correspondence), and the neural tangent kernel each arrived with a proof tailored to a particular family, and each extension — convolutions, recurrence, attention, weight tying — needed the argument redone.

The framework’s claim is architectural universality: those results are corollaries. Write the architecture as a program, apply the two rules, and the limit falls out — and because layer norm, softmax attention and residual connections are all coordinatewise maps and matmuls, transformers are in scope without extra work. The catch is the regularity conditions: the nonlinearities must be well behaved (roughly, pseudo-Lipschitz with polynomially bounded growth), the matrices Gaussian at init, and the step count fixed as width grows. Within those conditions the coverage is genuinely general; outside them the theorem does not speak.

The abc-parametrization: three exponents per layer

Once you can compute limits, you can ask which limit you get — and that depends on how the network is scaled. The abc-parametrization makes the choice explicit: for each layer l, write the effective weight as a multiplier times a trainable parameter with its own initialization scale and learning rate:

W_l  =  n^(−a_l) · w_l          # a_l : forward multiplier
w_l  ~  N(0, n^(−2b_l))         # b_l : init scale
lr(w_l) = η · n^(−c_l)          # c_l : learning rate

These are not independent. Shifting a_l → a_l + θ, b_l → b_l − θ and c_l → c_l − 2θ in every layer leaves the effective weight n^(-a-b) and the effective SGD update n^(-c-2a) unchanged — the trained function is identical at every step. A parametrization is therefore a point in a quotient space, and any table of exponents is one representative. (Below is plain SGD; Adam’s exponents differ, because Adam normalizes the update.)

Three fates: blow-up, frozen features, feature learning

Now sweep the exponents. Every choice lands in one of three buckets, sorted by a single quantity: how much one optimizer step changes a hidden preactivation, as a function of width.

If that change grows with n, the parametrization is unstable — activations or logits diverge and wide models fall over. If it shrinks to zero so that no hidden representation moves in the limit, you are in the kernel regime. The network still learns, because the readout sums n hidden units and n vanishing contributions add up to a finite change in the output — but the features are frozen at initialization, and training is exactly a linear model over a fixed kernel. If the change is Θ(1), representations move by an order-one amount at any width: the feature-learning regime. The classification result is that these are the only options — a stable, non-degenerate parametrization is either kernel-like or feature-learning, with nothing in between.