How AI settled the complexity of the oldest SGD algorithm

Abstract: In 1937, Stefan Kaczmarz proposed a simple algorithm for solving systems of linear equations. This algorithm turned out to be the earliest known example of stochastic gradient descent, a ubiquitous computing paradigm that drives the training of modern AI models such as ChatGPT and Gemini. Now, those AI models have joined forces to discover the worst-case complexity of the Kaczmarz algorithm. This paper tells the story of how it happened.

• The paper demonstrates that the randomized Kaczmarz algorithm attains an optimal O(1/ε) convergence rate for its last iterate through a novel operator-theoretic proof.
• It employs a hybrid human-AI collaboration, leveraging LLMs like Gemini and ChatGPT to integrate probabilistic averaging and functional analysis techniques.
• The analysis bridges theory and practice by resolving long-standing gaps in SGD performance for least-squares problems, independent of the matrix condition number.

Optimal Complexity of the Kaczmarz Algorithm Resolved by AI Collaboration

Introduction

The paper "How AI settled the complexity of the oldest SGD algorithm" (2606.29593) presents a comprehensive complexity analysis for the randomized Kaczmarz algorithm, the first and prototypical form of stochastic gradient descent (SGD) dating back to 1937. The paper closes a long-standing gap in understanding the worst-case convergence rate of the algorithm's last iterate, establishing that it achieves the optimal $O(1/\epsilon)$ convergence for the squared error, independent of the data matrix's condition number. Notably, the proof was produced through a hybrid human–AI collaboration leveraging state-of-the-art LLMs, specifically Gemini and ChatGPT, constituting a defining example of "AI-led research".

Background and Problem Statement

The Kaczmarz algorithm is an iterative procedure for solving linear systems $Ax = b$, which in each iteration projects the current iterate onto the solution set of a randomly chosen equation. Strohmer and Vershynin’s 2009 work introduced randomization, showing expected geometric convergence with a rate depending on the condition number of $A$. However, for practical tasks such as large-scale learning—often modeled by SGD on least-squares problems—theoretical analysis of the condition number-independent convergence rate for the Kaczmarz method’s last iterate had remained incomplete. While averaging iterates yields an $O(1/\epsilon)$ error bound, such averaging is not practised in real SGD or Kaczmarz deployments, and empirical observations consistently showed strong last-iterate performance unexplained by existing theory.

Attempts to bound the last-iterate convergence rate for Kaczmarz, a canonical quadratic minimization under noiseless sampling, have posed a central instance of a broader class of open problems regarding the "last-iterate" behavior of SGD. This includes machine learning-relevant settings such as regression with noiseless data and minimal hypotheses (e.g., Attia et al. (Hohloch, 2024), Berthier et al. (Zech et al., 2020)).

Resolution and Main Technical Contribution

The core contribution is a sharp $O(1/\epsilon)$ convergence guarantee for the randomized Kaczmarz algorithm's last iterate, for general right-hand-side vectors and all invertible data matrices, that matches the rate known for average iterates. Crucially, the bound holds with constants independent of the problem size $n$ and matrix $A$, up to natural homogeneity factors.

Theorem (Main Result):

For any $n\times n$ invertible matrix $A$ and vector $b$, the randomized Kaczmarz algorithm, after $Ax = b$0 iterations, returns $Ax = b$1 satisfying

$Ax = b$2

for some universal constant $Ax = b$3.

This result decisively resolves a question open since the method was identified as an SGD instance, confirming the practical experience that last-iterate Kaczmarz does not require iterates averaging for optimal convergence rate.

Technical Insights and Proof Structure

The proof employs an operator-theoretic framework, transforming the iteration dynamics into recurrence relations on linear operators over symmetric matrices. By expressing the expected error propagation as an application of a contraction mapping $Ax = b$4, defined as

$Ax = b$5

where $Ax = b$6 is a random orthogonal projection, the authors reduce the convergence analysis to bounding the "discrete derivative"

$Ax = b$7

in operator norm.

A substantial insight, suggested by Gemini and confirmed by ChatGPT, connects this operator to structures in functional analysis such as symmetric quantum Markov semigroups and leverages their analyticity (Ritt operator) properties. The crucial estimate established is:

$Ax = b$8

The proof utilizes probabilistic averaging, spectral norm controls, and, notably, complex analysis arguments (such as Cauchy's estimate in the context of matrices), providing an elementary and broadly applicable argument not present in standard SGD or randomized numerical linear algebra analysis.

This approach extends not only to Kaczmarz under importance sampling but, mutatis mutandis, to SGD for noiseless regression tasks on general features with a fourth-moment upper bound, matching lower bounds shown in prior work.

AI-Led Research Process

The resolution originated from a process in which autonomous LLMs, Gemini and ChatGPT, contributed critical components, with human researchers in a facilitating role. The main steps were:

1. Partial reductions by humans and AI: Initial attempts reduced the convergence assertion to a scalar recursion, which AI—specifically ChatGPT—verified for an exponent $Ax = b$9, short of the critical value $A$0 required for $A$1 convergence.
2. Autonomous attempts: Both Gemini and ChatGPT were tasked to solve the last-iterate complexity question in a "closed world" prompt without prior context. Gemini returned an affirmative assertion with heuristic operator connections and incomplete references; ChatGPT withheld a claim, stating the problem was open.
3. Cross-model synthesis: Humans distilled crucial hints from Gemini, particularly the operator-theoretic perspective and its connections to functional analysis, and supplied them to ChatGPT. Equipped with these, ChatGPT generated a novel, elementary proof—subsequently independently verified by the authors.
4. Proof verification and generalization: The final proof, relying on leave-one-out operator averages and complex analysis, was confirmed to be optimal and extended to related SGD settings.

Implications and Future Prospects

Theoretical Significance

• Closing the theory–practice gap: The result justifies and explains the empirically observed last-iterate performance of randomized Kaczmarz and provides a matching rate to averaged variants without additional assumptions.
• SGD theory: The operator analysis yields a worst-case $A$2 last-iterate bound for SGD on quadratic objectives with exact data, resolving open questions and matching established lower bounds for this setting [Berthier et al., (Zech et al., 2020)].
• Methodological synthesis: The proof imports tools from functional analysis and complex-variable techniques into stochastic optimization, suggesting further cross-fertilization could be fruitful.

Practical Impact

• Algorithmic design: Given that averaging iterates incurs extra memory and computation, the theoretical optimality of last-iterate Kaczmarz and SGD reinforces their use in large-scale applications such as distributed linear solvers and LLM training.
• AI-driven mathematical discovery: The case study demonstrates that LLMs, even without full autonomy, can function as "theorem-proving collaborators", especially when their partial, non-rigorous suggestions are curated and cross-fed by human researchers.

Future Directions

• Extension to more general SGD settings: While the proof applies to noiseless least-squares settings, analogous operator-theoretic approaches may be explored for noisy, non-quadratic, or non-convex objectives.
• AI autonomy and collaboration: The documented process provides a model for integrating LLMs, operator-theoretic intuition, and human oversight. It also sets a benchmark for further advancements in AI agents for open mathematical research tasks.

Conclusion

The paper delivers an optimal, condition-number-independent complexity analysis for the last iterate of the randomized Kaczmarz algorithm, confirming the empirically observed performance with a simple, general proof. The argument, sourced by a collaborative process between leading AI LLMs and human researchers, exemplifies a new paradigm of AI-led research and suggests both a template for automating mathematical discovery and new technical tools for stochastic optimization. This development not only completes the understanding of the foundational Kaczmarz/SGD method but also illuminates the evolving interplay between artificial intelligence and mathematical sciences (2606.29593).