Papers
Topics
Authors
Recent
Search
2000 character limit reached

Tensor Programs Technique Overview

Updated 2 June 2026
  • Tensor Programs technique is a formal framework that encodes neural computations as straight-line programs over Gaussian matrices, enabling rigorous infinite-width analysis.
  • It systematically derives closed-form expressions for Neural Tangent Kernels, Gaussian Process kernels, and network training dynamics across diverse architectures.
  • The method extends to complex models such as RNNs, Transformers, and adaptive optimizers, offering a universal, mechanizable approach to deep network analysis.

The Tensor Programs technique is a formal framework for the analysis and understanding of wide neural networks, capable of rigorously tracking the statistics of high-dimensional random variables that arise in both the forward and backward passes of arbitrarily structured neural architectures. By encoding neural computations as straight-line programs over random Gaussian matrices and vectors, Tensor Programs enable a systematic derivation of infinite-width limits, yielding closed-form expressions for quantities such as the Neural Tangent Kernel (NTK), Gaussian Process kernels of neural networks, and dynamics under various optimizers including SGD and Adam. This methodology extends to feedforward, recurrent, convolutional, attention-based, and batch-normalized architectures, furnishing the first general, mechanically implementable, and architecture-agnostic approach for computing kernel and dynamical behavior in modern deep networks (Yang, 2019, Yang, 2020, Yang et al., 2021, Yang et al., 2023, Sakai et al., 1 Jun 2025).

1. Formalism of Tensor Programs

A Tensor Program is defined as a straight-line sequence of variable assignments of the following primary types:

  • G-vars: Vectors in Rn\mathbb{R}^n, representing pre-activations or gradients, initialized or generated via random matrix operations.
  • H-vars: Vectors in Rn\mathbb{R}^n resulting from coordinatewise nonlinearities applied to G-vars.
  • A-vars: Random Gaussian matrices, typically n×nn \times n, representing weight matrices.

Statement primitives include:

  • MatMul: For A∈Rn×nA\in\mathbb{R}^{n\times n} and h∈Rnh\in\mathbb{R}^n, g:=Ahg := A h, yielding g∼N(0,σA2E[hh⊤])g \sim \mathcal{N}(0, \sigma_A^2 \mathbb{E}[hh^\top]).
  • Nonlin/Nonlin+^+: For g1,...,gk∈Rng_1,...,g_k \in \mathbb{R}^n, h:=Ï•(g1,...,gk)h := \phi(g_1,...,g_k), with Rn\mathbb{R}^n0 acting coordinatewise.
  • Moment/Avg: For Rn\mathbb{R}^n1, compute Rn\mathbb{R}^n2.
  • OuterNonlin (NEXORT language): Permits nonlinear functions across multiple coordinates and samples, required to handle adaptive optimizer updates (Yang et al., 2023).

This programmatic representation is universal for neural network computations expressible as directed acyclic graphs of linear and nonlinear operations, encompassing forward, backward, and parameter update computations (Yang et al., 2021, Yang, 2019).

2. The Master Theorem and Infinite-Width Behavior

The core analytical result is the Master Theorem, which asserts that as Rn\mathbb{R}^n3, the empirical distribution of any finite collection of program variables converges to a deterministic law:

Rn\mathbb{R}^n4

where Rn\mathbb{R}^n5 are jointly Gaussian random variables determined inductively by the recursive structure of the program—for MatMul lines, by the covariance of the input vectors and the scaling of the random matrix; for Nonlin operations, by applying the nonlinearity to the relevant Gaussians (Yang, 2019, Yang, 2020).

Crucially, Tensor Programs generalize beyond simple feedforward architectures: dependencies between forward and backward passes, such as the weight-transpose coupling intrinsic to backpropagation, are accounted for via explicit Gaussian conditioning and rank-stability arguments. The Master Theorem's machinery enables derivation of exact infinite-width limits for quantities such as the NTK (via sum-of-gradient outer products), as well as for output covariances relevant to Gaussian Processes (Yang, 2020, Yang et al., 2021).

3. Treatment of Complex Architectures and Operations

Tensor Programs extend to all standard components in modern neural architectures:

  • Recurrent Networks: By encoding time-unrolled computations, recursive kernel recurrences for RNN hidden states and outputs are mechanically derivable (Yang, 2020).
  • Attention and Transformers: Initial formulations provided Gaussian process limits for attention with infinite head counts or specific scalings. Recent advances extend the formalism to attention layers at finite head count and realistic scaling, uncovering that the infinite-width limit is not Gaussian but maintains a hierarchical conditional-Gaussian law over dot-product similarity scores (Sakai et al., 1 Jun 2025). In this formulation, output coordinates of an attention layer are conditionally Gaussian given the limiting random similarity scores, but exhibit non-Gaussianity marginally.
  • Normalization Layers: Operations such as BatchNorm are handled by introducing auxiliary variables for means and variances, which become deterministic projections and shifts in the infinite-width limit (Yang, 2020).
  • Adaptive Optimization (e.g., Adam): The NEXORT language extends the Tensor Program model, using "OuterNonlin" to encode per-coordinate, possibly nonlinear, update schemes (such as Adam's moving averages and normalizations). This extension preserves the dichotomy between "kernel regime" and "feature-learning regime," allowing closed-form kernel and maximal update dynamical limits for a broad class of optimizers (Yang et al., 2023).

4. Mechanization and Graphical Notation

Tensor Programs admit a graphical notation where variables correspond to nodes, MatMul operations to directed edges, Nonlin gates to boxes, and Moment/Avg to special aggregating nodes. This notation elucidates the directed computational structure and highlights reused parameters (e.g., weight-sharing), path equivalences, and the propagation of randomness.

This mechanized, graph-compatible formalism enables the systematic construction and verification of the infinite-width limits for arbitrarily complex architectures, supporting full embedding of optimization dynamics—including the entire sequence of SGD or adaptive steps—within a finite Tensor Program whose analysis yields dynamical laws for outputs, hidden states, and parameter updates (Yang et al., 2021).

5. Algorithms for Infinite-Width Limit Computation

The practical application of the Tensor Programs technique follows an explicit pipeline:

  1. Program construction: Encode the neural network's computation—forward, backward, and update steps—as a Tensor Program using the prescribed primitives (Yang, 2019).
  2. Recursion derivation: Translate each operation into a mapping for means and covariances (matrices and vectors), respecting the dependencies introduced by parameter sharing and nonlinearities.
  3. Master Theorem application: Compute the limiting joint law of all variables by inductively applying the Master Theorem's recursion. For architectures with attention, special treatment of dot-product lines must be included, resulting in conditional Gaussian laws indexed by the dot-product statistics (Sakai et al., 1 Jun 2025).
  4. Kernel extraction: For networks initialized at infinite width, derive the NTK or GP kernel by recursive evaluation of the kernel update equations, optionally deriving explicit training dynamics (e.g., kernel gradient descent ODEs for outputs) (Yang, 2020, Yang et al., 2021).
  5. Implementation: Open-source reference implementations exist for both Gaussian Process and NTK kernels across MLPs, RNNs, batchnorm, and transformers (Yang, 2019, Yang, 2020).

6. Adaptive Optimization and the NEXORT Extension

With the introduction of the NEXORT language [Editor’s term: NEXORT], Tensor Programs encompass entrywise adaptive optimizers. NEXORT includes "OuterNonlin" operations, which permit updates depending not just on current gradients but on the history of per-parameter gradients and state variables, required for optimizers like Adam. The infinite-width limit then tracks the distribution of all program variables under the optimizer's nonlinear update rules.

The Master Theorem and its generalizations yield:

  • Neural Tangent limit: Output trajectories follow kernel dynamics under a (possibly nonlinear) operator defined via the entrywise optimizer.
  • Maximal Update limit: Feature-learning dynamics under correct scaling, supporting μP ("maximal update parametrization") as well as standard mean-field and kernel regimes.

These results hold broadly: the infinite-width stationary dynamics dichotomizes into operator (kernel) or feature-learning regimes, with no intermediate possibilities, even for sophisticated adaptive optimizers (Yang et al., 2023).

7. Extensions, Impact, and Future Directions

Tensor Programs now provide a universal, scalable toolkit for infinite-width analysis in deep learning. Notable advancements include:

  • Exact kernels for arbitrary architectures: Any architecture codable in the Tensor Program language admits an exact, closed-form infinite-width kernel (Yang, 2020).
  • Rigorous dynamical laws under training: Embedding SGD or momentum-based optimization trajectories as Tensor Programs yields provably correct kernel gradient descent dynamics ("architectural universality of NTK") (Yang et al., 2021).
  • Non-Gaussian infinite-width limits in attention: For the first time, provably non-Gaussian behavior in transformers' attention layers is characterized, opening the path to a comprehensive theory of transformers at and beyond initialization (Sakai et al., 1 Jun 2025).
  • Adaptive optimization: Universal treatment of optimizers such as Adam recovers new nonlinear effective kernels and demonstrates the absence of new dynamical regimes (Yang et al., 2023).

Potential future directions include the propagation of non-Gaussian statistics across multiple layers in deeply stacked architectures, dynamic analyses over entire training trajectories with attention mechanisms, and the development of further automated tools for kernel extraction and large-n limit analysis.


References:

(Yang, 2019, Yang, 2020, Yang et al., 2021, Yang et al., 2023, Sakai et al., 1 Jun 2025)

Topic to Video (Beta)

No one has generated a video about this topic yet.

Whiteboard

No one has generated a whiteboard explanation for this topic yet.

Follow Topic

Get notified by email when new papers are published related to Tensor Programs Technique.