Pith. sign in

REVIEW 5 cited by

Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought

Not yet reviewed by Pith; the record is open.

Re-run · record.json Download PDF Read on arXiv ↗

This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.

SPECIMEN: schema-true, not a live event

T0 review · schema-true

One-sentence machine reading of the paper's core claim.

pith:XXXXXXXX · record.json · timestamp

arxiv 2502.21212 v1 pith:R7KUQSP3 submitted 2025-02-28 cs.LG cs.AI

Jianhao Huang , Zixuan Wang , Jason D. Lee This is my paper
classification cs.LGcs.AI
keywords transformerslinearperformancepromptingtransformerchaindescentgradient
verification ladder T0 review T1 audit T2 compute T3 formal
0 comments
read the original abstract

Chain of Thought (CoT) prompting has been shown to significantly improve the performance of large language models (LLMs), particularly in arithmetic and reasoning tasks, by instructing the model to produce intermediate reasoning steps. Despite the remarkable empirical success of CoT and its theoretical advantages in enhancing expressivity, the mechanisms underlying CoT training remain largely unexplored. In this paper, we study the training dynamics of transformers over a CoT objective on an in-context weight prediction task for linear regression. We prove that while a one-layer linear transformer without CoT can only implement a single step of gradient descent (GD) and fails to recover the ground-truth weight vector, a transformer with CoT prompting can learn to perform multi-step GD autoregressively, achieving near-exact recovery. Furthermore, we show that the trained transformer effectively generalizes on the unseen data. With our technique, we also show that looped transformers significantly improve final performance compared to transformers without looping in the in-context learning of linear regression. Empirically, we demonstrate that CoT prompting yields substantial performance improvements.

Discussion (0). Sign in to comment.

Forward citations

Cited by 5 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. Scaling Latent Reasoning via Looped Language Models

    cs.CL 2025-10 unverdicted novelty 7.0 of 10

    Looped language models with latent iterative computation and entropy-regularized depth allocation achieve performance matching up to 12B standard LLMs through superior knowledge manipulation.

  2. Learning Compositional Functions with Transformers from Easy-to-Hard Data

    cs.LG 2025-05 conditional novelty 7.0 of 10

    A transformer with O(log k) layers provably learns the k-fold permutation composition task in poly(N,k) samples with curriculum or mixed easy-to-hard data, despite an SQ lower bound requiring N^{Omega(k)} samples on h...

  3. The Power of Power Law: Asymmetry Enables Compositional Reasoning

    cs.AI 2026-04 unverdicted novelty 6.0 of 10

    Power-law data sampling creates beneficial asymmetry in the loss landscape that lets models acquire high-frequency skill compositions first, enabling more efficient learning of rare long-tail skills than uniform distr...

  4. Revisiting Test-Time Scaling: A Survey and a Diversity-Aware Method for Efficient Reasoning

    cs.CL 2025-06 conditional novelty 5.0 of 10

    ADAPT, a diversity-aware prefix fine-tuning method, improves best-of-N sampling efficiency for a 1.5B reasoning model, reaching 80% accuracy at N=32 versus N=256 for the baseline.

  5. Deciphering Trajectory-Aided LLM Reasoning: An Optimization Perspective

    cs.CL 2025-05 conditional novelty 5.0 of 10

    Reasoning trajectories are formalized as pseudo-gradient descent on LLM parameters, making LLM reasoning training a MAML-style meta-learning problem.

Pith tools