Composing two functions usually means feeding the output of the first into the second. A transformer layer introduces a small twist: a token is transformed using the other tokens as context, and those other tokens are transformed too. The next layer sees a new point and a new context.
In Transformers are Universal In-context Learners, with Takashi Furuya and Maarten de Hoop, we called the resulting operation in-context composition and denoted it by \(\diamond\) [1]. Its definition uses a pushforward of the context measure. There is a simple way to understand the whole construction: enlarge the state to include both the measure and the query. Then this apparently twisted composition becomes ordinary composition. The same viewpoint connects successive layers to particle transport and to an advection equation.
A function of a distribution and a point
A context is a probability measure \(\mu\) on \(\mathbb R^d\). For a finite collection of tokens, its empirical distribution is
\[ \mu=\frac1n\sum_{i=1}^n\delta_{x_i}. \]This representation forgets the ordering of the tokens, which is appropriate for unmasked self-attention without positional information. It also makes sense for a continuous distribution or a weighted point cloud.
An in-context map takes both a context and a query:
\[ f(\mu,x)=f_\mu(x). \]For fixed \(\mu\), \(f_\mu\) is an ordinary spatial map. The query \(x\) can be one of the context tokens, but it can also be a separate point. Merely evaluating a query does not add it to the context.
To keep all formulas well defined, we work with compactly supported probability measures and maps continuous in the spatial variable. Their pushforwards then remain compactly supported. All maps in the main discussion have values in the same \(\mathbb R^d\); different embedding dimensions are also possible when successive domains and ranges match. Regularity with respect to \(\mu\) will enter later.
Pushforward: move the points, carry their mass
For a measurable map \(T:\mathbb R^d\to\mathbb R^d\), the pushforward \(T_\#\mu\) is the distribution after applying \(T\). For a Borel set \(B\),
\[ (T_\#\mu)(B)=\mu(T^{-1}(B)). \]Here \(T^{-1}(B)\) is an inverse image of a set; \(T\) need not be invertible. Equivalently, for every bounded continuous test function \(\varphi\),
\[ \int\varphi(y)\,d(T_\#\mu)(y) =\int\varphi(T(x))\,d\mu(x). \]If \(X\) has law \(\mu\), then \(T(X)\) has law \(T_\#\mu\). In the empirical case,
\[ T_\#\mu=\frac1n\sum_{i=1}^n\delta_{T(x_i)}. \]The weights stay the same; the positions move, and some positions may merge. The identity behind everything that follows is
\[ (S\circ T)_\#\mu=S_\#(T_\#\mu). \]The next layer must see the updated context
Apply \(g\) first. Its spatial map is \(g_\mu\), so the new query and context are \(g_\mu(x)\) and \((g_\mu)_\#\mu\). Applying \(f\) next therefore means
\[ \begin{aligned} \nu&=(g_\mu)_\#\mu,\\ (f\diamond g)(\mu,x)&=f(\nu,g_\mu(x)). \end{aligned} \]Equivalently, as an equality of spatial maps,
\[ (f\diamond g)_\mu=f_\nu\circ g_\mu. \]This is the definition in §2.2, equation (10), of our paper [1]. The order follows ordinary composition: \(g\) acts first. The difference from the frozen-context expression \(f_\mu\circ g_\mu\) is that the second layer uses \(\nu\).
For a finite context, write
\[ y_i=g(\mu,x_i),\qquad \nu=\frac1n\sum_{i=1}^n\delta_{y_i}. \]The second layer outputs \(f(\nu,y_i)\). Thus \(\diamond\) exactly reproduces ordinary composition of the maps acting on the whole token array. It is the pointwise, measure-based description of that composition which needs a new symbol.
Ordinary composition on a lifted space
Let \(\mathcal P_c(\mathbb R^d)\) denote the compactly supported probability measures. Define the measure update \(P_f\) and the lifted map \(\widehat f\) by
\[ \begin{aligned} P_f(\mu)&=(f_\mu)_\#\mu,\\ \widehat f(\mu,x)&=(P_f(\mu),f_\mu(x)). \end{aligned} \]The lifted state lies in \(\mathcal P_c(\mathbb R^d)\times\mathbb R^d\): a population together with a tagged point. Its two components are updated consistently by the same spatial map.
Theorem 1 (the lift removes the twist). For any such in-context maps,
\[ \widehat{f\diamond g}=\widehat f\circ\widehat g, \qquad P_{f\diamond g}=P_f\circ P_g. \]Consequently, \(\diamond\) is associative, and \(e(\mu,x)=x\) is its identity.
Proof: compose the pushforwards
Set \(\nu=P_g(\mu)\). The spatial identity above gives
\[ \begin{aligned} P_{f\diamond g}(\mu) &=(f_\nu\circ g_\mu)_\#\mu\\ &=(f_\nu)_\#\nu=P_f(\nu). \end{aligned} \]The query component is \(f_\nu(g_\mu(x))\), exactly the second component of \(\widehat f(\widehat g(\mu,x))\). This proves the two identities.
The lift is injective: its second component determines \(f(\mu,x)\). Ordinary associativity of lifted maps therefore implies
\[ (f\diamond g)\diamond h=f\diamond(g\diamond h). \]Finally, \(P_e(\mu)=\mu\), so \(e\diamond f=f\diamond e=f\). \(\square\)
Thus in-context maps form a monoid: they can be composed associatively and have an identity. The lift identifies them with a family of ordinary maps closed under composition. If both maps ignore the context, \(f(\mu,x)=F(x)\) and \(g(\mu,x)=G(x)\), then \((f\diamond g)(\mu,x)=F(G(x))\).
A revealing example: center the cloud twice
Let \(m(\mu)=\int y\,d\mu(y)\) be the mean and consider centering:
\[ c(\mu,x)=x-m(\mu). \]The updated context has mean zero. Hence
\[ c\diamond c=c. \]By contrast, freezing the original context would give \(c_\mu(c_\mu(x))=x-2m(\mu)\). The first formula is what one expects from centering a dataset twice.
Order also matters. For a fixed translation \(\tau_a(\mu,x)=x+a\),
\[ c\diamond\tau_a=c,\qquad (\tau_a\diamond c)(\mu,x)=c(\mu,x)+a. \]Translating and then centering removes the translation; centering and then translating retains it.
There is a subtle loss of information here. Each spatial map \(c_\mu\) is a bijective translation. Nevertheless, the lifted centering map is not invertible: simultaneously translating the original context and query gives the same centered output. Pointwise invertibility for each fixed context does not guarantee invertibility of the full in-context transformation.
Self-attention as a transport rule
A single attention head gives a concrete example. Let \(Q,K:\mathbb R^d\to\mathbb R^r\) and \(V:\mathbb R^d\to\mathbb R^d\) be linear maps, and write
\[ \begin{aligned} w(x,y)&=\exp\!\left(\frac{\langle Qx,Ky\rangle}{\sqrt r}\right),\\ A(\mu,x)&=\frac{\int w(x,y)\,Vy\,d\mu(y)} {\int w(x,y)\,d\mu(y)}. \end{aligned} \]Compact support makes these integrals finite and the denominator positive. For an empirical measure, the integrals reduce to the usual softmax-weighted sum. A residual layer can be written as \(g_h(\mu,x)=x+hA(\mu,x)\). Its context update is \((g_{h,\mu})_\#\mu\), so stacking layers uses \(\diamond\).
This measure-theoretic description of attention has earlier roots in the work of Vuckovic, Baratin, and Tachet des Combes [2]. It makes the connection with interacting particles explicit: each token moves according to a rule depending on the distribution of all tokens. The dynamical viewpoint is developed in detail by Geshkovski, Letrouit, Polyanskiy, and Rigollet [3].
Discrete advection and the continuity equation
More generally, let \(v[\mu](x)\) be a velocity depending on the current distribution. A residual layer with step size \(h>0\) is
\[ g_h(\mu,x)=x+h\,v[\mu](x). \]Repeated in-context composition gives the coupled recursion
\[ \begin{aligned} \mu_{k+1}&=(\operatorname{Id}+h\,v[\mu_k])_\#\mu_k,\\ x_{k+1}&=x_k+h\,v[\mu_k](x_k). \end{aligned} \]For an empirical context, every particle follows the same explicit Euler update, with the velocity recomputed from the new cloud at each step. This is a discrete Lagrangian transport scheme: follow the particles and carry their mass. Here \(k\) counts layers; \(t=kh\) will play the role of continuous depth, as in the residual-network and neural-ODE viewpoint [4].
Proposition 1 (consistency with the continuity equation). Let \(\varphi\in C_c^2(\mathbb R^d)\) and assume \(\int|v[\mu_k]|^2\,d\mu_k<\infty\). Then
\[ \begin{aligned} \frac1h\left(\int\varphi\,d\mu_{k+1} -\int\varphi\,d\mu_k\right) &=\int\nabla\varphi\cdot v[\mu_k]\,d\mu_k+R_k,\\ |R_k|&\le\frac h2\|D^2\varphi\|_\infty \int|v[\mu_k]|^2\,d\mu_k. \end{aligned} \]Proof: test the pushforward and use Taylor's formula
The pushforward identity rewrites the difference of integrals as
\[ \int\bigl[\varphi(x+h\,v[\mu_k](x))-\varphi(x)\bigr]\,d\mu_k(x). \]Taylor's formula gives the linear term \(h\nabla\varphi(x)\cdot v[\mu_k](x)\) and a remainder bounded by \(\tfrac12h^2\|D^2\varphi\|_\infty|v[\mu_k](x)|^2\). Integrate and divide by \(h\). \(\square\)
Letting \(h\) tend to zero suggests the weak equation
\[ \frac{d}{dt}\int\varphi\,d\mu_t =\int\nabla\varphi\cdot v[\mu_t]\,d\mu_t, \]or, in distributional notation,
\[ \partial_t\mu_t+\nabla\cdot(\mu_t\,v[\mu_t])=0. \]The Taylor estimate is a consistency statement. Convergence additionally needs control of the velocity and its dependence on the measure; Lipschitz and growth assumptions provide a standard framework for well-posedness and convergent particle schemes [5].
If \(\mu_t\) has density \(\rho_t\), the equation reads \(\partial_t\rho+\nabla\cdot(\rho v)=0\). This is the mass-conserving continuity equation, also called a transport equation. Expanding it gives \(\partial_t\rho+v\cdot\nabla\rho+\rho\,\nabla\cdot v=0\). The last term accounts for compression or expansion; the passive-scalar advection equation omits it and coincides with this equation only when \(v\) is divergence free.
The characteristic point satisfies \(\dot X_t(x)=v[\mu_t](X_t(x))\), with \(\mu_t=(X_t)_\#\mu_0\). For an autonomous velocity rule with globally defined, unique characteristic solutions, let \(F_t(\mu,x)\) denote this characteristic starting from context \(\mu\). Restarting at time \(t\) yields
\[ F_{s+t}=F_s\diamond F_t. \]The nonlinear evolution is a semigroup under in-context composition. Its lift is an ordinary semigroup on measure–point pairs. If the velocity also depends explicitly on depth, one instead uses a two-time evolution family.
An exactly solvable flow: attraction to the mean
Take \(v[\mu](x)=m(\mu)-x\). A discrete step gives
\[ g_h(\mu,x)=(1-h)x+h\,m(\mu). \]For \(0<h<1\), this pulls every point towards the mean. The mean stays fixed, and the covariance is multiplied by \((1-h)^2\). After \(k\) steps,
\[ g_h^{\diamond k}(\mu,x) =m(\mu)+(1-h)^k(x-m(\mu)). \]Passing to \(t=kh\), the continuous formula is
\[ F_t(\mu,x)=m(\mu)+e^{-t}(x-m(\mu)). \]The distribution contracts around its conserved mean, and its covariance decays as \(e^{-2t}\). The identity \(F_{s+t}=F_s\diamond F_t\) can be checked directly because every intermediate distribution has the same mean. This is a simple model of consensus; general attention dynamics need not preserve the mean or follow this particular contraction.
Why the Wasserstein viewpoint is natural
A measure-based description should be stable when the context points move. The \(1\)-Wasserstein distance compares distributions by coupling their points:
\[ W_1(\mu,\nu)=\inf_{\pi\in\Pi(\mu,\nu)} \int|x-y|\,d\pi(x,y), \]where \(\Pi(\mu,\nu)\) is the set of joint laws with marginals \(\mu\) and \(\nu\).
Proposition 2 (stability of the context update). Suppose
\[ |f_\mu(x)-f_\nu(y)| \le L_x|x-y|+L_\mu W_1(\mu,\nu). \]Then
\[ W_1(P_f(\mu),P_f(\nu)) \le(L_x+L_\mu)W_1(\mu,\nu). \]Proof: transport a coupling through the two spatial maps
For any \(\pi\in\Pi(\mu,\nu)\), the law \((f_\mu,f_\nu)_\#\pi\) couples the two updated measures. Consequently,
\[ \begin{aligned} W_1(P_f(\mu),P_f(\nu)) &\le\int|f_\mu(x)-f_\nu(y)|\,d\pi(x,y)\\ &\le L_x\int|x-y|\,d\pi+L_\mu W_1(\mu,\nu). \end{aligned} \]Take the infimum over \(\pi\). \(\square\)
Such estimates connect finite clouds to continuous contexts and allow errors to be propagated across layers. Attention regularity is studied in [6]; our universality paper uses continuity in the context measure to formulate approximation uniformly over its number of tokens [1].
Related viewpoints and a few boundaries
The pushed measure \(P_f(\mu)\) records where the population ends up, but it does not retain every pointwise correspondence. For a distribution symmetric about the origin, the identity and reflection \(x\mapsto-x\) give the same updated measure while moving a tagged query differently. This is why the measure–point lift contains more information than the measure update alone.
Architectures acting compositionally on distributions also appear in our earlier Stochastic Deep Networks with Gwendoline de Bie and Marco Cuturi [7]. The distinctive feature emphasized here is to keep the tagged-point map visible together with its induced distribution update.
Transport by pushforward does not by itself make a layer an optimal transport map, or make its continuous-depth limit a Wasserstein gradient flow. Those require extra structure. Sinkformers, with Michael Sander, Pierre Ablin, and Mathieu Blondel, studies a particular attention normalization and its links to transport and gradient-flow interpretations [8].
Finally, the choice of state depends on the architecture. In causal self-attention, retain each token's position \(s\) and push the joint context by \((y,s)\mapsto(g(\mu,y,s),s)\): the embedding moves, while its position label stays fixed, as in equation (13) of [1]. In cross-attention to a fixed external memory, that memory is unchanged, so ordinary composition with the memory held fixed is appropriate. The diamond rule describes a context which evolves together with its queries.
References
- Takashi Furuya, Maarten V. de Hoop, and Gabriel Peyré. Transformers are Universal In-context Learners. ICLR, 2025.
- James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. A Mathematical Theory of Attention. arXiv:2007.02876, 2020.
- Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on Transformers. Bulletin of the American Mathematical Society 62(3), 427–479, 2025.
- Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural Ordinary Differential Equations. NeurIPS, 2018.
- Benedetto Piccoli and Francesco Rossi. Transport equation with nonlocal velocity in Wasserstein spaces: convergence of numerical schemes. Acta Applicandae Mathematicae 124, 73–105, 2013.
- James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. On the Regularity of Attention. arXiv:2102.05628, 2021.
- Gwendoline de Bie, Gabriel Peyré, and Marco Cuturi. Stochastic Deep Networks. ICML, PMLR 97, 1556–1565, 2019.
- Michael E. Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transformers with Doubly Stochastic Attention. AISTATS, PMLR 151, 3515–3530, 2022.