Attention as a nonlocal integral, LayerNorm as a projection: Transformers made continuous

A Mathematical Explanation of Transformers

Xue-Cheng Tai, Hao Liu, Lingfeng Li, Raymond H. Chan

cs.LG, cs.AI, math.NA

2025-10-05

The encoder is Lie splitting of an integro-differential equation: attention is a nonlocal integral, LayerNorm and ReLU are projections; Nt=6, J=2 recovers Vaswani 2017.

What problem this solves

The Transformer is the backbone of large language models, yet most accounts of a block still stop at "weighted sum of values." Earlier theory treats the model as an interacting-particle ODE, a high-order spline, or an approximator on a low-dimensional manifold. Those views explain pieces of behavior. They do not put self-attention, LayerNorm, the feedforward block, and residual links on one continuous equation whose time stepping is the encoder itself.

Tai, Liu, Li, and Chan write the encoder as an integro-differential equation on two spatial variables: token index x and feature coordinate y. Time t is depth. After Lie splitting and a uniform grid, the substeps match the modules in Vaswani et al. (2017).

Method

The hidden state is a function u(x, y, t). The right-hand side splits into three operators.

Time stepping uses Lie splitting with step size 1. One step runs attention plus residual, LayerNorm, J linear-plus-ReLU stages, an average of the pre- and post-FFN states (the residual after the FFN), then LayerNorm again. With J=2 this is one encoder block. Nt time steps are Nt blocks; Nt=6 recovers the original 6-layer encoder.

A uniform grid turns integrals into matrix products, so the single-head formula becomes Softmax(QK^T/√Ny)V. Multi-head attention is an extra continuous head coordinate, discretized into Nh heads and summed. ViT is the same evolution with a learned patch embedding in front and a classification head at the end. For images and video, the general QKV integrals become convolutions, which recovers a convolutional Transformer in the style of CvT.

Training is an optimal-control problem: kernels, biases, and (σ1, σ2) are controls, and the loss is on the terminal state.

Results

The paper is a structural identification, not a leaderboard. What can be checked is the exact match after discretization:

Continuous objectDiscrete counterpart
One Lie stepOne encoder block
Nt=6The 6-layer encoder in Vaswani 2017
J=2The standard two-layer FFN
Projection onto S1LayerNorm
Projection onto nonnegative functionsReLU
Integral over a head coordinateMulti-head attention
Convolutional kernels in place of general integralsA one-stage CvT

The authors contrast this with particle-system readings. Those treat tokens as interacting particles and study clustering. Here the whole block is three operators on one integro-differential equation. The same continuous-plus-splitting language was used by this group for CNNs and UNets; Transformer is now on that list.

Why it matters

The design hint is concrete: change the splitting, the time scheme, or the integral kernel, and the discrete net changes with it. PDE tools for stability and conservation become available in principle. The residual average is labeled as numerical relaxation, the same family of tricks used to preserve norms in Runge-Kutta methods.

The work still stops at recovering known stacks. There is no new architecture, no new optimizer, and no well-posedness proof. Training practice gains almost nothing today. The use is theoretical: if you edit attention or normalization, you can name which term in the equation you moved.

Limitations

The authors flag the gaps themselves. The FFN analysis covers ReLU only. The derivation is the encoder block, so causal masks and cross-attention are out of scope. Positional encodings are not in the continuous equation. Well-posedness and regularity of the integro-differential equation are left open.

A few further mismatches. The step size is frozen at 1, so depth and the time grid are the same object; "let numerical analysis pick hyperparameters" has no example. The FFN skip is an average, which differs from u+FFN(u) by a scale; they absorb it by writing the linear map as (I+W). That is retrofitting, not a prediction. There is no experiment showing that Strang splitting or a different kernel trains a better net. Particle-system theory already discusses clustering; this paper barely touches dynamics.

Terms

Source

What people are saying

Related papers

All paper explainers