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.

Why maximal update is a corner, not a tweak

Which corner do you want? The kernel regime is mathematically beautiful and empirically the wrong answer for pretraining: a model whose representations never move cannot discover features, and that discovery is the whole reason deep learning beats kernel methods. So you want the feature-learning side — where the exponents still have slack, since some layers can learn while others sit frozen.

Maximal update parametrization is the extreme point: the unique choice (up to the abc-symmetry above) that is stable and in which every layer learns features maximally, meaning each one contributes an order-one change to the representation and none can be made to contribute more without breaking stability. Hyperparameter transfer is a consequence rather than the definition. Once every width-dependent quantity is pinned at Θ(1), narrow and wide models run the same dynamics at the same learning rate, so the optimum stops moving. The per-layer rules that implement this — and the coordinate check, which is just the master theorem’s empirical shadow — are in the companion article.

What the framework buys beyond hyperparameter transfer

Four things. First, it changes what “infinite-width network” means. For years the phrase was a synonym for “kernel machine”; the framework shows that is an artifact of a scaling choice, and that the feature-learning limit is itself a well-defined training algorithm you could in principle implement.

Second, it turns kernel derivation for a new architecture into bookkeeping. Third, it produces concrete design corrections rather than vibes — attention logits want 1/d_head scaling once q and k are correlated by training, not the 1/√d_head that is right only at initialization, and that falls out of the same accounting. Fourth, it makes negative predictions, which is how you know it has content: under standard parametrization the optimal learning rate must drift with width. That is not a tuning failure to engineer around; it is what the parametrization implies.

What the asymptotics do not tell you

Depth is not width. The master theorem works because n coordinates give you a law of large numbers to take. Stacking layers gives you no such average — each layer is another dependent operation on the same vectors — so the large-depth limit is a different mathematical object, not a harder version of the same proof. Corrections like 1/√L residual-branch scaling are well-supported practice, not corollaries of this theorem.

The limit is a limit. Results are stated for a fixed number of steps as n → ∞. They say nothing about the O(1/√n) finite-width correction — the numeric check above showed several percent at width 4096 — nothing about long training horizons, nothing about final loss or generalization, and nothing about what does not live on the width axis: dataset size, regularization, batch size. And correct scaling is not a correct value. The theory tells you the learning rate should not move with width; it does not tell you what the learning rate is. You still sweep — you just get to sweep somewhere cheap.