Tensor Programs: The Feynman
Diagrams of Deep Learning
Greg YangMSR AIOutlineCorrelation Functions in Physics& MLCorrelation Functions in Physics
Source: https://en.wikipedia.org/wiki/Correlation_function_(statistical_mechanics)• Given a distribution on functions
𝜙: ℝ𝑘 → ℝ, the 2-point correlation
function is defined to be𝐶2 𝑥1 , 𝑥2 = 𝔼𝜙 𝜙 𝑥1 𝜙 𝑥2
• Measure of order in a physicalsystem (statmech)• Measures “amplitude” of a particle
going from x to y (QFT)• Often calculated with Feynman
Diagrams

Correlation Functions in ML• Given a neural network, the last layer
before the output is called the“embedding”• For an input
𝑥, write its embedding as𝑚
Φ 𝑥 ∈ℝ• Consider the “2-pointfunction”
𝑚𝐶2 𝑥1 , 𝑥2

1= ෍ Φ 𝑥1 𝛼 Φ 𝑥2 𝛼
𝑚𝛼=1

• The 2-point function holds semanticalinformation• The more semantically similar two inputs
are to each other, the larger the correlationof the embeddings

Credit:https://towardsdatascience.com/understanding-feature-engineering-part-4-deep-learning-methods-for-text-data-96c44370bbfa

Word2vec example• Woman – queen = man – king
• Cats – cat = dogs – dog• France – Paris = England - London

More frequent words also
usually have larger norms

Credit: https://samyzaf.com/ML/nlp/word2vec2.png

Thus correlation functions are quite
important in ML as well (though they
don’t go by that name).

Some Physics of Neural Networks
via Correlation Functions

Neural Network Terminologies
DepthWidth• Width = #neurons in ahidden layer• Depth = #layers

https://medium.com/datadriveninvestor/when-not-to-use-neural-networks-89fb50622429

Wide Neural Networks• Why do we care?
• The wider the better!

???

• In fact this contradicts
classical learning theoryorthodoxy• So what happens at theinfinite width limit?

𝑤 = width multiplier

Tan & Le 2019

Neural Networks at the InfiniteWidth LimitGaussian Process Behavior

Architecture

Slide borrowed from Jascha Sohl-Dickstein

Gaussian Process Behavior

Slide borrowed from Jascha Sohl-Dickstein

Gaussian Process Behavior

Slide borrowed from Jascha Sohl-Dickstein

Gaussian Process Behavior

Slide borrowed from Jascha Sohl-Dickstein

Gaussian Process Behavior

Slide borrowed from Jascha Sohl-DicksteinRecap: Gaussian Process• Given 𝜇: 𝑋 → ℝ and 𝐾: 𝑋 2 → ℝ, a Gaussian Process 𝒢𝒫(𝜇, 𝐾) is a
“Gaussian distribution” over the function space over 𝑋• 𝑓 ∼ 𝒢𝒫(𝜇, 𝐾) if for any finite set of positions 𝑥1 , … , 𝑥𝑘 , the tuple(𝑓 𝑥1 , … , 𝑓 𝑥𝑘 ) is distributed like 𝒩 𝝁, 𝑲 where
𝜇 𝑥1
𝐾 𝑥1 , 𝑥1 ⋯ 𝐾 𝑥1 , 𝑥𝑘⋮⋮⋱⋮
𝝁=,𝑲=𝜇 𝑥𝑘
𝐾 𝑥𝑘 , 𝑥1 ⋯ 𝐾 𝑥𝑘 , 𝑥𝑘

Slide borrowed from Jascha Sohl-DicksteinRecap: Gaussian Process• Given 𝜇: 𝑋 → ℝ and 𝐾: 𝑋 2 → ℝ, a Gaussian Process 𝒢𝒫(𝜇, 𝐾) is a
“Gaussian distribution” over the function space over 𝑋• 𝑓 ∼ 𝒢𝒫(𝜇, 𝐾) if for any finite set of positions 𝑥1 , … , 𝑥𝑘 , the tuple(𝑓 𝑥1 , … , 𝑓 𝑥𝑘 ) is distributed like 𝒩 𝝁, 𝑲 where
𝜇 𝑥1
𝐾 𝑥1 , 𝑥1 ⋯ 𝐾 𝑥1 , 𝑥𝑘
⋮
⋮
⋱
⋮
𝝁=,𝑲=
𝜇 𝑥𝑘
𝐾 𝑥𝑘 , 𝑥1 ⋯ 𝐾 𝑥𝑘 , 𝑥𝑘
• Of course, 𝐾 just specifies the 2-point correlations of “fields” drawn
from 𝒢𝒫(𝜇, 𝐾)

Slide borrowed from Jascha Sohl-DicksteinUniversality of Gaussian Process Behavior• For any neural network architecture, feedforward or recurrent, thenetwork, when randomly initialized, becomes a Gaussian Process as
its widths tend to ∞.
• Intuition: GP kernel = 2-point function of the last layer embedding
Simple RNN

GRUTransformer

Batchnorm

Yang 2019

What Happens when Depth →∞?

Architecture

Slide borrowed from Jascha Sohl-Dickstein

Infinite-width GP kernelof an 𝐿-layer NN

Expand layer 𝐿 kernel around fixed opint𝐿𝐿∗
𝐾 = 𝐾 + 𝑂 exp −𝜉𝑐• Edge of Chaos• 𝜉𝑐 → ∞• Training is stable• Dissimilar inputsconverge• 𝐾 ∗ = 𝛽𝟏𝟏⊤
• Typical network is close
to a constant function• Similar inputs diverge• 𝐾 ∗ = 𝛼𝐼 + 𝛽𝟏𝟏⊤
• Typical network is close
to white noise

Schoenholz et al. 2017
Slide borrowed from Jascha Sohl-Dickstein

Using Edge of Chaos for Profit!!!

!!!

Slide borrowed from Jascha Sohl-DicksteinSmall Learning Rate TrainingNeural Tangent Kernel• Another kernel (i.e. 2-point correlation function), the Neural TangentKernel is associated to the small learning rate training of neuralnetworks• Naïve Tayler expansion of neural network around parameters𝑓 𝑥; 𝜃 − 𝑓 𝑥; 𝜃0 ≈ ∇𝜃 𝑓 𝑥; 𝜃0 , 𝜃 − 𝜃0
• If 𝜃 doesn’t move much (e.g. if we use small learning rate), then NN =a linear model with a kernel called the NTK:𝑁𝑇𝐾 𝑥, 𝑥 ′ = ∇𝜃 𝑓 𝑥; 𝜃0 ⊤ ∇𝜃 𝑓(𝑥; 𝜃0 )• The NTK converges to a deterministic kernel as width → ∞ whenparameters are randomly initializedUniversality of Neural Tangent Behavior• For any neural network architecture, feedforward or recurrent, theNTK of the network converges deterministically, when the network israndomly initialized, as the network widths tend to ∞.
• If trained under small learning rate, then the network evolves like alinear model with the NTK kernel.
• These kernels have now achieved state-of-the-art in classification
tasks with little data

How to Compute the Correlation
Functions in a Wide Neural Net?

• The kernel of the limiting Gaussian process of an NN is essentially the2-point function of its last layer embeddings Φ 𝑥 when the width 𝑛is large:𝑛1′
𝐾 𝑥, 𝑥 ≈ ෍ Φ 𝑥 Φ 𝑥 ′ ≡ ⟨Φ 𝑥 Φ 𝑥 ′ ⟩𝑛𝑖=1• In our example NN with 𝐿 layers, Φ 𝑥 = 𝑦 𝐿 (𝑥)• So suffices to compute the 𝑛 → ∞ limit of ⟨𝑦𝑖𝐿 𝑥 𝑦𝑖𝐿 𝑥 ′ ⟩Feynman Diagrams?
• In the spirit of Feynman diagrams, we shall• Taylor expand all nonlinearities• Say 𝜙 𝑥 = 𝑎0 + 𝑎1 𝑥 + 𝑎2 𝑥 2 + ⋯• Algebra• Apply Wick’s theorem• Let 𝑥 and 𝑥′ be two inputs

𝑦𝑖𝐿 𝑥 𝑦𝑖𝐿 (𝑥 ′ )= ⟨𝑎02 + 𝑎0 𝑎1 𝑧𝑖𝐿−1 𝑥 + 𝑎0 𝑎1 𝑧𝑖𝐿−1 (𝑥 ′ ) + 𝑎12 𝑧𝑖𝐿−1 (𝑥)𝑧𝑖𝐿−1 (𝑥 ′ ) + ⋯ ⟩= ⟨𝑎02 +𝑎0 𝑎1 𝑊𝑖𝑗𝐿−1 𝑦𝑗𝐿−1 (𝑥) + 𝑎0 𝑎1 𝑊𝑖𝑗𝐿−1 𝑦𝑗𝐿−1 (𝑥 ′ ) + 𝑎12 𝑊𝑖𝑗𝐿−1 𝑦𝑗𝐿−1 (𝑥)𝐿−1 ′𝑊𝑖𝐿−1(𝑥 ) + ⋯ ⟩′ 𝑗 ′ 𝑦𝑗 ′= ⋯ 𝑔𝑒𝑡𝑠 𝑐𝑜𝑚𝑝𝑙𝑖𝑐𝑎𝑡𝑒𝑑 ⋯Tensor Program: An Intuitive IntroIntuition
If 𝑊 is Gaussian matrix, 𝑊𝑥 shouldhave approx. iid Gaussian entries
𝑊𝑥 𝑖 ≈ 𝑁 0, ⟨𝑥𝑗 𝑥𝑗 ⟩ .
If 𝑊 is Gaussian matrix, (𝑊𝑥, 𝑊𝑥 ′ )should have approx. iid jointly
Gaussian entries( 𝑊𝑥 𝑖 , 𝑊𝑥 ′ 𝑖 )⟨𝑥𝑗 𝑥𝑗 ⟩ ⟨𝑥𝑗 𝑥𝑗′ ⟩≈ 𝑁 0,⟨𝑥𝑗′ 𝑥𝑗 ⟩ ⟨𝑥𝑗′ 𝑥𝑗′ ⟩Then for any 𝑙,𝑦𝑖𝑙+1 𝑥 𝑦𝑖𝑙+1 (𝑥 ′ ) = 𝔼𝜙 𝜁 𝜙 𝜁 ′where𝜁, 𝜁 ′ ∼ 𝒩 0,𝑧𝑗𝑙 𝑥 𝑧𝑗𝑙 𝑥 ′𝑧𝑗𝑙 𝑥 ′ 𝑧𝑗𝑙 𝑥

𝑧𝑗𝑙 𝑥 ′ 𝑧𝑗𝑙 𝑥 ′and𝑧𝑖𝑙 𝑥 𝑧𝑖𝑙 𝑥 ′ = ⟨𝑦𝑖𝑙 𝑥 𝑦𝑖𝑙 𝑥 ′ ⟩with boundary condition
𝑦𝑖0 𝑥 𝑦𝑖0 𝑥 ′ = 𝑥𝑗 𝑥𝑗′ .

Let 𝐾 𝑙 =Let’s try to do some computations!

𝑧𝑗𝑙 𝑥 𝑧𝑗𝑙 𝑥

𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥

𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥 ′𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥

𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥 ′𝐾 𝑙+1 = 𝑉𝜙 𝐾 𝑙, then≔ 𝔼𝜁∼𝒩 0,𝐾𝑙 𝜙 𝜁Ԧ 𝜙 𝜁Ԧ⊤For simplicity, assume𝜎𝑤 = 1, 𝜎𝑏 = 0

This motivates us to introduce a simple language for expressing computation
for which the correlation function is easy to keep track ofGiven some initial set of matrices and vectors, we can Think of theseas the neuralform new ones by:
networkparameters• (MatMul) Given a matrix 𝑊, and vector 𝑥, we canRoughly Gaussian
generate
vectors, by our intuition:⊤𝑊𝑥 or 𝑊 𝑥.
Think: preactivations• (Nonlin) Given a bunch of vectors 𝑥 1 , … , 𝑥 𝑘 of thesame size, and a function 𝜙: ℝ𝑘 → ℝ, we cangenerate a new vector
Roughly, image of1
𝑘𝜙(𝑥 , … , 𝑥 )Gaussian vectors.
Think: activations(with 𝜙 applied coordinatewise)
A sequence of vectors generated this way, along with theinitial vectors and matrices, is known as a tensor programExample Tensor Program: forward andbackpropagation
1

𝐼𝑛𝑝𝑢𝑡: 𝑊𝑖𝑗𝑙 ∼ 𝑁 0, 𝑛 ,𝑤𝑖 , 𝑏𝑖𝑙 ∼ 𝑁(0, 1).

Think of 𝑤 as last layer
gradient𝑥෪1 = 𝑊 1 𝑥 + 𝑏1𝑥 1 = 𝜙 𝑥෪1𝑥෪2 = 𝑊 2 𝑥 1 + 𝑏2𝑥 2 = 𝜙 𝑥෪2⋮𝑥 𝐿 = 𝜙(𝑥෪𝐿 )𝑑𝑥෪𝐿 = 𝑤 ⊙ 𝜙 ′ 𝑥෪𝐿𝑑𝑥 𝐿−1 = 𝑊 𝐿⊤ 𝑑𝑥෪𝐿⋮𝑑 𝑥෪1 = 𝑑𝑥 2 ⊙ 𝜙 ′ 𝑥෪1𝑑𝑥 = 𝑊 1⊤ 𝑑 𝑥෪1𝑥෪1 ′ = 𝑊 1 𝑥 + 𝑏1𝑥 1 ′ = 𝜙 𝑥෪1 ′𝑥෪2 ′ = 𝑊 2 𝑥 1 ′ + 𝑏2𝑥 2 ′ = 𝜙 𝑥෪2 ′⋮𝑥 𝐿 ′ = 𝜙(𝑥෪𝐿 ′)𝐿 ′ = 𝑤 ⊙ 𝜙 ′ 𝑥෪𝐿′𝑑 𝑥෪
𝐿′
𝑑𝑥 𝐿−1 ′ = 𝑊 𝐿⊤ 𝑑𝑥෪
⋮𝑑 𝑥෪1 ′ = 𝑑𝑥 2 ′ ⊙ 𝜙 ′ 𝑥෪1 ′𝑑𝑥′ = 𝑊 1⊤ 𝑑 𝑥෪1 ′Two Surprises of Tensor Programs

Surprise I:
Universality of Tensor Program Expressivity
Almost all modern architectures:
Any composition of• Multilayer perceptron• Recurrent neural networks• LSTM• GRUs• Skip connection• Convolution• Pooling• Batch normalization• Layer normalization• Attention• etc
Both forward and backward
propagation
Even SGD training!

Surprise II:
Any tensor program has a “mean field theory”
that allows us to calculate its correlation
functionsIf- the initial set of vectors and matrices of a tensor program are all randomized as Gaussians, and- all the vectors in a tensor program are 𝑔1 , … , 𝑔𝑀 ,then for any 𝜓: ℝ𝑀 → ℝ
e.g. 𝑔2 = 𝑊 𝑔1
𝑔1

Average
1
෍𝑛

𝑔3

𝑔𝑀

𝜓()

𝜓()

𝜓()

1

𝜓()

𝜓()

𝜓()

For correlation, take 𝜓
to be the productfunction.

e.g. 𝑔3 = 𝜙 𝑔1 , 𝑔2𝑛→∞𝑛𝑔2𝑎. 𝑠.

𝔼𝑍 ~ 𝒩 𝜇,Σ𝜓(???

𝜇 and Σ are some mean andcovariance matrices that arecomputed recursively like inthe example we did.

Roughly speaking, each vector 𝑔𝑖 will be coordinatewise iid, but for each 𝛼, thecoordinate-slice (𝑔𝛼1 , … , 𝑔𝛼𝑀 ) has a nonseparable joint distribution which is some image of𝒩 𝜇, Σ .

)All of the correlation
functions wementioned can becomputed by applying
this theorem to an
appropriate tensor
programExample

Let 𝐾 𝑙 =𝐾𝑙+1

𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥

𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥 ′𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥

𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥 ′= 𝑉𝜙 𝐾, then≔ 𝔼𝜁∼𝒩 0,𝐾𝑙 𝜙 𝜁Ԧ 𝜙 𝜁Ԧ

𝑙0

𝐾 =

𝑥𝑖 𝑥𝑖

𝑥𝑖 𝑥𝑖′

𝑥𝑖′ 𝑥𝑖𝑥𝑖′ 𝑥𝑖′⊤and ∀𝑖, (𝑧𝑖𝑙 𝑥 , 𝑧𝑖𝑙 (𝑥 ′ )) ≈ 𝒩(0, 𝐾 𝑙 )𝑧 0 𝑥 𝑧 0 𝑥′)

𝜓()

𝜓()

1

𝜓()

𝜓()

𝜓()

𝑛→∞

Average𝑛1
෍𝑛𝜓(For simplicity, assume
𝜎𝑤 = 1, 𝜎𝑏 = 0𝑎. 𝑠.

𝔼

𝑍 0 ~ 𝒩 0,𝐾𝑍10
𝜓
(0𝑍20)Example

Let 𝐾 𝑙 =𝐾𝑙+1

𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥 ′𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥 ′= 𝑉𝜙 𝐾𝑙, then≔ 𝔼𝜁∼𝒩 0,𝐾𝑙 𝜙 𝜁Ԧ 𝜙 𝜁Ԧ

0

𝐾 =

𝑥𝑖 𝑥𝑖

𝑥𝑖 𝑥𝑖′

𝑥𝑖′ 𝑥𝑖𝑥𝑖′ 𝑥𝑖′⊤

and ∀𝑖, (𝑧𝑖𝑙 𝑥 , 𝑧𝑖𝑙 (𝑥 ′ )) ≈ 𝒩(0, 𝐾 𝑙 )𝑧 0 𝑥 𝑧 0 𝑥′ 𝑧1 𝑥 𝑧1 𝑥′)

𝜓()

𝜓()

1

𝜓()

𝜓()

𝜓()

𝑛→∞

Average𝑛1
෍𝑛𝜓(For simplicity, assume
𝜎𝑤 = 1, 𝜎𝑏 = 0𝑎. 𝑠.

𝔼

𝑍 0 ~ 𝒩 0,𝐾 0 ,1𝑍 ∼𝒩(0,𝑉𝜙 𝐾 0 )𝜓 ( 𝑍10𝑍20𝑍11𝑍21)

Example

Let 𝐾 𝑙 =𝐾𝑙+1

𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥

𝑦𝑗𝑙 𝑥 𝑦𝑗𝑙 𝑥 ′𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥

𝑦𝑗𝑙 𝑥 ′ 𝑦𝑗𝑙 𝑥 ′= 𝑉𝜙 𝐾, then≔ 𝔼𝜁∼𝒩 0,𝐾𝑙 𝜙 𝜁Ԧ 𝜙 𝜁Ԧ𝑙0

𝐾 =𝑥𝑖 𝑥𝑖

𝑥𝑖 𝑥𝑖′

𝑥𝑖′ 𝑥𝑖𝑥𝑖′ 𝑥𝑖′⊤and ∀𝑖, (𝑧𝑖𝑙 𝑥 , 𝑧𝑖𝑙 (𝑥 ′ )) ≈ 𝒩(0, 𝐾 𝑙 )𝑧 0 𝑥 𝑧 0 𝑥′ 𝑧1 𝑥 𝑧1 𝑥′ 𝑧 2 𝑥 𝑧 2 𝑥′)

𝜓()

𝜓()

1

𝜓()𝜓()

𝜓()𝑛→∞

Average𝑛1
෍𝑛𝜓(For simplicity, assume𝜎𝑤 = 1, 𝜎𝑏 = 0𝑎. 𝑠.

𝔼𝑍 0 ~ 𝒩 0,𝐾0 ,1𝑍 ∼𝒩(0,𝑉𝜙 𝐾 0 ),𝑍 2 ∼𝒩(0,𝑉𝜙 ∘𝑉𝜙 𝐾 0 )𝜓 ( 𝑍10𝑍20𝑍11𝑍21𝑍12𝑍22)

Tensor Programs Summary

languagetechniqueframework

framework
language

A language for specifying
computation graph involving
matrices and vectors

e.g.

Express deep neural network
forward computation in atensor program

techniqueBehavioral insight of thecomputation when thedimensions → ∞Gaussian process behavior
when the widths → ∞Consequences• Universality of Gaussian Process behavior• Universality of Neural Tangent Kernel behavior• Phase transitions for different neural architectures• New proofs of semicircle and Marchenko-Pastur laws
• Jacobian singular value of wide neural networks• State Evolution for Approximate Message Passing algorithm inCompressed Sensing• Gradient Independence phenomenon• More in the paper!

Tensor Programs and Feynman
Diagrams

Differences and Similarities
Feynman Diagrams

Tensor Programs• Perturbative (small 𝑔)• Series Expansion• Diagrams reflect particleinteraction• Handles interaction from higher
order terms in the action• Organizing principle for particle
physics• Infinite-Width Limit (large 𝑛)• Recursive Computation• Programs reflect neural network
computation• Handles interactions fromcoordinatewise nonlinearities• I hope would become theorganizing principle for thetheory of deep learning

Concrete Example in Random Matrix Theory
Feynman Diagrams

Tensor Programs• Symmetric matrices 𝐴 with
distribution∝ exp −𝑡𝑟 𝐴2 − 𝑡𝑟 𝐴4• Symmetric matrices 𝐴 with
𝐴𝑖𝑗 = 𝑊𝑖𝑗 ෍ 𝑊𝑖 ′𝑗𝑖′

෍ 𝑊𝑖𝑗 ′𝑗where 𝑊𝑖𝑗 is a symmetric
Gaussian matrix (i.e. GOE)
• In general, can study jacobian
singular values of random deepneural networks

My Question to Physicists• “Tensor Programs” seems like the “right” language for the theory ofdeep learning• Some very shiny, attractive mathematics is now possible as a result• Does it help perform calculations in physics?

Papers
Scaling Limits of Wide Neural Networks
https://arxiv.org/abs/1902.04760

Wide NNs are GPs
https://arxiv.org/abs/1910.12478