Pith. sign in

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 →

arxiv 2412.16683 v1 pith:TA3TBAMS submitted 2024-12-21 math.DS

classification math.DS MSC 34D0537C1037C20
keywords in-contextlearninggradientflowlinearattentioninvariantmanifoldscriticalpointssaddleordinarydifferentialequationsstabilityofdynamicalsystems
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper is about what happens inside a linear attention model while it is trained on the in-context learning loss: it derives, in full generality, the system of ordinary differential equations that the gradient flow follows. This is a closed-form description of the training dynamics for the full linear attention model, not a restricted subcase. For two instances the paper goes further and classifies the entire dynamics: an invariant foliation confines every trajectory, the only attracting critical points are identified explicitly, and the saddles that separate their basins are shown to form a null set. The upshot is that, in these two systems, which minimum the training reaches is determined by the initial condition through a single invariant $\kappa$, and almost all initial conditions lead to one of two attractors. If these results are correct they give the first complete qualitative picture of gradient-flow training for linear in-context learning.

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.

Watch

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

Editorial extensions of the paper, not claims the author makes directly.

  • 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.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 4 minor

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)
  1. [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.
  2. [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.
  3. [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)
  1. [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'.
  2. [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.
  3. [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.
  4. [Remark 5.4] The word 'bassins' should be 'basins'.

Circularity Check

0 steps flagged · score 0.0 of 10

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 0 free parameters · 4 assumptions · 0 invented entities

No free parameters are fitted to data. The load-bearing assumptions are the Gaussian linear-attention model from prior work and an unstated convergence theorem for analytic gradient flows. No new particles, forces, or entities are postulated.

assumptions (4)
  • standard math Isserlis' theorem for fourth and sixth moments of multivariate normal variables
    Used throughout the appendix to evaluate expectations in Lemmas 3.2-3.5; cited as [9].
  • 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
    This is the setup inherited from [24]; if this is not the right model for ICL training, all downstream equations describe a different process.
  • standard math Bounded analytic gradient flows converge to critical points (Lojasiewicz-type theorem)
    Needed to justify global convergence and basin structure in Remark 4.3 and Section 5.3, but not stated or cited in the paper.
  • domain assumption The two analyzed instances (z=Z=0 and d=1) are representative of the structural features of the general system
    The paper asserts these cases contain all necessary features but does not prove that conclusions transfer to higher dimensions or to the full system.

how reviews work

0 comments
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

Figures reproduced from arXiv: 2412.16683 by the authors.

Figure 1
Figure 1. From left to right [PITH_FULL_IMAGE:figures/full_fig_p026_1.png] view at source ↗
Figure 2
Figure 2. Some trajectories when κ ̸= 0 and when κ = 0. Remark 5.4. The relative positions of the points and manifolds are as follows. We use the convention from Remark 5.3 On each 3-dimensional leaf Kκ, the unstable manifold (curve) of the saddle point Bκ(+ + ++) joins on one side the attractor A(+00+) and on the other side the attractor Aκ(0 + +0). The two-dimensional (relative to Kκ) stable manifold of Bκ(+ + ++) separates… view at source ↗
Figure 3
Figure 3. A 3D view on the trajectories κ ̸= 0. The 2-dimensional surfaces, the stable manifolds of points B do not intersect and delineate the basins of attraction of points A. ai = σi(z − τiZ), bi = v + σiU, ci = τi(v + σiU) and di = −(z − τiZ), keeping in mind that Ai needs to be linearly independent, for instance a1 = a2, b1 = b2 but c1 = −c2 and d1 = −d2. Remark 5.6. When κ ̸= 0, then "inside" Kκ there is a three-dimensi… view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

24 extracted references · 19 canonical work pages

  1. [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) ...

  2. [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

  3. [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

  4. [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

  5. [4]

    Training dynamics of multi-head softmax attention for in-context learning: Emer- gence, convergence, and optimality

    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

  6. [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

  7. [6]

    Can looped transformers learn to implement multi- step gradient descent for in-context learning? InProceedings of Interna- tional Conference on Machine Learning (ICML), 2024

    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

  8. [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

Show all 24 references
  1. [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

  2. [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

  3. [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

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

  9. [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

  10. [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

  11. [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

  12. [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

  13. [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

  14. [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

  15. [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

  16. [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

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.