REVIEW 3 major objections 4 minor 24 references
Dynamical Behaviors of the Gradient Flows for In-Context Learning
T0 review · 3 major / 4 minor · reviewed 2026-08-11 · deepseek-v4-flash
Pith's one-line read Training a linear attention model by gradient flow is shown to be a completely described dynamical system in two important cases: the full gradient-flow equations are derived in closed form, and for the simplified and one-dimensional…
desk verdict A careful, mostly sound derivation of the full gradient flow for linear ICL, with solid local analysis but global basin claims that rest on an unstated stable-manifold/Lojasiewicz argument. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central machinery is a change of bookkeeping: the loss is written in summation form rather than matrix form, so each partial derivative reduces to Gaussian moment computations carried out with Isserlis' theorem, and the resulting explicit ODE system is Theorem 3.1. The qualitative analysis then rides on two exact invariants. In the simplified system ($z=Z=0$) the quantity $v^2(t)-\operatorname{Tr}[U U^\top](t)$ is constant, so the flow is laminated by rotational hyperboloids $v^2=\operatorname{Tr}[UU^\top]+\kappa$; after diagonalizing $\Lambda$ and $\Gamma=(1+1/N)\Lambda+(1/N)\operatorname{Tr}[\Lambda]I$, the critical-point equations collapse to $v\Gamma U=I$ for the two symmetric attractors $A$ and to $v=0$ with $\sum_i \lambda_i^2 B_{ii}=0$ for the degenerate saddles $B$. In the $d=1$ full system the invariant is $v^2-U^2-Z^2+z^2=\kappa$, and the critical points split into type A with two zero variables satisfying $\alpha vU=1$ or $\alpha zZ=1$, and type B with $vU=zZ=\rho=1/(\alpha+\sqrt{\beta\delta}+\gamma)$, where $\alpha=(N+2)/N$, $\beta=(3N+6)/N$, $\gamma=(2N+7)/N$, and $\delta=(N+8)/N+15/N^2$. The Hessian at each critical point is symmetric, which makes the eigenvalue count straightforward.
What would settle it
A concrete check: take the simplified system with $d=2$, $\Lambda=\operatorname{diag}(1,2)$, $N=1$, and solve the linearized equations at a type-B point with $B_{11}=-4B_{22}$. The paper's formula predicts two nonzero eigenvalues of opposite sign and one zero eigenvalue from the $d^2-d$ off-diagonal block; if the computed product of the two nonzero eigenvalues is not negative, Theorem 4.1 fails. More globally, integrate (10)-(11) from an initial condition with $v>0$ and $U=0$; the paper predicts convergence to the attractor $A$ on that leaf, so a trajectory that instead enters a periodic orbit or converges to $B$ would refute the exhaustive description.
Extended reading notes
Core claim
The paper claims that the gradient flow for the linear in-context-learning loss (1) is exactly $(\dot U, \dot v, \dot Z, \dot z) = -(\partial L/\partial U, \partial L/\partial v, \partial L/\partial Z, \partial L/\partial z)$, with the four derivatives given in closed form in Lemmas 3.2 through 3.5. For the simplified system with $z=Z=0$, each trajectory lies on a leaf $v^2 = \operatorname{Tr}[UU^\top] + \kappa$, and on each leaf the only attracting critical points are the two points $A$ with $v\Gamma U = I$, while the points $B$ with $v=0$ and $\sum_i \lambda_i^2 B_{ii}=0$ are degenerate saddles whose stable manifolds separate the basins of attraction. For the full $d=1$ system, each trajectory lies on $K_\kappa = \{v^2 - U^2 - Z^2 + z^2 = \kappa\}$, the type-A points with only two nonzero variables satisfying $\alpha vU=1$ or $\alpha zZ=1$ are attractors, and the type-B points with all variables nonzero and $vU=zZ=\rho=(\alpha+\sqrt{\beta\delta}+\gamma)^{-1}$ are hyperbolic saddles; for $\kappa=0$ the origin is an additional saddle. The paper concludes that this gives an exhaustive qualitative description of these two gradient flows for the full range of parameters and initial conditions.
Load-bearing premise
The argument's load-bearing premise is that every bounded trajectory of these gradient flows eventually settles at a critical point and that the stable manifolds of the saddle points really separate the state space into the claimed basins of attraction; the paper relies on this standard fact without stating or proving it.
Editorial extensions
If this is right
- For the simplified system, the invariant $v^2=\operatorname{Tr}[UU^\top]+\kappa$ shows that the state space is foliated by one leaf per initial condition; when $\kappa>0$ the saddle points $B$ do not lie on the leaf, so every initial condition in that regime flows to one of the two attractors $A$.
- Because the stable manifolds of the saddles have positive codimension, almost all Lebesgue-almost-every initial conditions converge to an attractor of type $A$, so the training dynamics is globally convergent up to a null set even though the loss is nonconvex.
- In the $d=1$ system the same qualitative pattern holds on each leaf $K_\kappa$: two attractors and four saddles, with the two-dimensional stable manifolds of the saddles separating the basins; at $\kappa=0$ the origin is an additional saddle.
- The loss values at the critical points are ordered $L(A)=-1/(2\alpha)<L(B)=-\rho<L(O)=0$, so the flow can only converge to a restricted global minimum, never to a suboptimal critical point.
- The closed-form gradient-flow system of Theorem 3.1 is the missing ingredient for extending the analysis to all four weight matrices of the linear attention model.
Reading between the lines
- If the structural pattern seen in the two analyzed systems persists in the full model, the training landscape of linear ICL has no unique minimum; instead, the initialization fixes an invariant that selects among a discrete family of attractors, which would explain why different runs of the same architecture can converge to different in-context predictors.
- The paper's saddle-avoidance conclusion suggests a quantitative test of practical training: for linear attention trained by stochastic gradient descent, the probability of landing near a type-B saddle should be essentially zero, and the observed convergence rate should be governed by the negative eigenvalues of the linearization at type-A points rather than by the saddle's unstable direction.
- The $d=1$ invariant tori $T_{\rho,\mu}\subset K_\kappa$ suggest reducing the four-dimensional flow to a two-dimensional system on each torus; a normal-form calculation on this reduction would give explicit rates for the spiraling approach to $A$ and would predict the transient slowdown near $B$ that the present paper describes only qualitatively.
- One testable extension is to perturb the simplified system by adding small $z,Z$ terms: since the type-A points are hyperbolic attractors within their leaf, the full system should still have nearby hyperbolic attractors for sufficiently small coupling, so the basin structure is structurally stable.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper derives the gradient-flow ODE system for a linear in-context-learning loss with four parameter blocks (U, z, Z, v), computes the Gaussian moments explicitly (Lemmas 3.2–3.5, Theorem 3.1), and then analyzes two specializations. For the simplified system with z = Z = 0, it identifies the invariant leaf v² = Tr(UUᵀ) + κ, classifies the critical points into attracting A-points and degenerate saddle B-points, and claims global almost-everywhere convergence to A with B-separatrices as the exceptional set (Section 4). For d = 1 with all weights, it derives invariant manifolds Kκ, classifies critical points of type A, B, and O, and claims an exhaustive global qualitative description (Section 5).
Significance. If the global claims go through, the paper would be the first to write the full linear ICL gradient flow in closed form and to describe its invariant foliation and saddle structure beyond the restricted settings of [24]. The moment calculations are self-contained and parameter-free, and the invariant v² = Tr(UUᵀ) + κ is an elegant structural observation. The local stability analyses at A and B are largely explicit and checkable, and the paper does not rely on fitted constants or ad hoc assumptions.
major comments (3)
- [Section 4, Remark 4.3] The global statement in Remark 4.3 — that almost all trajectories converge to A-type attractors except for separatrices of B-type saddles forming a finite union of smooth manifolds of positive codimension — is not established by the preceding analysis. The proof supplies local stability (Theorem 4.1) and an invariant foliation, but passing to the global claim requires (i) a Lojasiewicz-type theorem ensuring every bounded gradient trajectory converges to a critical point, and (ii) a stable/center-manifold theorem for the critical set B. Fact (ii) is delicate because B is not an isolated saddle: for κ < 0 the intersection B ∩ leaf is a continuum (Section 4.4), and at each B point Theorem 4.1 gives d²−d plus d−1 zero eigenvalues, so the standard hyperbolic stable-manifold theorem does not apply. The phrase 'finite union of smooth manifolds' is also in tension with the continuum of B points. A citation or a self-contained proof of these two facts is required before the 'global behavior' claim can be accepted.
- [Section 5, Remark 5.4] The same global claim for the d = 1 system is asserted without proof. Theorem 5.2 establishes that each type-B point is a hyperbolic saddle on the 3-dimensional leaf Kκ, so the stable-manifold theorem is applicable once boundedness and convergence are known. However, the paper does not prove that every trajectory remains bounded on Kκ: Remark 5.1 only bounds the products vU, zZ, zU, vZ, not the individual variables. Without boundedness and a Lojasiewicz argument, the assertions that the stable manifolds of B separate the basins of A, that they do not intersect, and that for κ = 0 the stable manifolds of O connect to the B's are not consequences of the local analysis. This is load-bearing for the claimed 'exhaustive qualitative description' in Section 5.
- [Section 5.3, proof of Theorem 5.2] The Hessian at the origin O is displayed as the matrix [[0,0,0,1],[0,0,1,0],[0,1,0,0],[1,0,1,0]], which is not symmetric. Since H(S) = −∂²L/∂S² is symmetric by definition, the (4,3) entry should be 0; in the ordering (U, z, Z, v) the last row should be [1,0,0,0]. The reported eigenvalues (1,1,−1,−1) are unchanged, so this is a local error, but it should be corrected before publication.
minor comments (4)
- [Lemmas 3.3 and 3.5] There are several typographical errors in the displayed gradient formulas: 'zT' in Lemma 3.3 should be 'zᵀ', and in Lemma 3.5 'ΛℓP' should be '(Λ)ℓp' while '(z⊤λ)p' should be '(z⊤Λ)p'.
- [Remark 4.4] The formula for γi in Remark 4.4 is garbled: it should read γi = (1 + 1/N)λi + (1/N)∑λi, matching the definition of Γ in Section 4.
- [Section 5.1, Theorem 5.1 proof] The sentence beginning 'Where for the fourth equation, closing the system, may be chosen from any equation of the original ˙S = 0. The fourth equation' is incomplete and should be rewritten.
- [Remark 5.4] The word 'bassins' should be 'basins'.
Circularity Check
No circularity: all derivations are self-contained against the loss and Gaussian moment calculations; no fitted inputs or load-bearing self-citations.
full rationale
The paper derives the gradient-flow ODEs (Theorem 3.1) directly from the loss (1) via explicit Gaussian moment computations (Lemmas 3.2–3.5 and the Appendices). The simplified system (10)–(11) is the restriction z = Z = 0 of these derived equations, and the invariant v^2 = Tr[UU^T] + kappa is proved from the ODEs themselves (Lemma 4.1). Critical points are solved from the flow, not assumed, and stability is analyzed through explicit linearization and eigenvalue calculations. The d = 1 system is likewise derived from the same loss and analyzed by direct Hessian computations. The reliance on [24] is for problem setup and comparison, with no fitted parameter or imported conclusion being presented as a prediction. The only notable gap is the unstated use of a Lojasiewicz-type convergence theorem and a stable/center-manifold theorem in Remarks 4.3 and 5.4 to pass from local stability to global basin separation; this is a rigor/correctness concern, not circularity, because those theorems are external mathematical facts and the trajectory equations do not define their conclusions into existence. No fitted parameters, no 'predictions' that reduce to inputs, and no load-bearing self-citations appear anywhere in the derivation chain.
Assumptions & free parameters
assumptions (4)
- standard math Isserlis' theorem for fourth and sixth moments of multivariate normal variables
- domain assumption Linear attention model and loss L(ξ)=1/4 E[(ξ^T Hξ - A)^2] with prompt data x_k, x_q ~ N(0,Λ), w ~ N(0,I), y_k = w^T x_k
- standard math Bounded analytic gradient flows converge to critical points (Lojasiewicz-type theorem)
- domain assumption The two analyzed instances (z=Z=0 and d=1) are representative of the structural features of the general system
Cite this review
Pith. "Pith review of Dynamical Behaviors of the Gradient Flows for In-Context Learning." pith.science (2026). https://pith.science/paper/TA3TBAMS
@misc{pith2026241216683,
author = {Pith},
title = {Pith review of: Dynamical Behaviors of the Gradient Flows for In-Context Learning},
year = {2026},
howpublished = {\url{https://pith.science/paper/TA3TBAMS}},
note = {Machine review of arXiv:2412.16683}
}
read the original abstract
We derive the system of differential equations for the gradient flow characterizing the training process of linear in-context learning in full generality. Next, we explore the geometric structure of the gradient flows in two instances, including identifying its invariants, optimum, and saddle points. This understanding allows us to quantify the behavior of the two gradient flows under the full generality of parameters and data.
Figures
Reference graph
Works this paper leans on
-
[24]
dX m=1 wmx(m) q NX n=1 dX i=1 dX j=1 dX k=1 (wiUjk + Zkwiwj)x(i) n x(j) n x(k) q !# =E
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context.Journal of Machine Learning Research, 25(49):1–55, 2024. 39 Appendix A. Calculations of the Derivatives Appendix A.1. Derivative with respect tov Recall that, ∂L ∂v =E " 1 N NX n=1 dX i=1 dX j=1 dX k=1 (ziUjk + ziZkwj + vwiUjk + vZkwiwj)x(i) n x(j) n x(k) ...
work page 2024
-
[1]
Trans- formers learn to implement preconditioned gradient descent for in-context learning
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Trans- formers learn to implement preconditioned gradient descent for in-context learning. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2023
work page 2023
-
[2]
What learning algorithm is in-context learning? investigations with linear models
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. InProceedings of International Conference on Learning Representations (ICLR), 2023
work page 2023
-
[3]
Max-margin token selection in attention mechanism
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oy- mak. Max-margin token selection in attention mechanism. InProceedings of Advances in Neural Information Processing Systems (NeurIPS), pages 48314–48362, 2023
work page 2023
-
[4]
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emer- gence, convergence, and optimality. InProceedings of Annual Conference on Learning Theory (COLT), 2024
work page 2024
-
[5]
What can transformers learn in-context? a case study of simple function classes
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), pages 30583–30598, 2022
work page 2022
-
[6]
Khashayar Gatmiry, Nikunj Saunshi, Sashank J Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi- step gradient descent for in-context learning? InProceedings of Interna- tional Conference on Machine Learning (ICML), 2024. 37
work page 2024
-
[7]
Can a transformer represent a kalman filter? arXiv preprint arXiv:2312.06937, 2023
Gautam Goel and Peter Bartlett. Can a transformer represent a kalman filter? arXiv preprint arXiv:2312.06937, 2023
arXiv 2023
Show all 24 references
-
[8]
In-context convergence of transformers
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Proceedings of International Conference on Machine Learning (ICML), 2024
2024
-
[9]
On a formula for the product-moment coefficient of any order of a normal frequency distribution in any number of variables
Leon Isserlis. On a formula for the product-moment coefficient of any order of a normal frequency distribution in any number of variables. Biometrika, 12(1/2):134–139, 1918
1918
-
[10]
Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. In Proceedings of International Conference on Machine Learning (ICML), 2024
2024
-
[11]
Can language models learn from explanations in context?arXiv preprint arXiv:2204.02329, 2022
Andrew K Lampinen, Ishita Dasgupta, Stephanie CY Chan, Kory Matthewson, Michael Henry Tessler, Antonia Creswell, James L McClel- land, Jane X Wang, and Felix Hill. Can language models learn from explanations in context?arXiv preprint arXiv:2204.02329, 2022
2022 arXiv
-
[12]
Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond
Yingcong Li, Ankit Singh Rawat, and Samet Oymak. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2024
2024
-
[13]
One-layer transformer provably learns one-nearest neighbor in context
Zihao Li, Yuan Cao, Cheng Gao, Yihan He, Han Liu, Jason M Klusowski, Jianqing Fan, and Mengdi Wang. One-layer transformer provably learns one-nearest neighbor in context. InProceedings of Advances in Neural Information Processing Systems (NeurIPS), 2024
2024
-
[14]
Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In Proceedings of International Conference on Learning Representations (ICLR), 2023
2023
-
[15]
Almost sure convergence rates analysis and saddle avoidance of stochastic gradient methods.Journal of Machine Learning Research, 25(271):1–40, 2024
Jun Liu and Ye Yuan. Almost sure convergence rates analysis and saddle avoidance of stochastic gradient methods.Journal of Machine Learning Research, 25(271):1–40, 2024. 38
2024
-
[16]
Local to global: Learning dynamics and effect of initialization for transformers
Ashok Vardhan Makkuva, Marco Bondaschi, Chanakya Ekbote, Ad- way Girish, Alliot Nagle, Hyeji Kim, and Michael Gastpar. Local to global: Learning dynamics and effect of initialization for transformers. In Proceedings of International Conference on Machine Learning (ICML), 2024
2024
-
[17]
Geometric theory of dynamical systems: an introduction
J Jr Palis and Welington De Melo. Geometric theory of dynamical systems: an introduction. Springer Science & Business Media, 2012
2012
-
[18]
Scan and snap: Understanding training dynamics and token composition in 1- layer transformer
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1- layer transformer. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2023
2023
-
[19]
JoMA: Demystifying multilayer transformers via joint dynamics of mlp and attention
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. JoMA: Demystifying multilayer transformers via joint dynamics of mlp and attention. InProceedings of International Conference on Learning Representations (ICLR), 2023
2023
-
[20]
Transformers learn in-context by gradient descents
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacra- mento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymy- rov. Transformers learn in-context by gradient descents. InProceedings of International Conference on Machine Learning (ICML), 2023
2023
-
[21]
A theoretical understanding of self-correction through in-context alignment
Yifei Wang, Yuyang Wu, Zeming Wei, Stefanie Jegelka, and Yisen Wang. A theoretical understanding of self-correction through in-context alignment. arXiv preprint arXiv:2405.18634, 2024
2024 arXiv
-
[22]
Chain-of-thought prompting elicits reasoning in large language models
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. InProceedings of Advances in Neural Information Processing Systems (NeurIPS), pages 24824–24837, 2022
2022
-
[23]
In-contextlearning with representations: Contextual generalization of trained transformers
TongYang, YuHuang, YingbinLiang, andYuejieChi. In-contextlearning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv:2408.10147, 2024
2024 arXiv
Reviewed August 11, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.