Machine Learning

Motivation

An LLM agent acting in an environment spends most of its wall-clock time waiting. Each step — a model call, a tool or MCP request, a browser action, sometimes a human reply — must complete before the next can be issued, and the round trips dominate end-to-end latency: a chess game between two reasoning agents runs for hours, and an operating-system tuning task for tens of minutes. When a training or prompt-optimization loop repeats such a run thousands of times, the waiting is the cost.

Speculative actions (Ye, Ahuja, Liargkovas, Lu, Kaffes, Peng, ICLR 2026) transplants a classical systems idea — speculative execution in microprocessors, and speculative decoding for LLM inference — to the agent's environment loop. A cheap, fast speculator guesses the action a slow, authoritative actor is about to produce, the guess is used to launch the next environment call early, and the work is committed only when the actor's real action confirms the guess. The interface stays sequential and lossless; the internals run in parallel.

What makes this a formalization target rather than an engineering report is the paper's §5 cost–latency analysis. Speculating more branches buys hit probability but costs tokens, and the paper derives closed-form expressions for both sides of that trade — a self-contained piece of applied probability sitting underneath a systems paper. This mission asks for those expressions, machine-checked.

Setting

Fix a horizon $T$ and index steps $t = 0, 1, \dots, T-1$. At each step a policy maps the state to an API call; the actor executes it with latency $\mathrm{Exp}(\beta)$, while the speculator proposes candidate actions with latency $\mathrm{Exp}(\alpha)$, where $\beta < \alpha$ (the speculator is faster in expectation). A speculative branch hits when the action it guesses implies the same next call the actor's true action would have implied; branches hit independently across steps with probability $p$.

Two knobs define the two regimes analyzed. Breadth $k$: at each step, launch $k$ independent one-step speculations in parallel, each immediately followed by a real call. At least one of the $k$ succeeds with probability

Depth: follow a single branch, extending it whenever a speculative or real call returns and pruning subtrees the actor contradicts.

The quantity driving both results is $S_n$, the expected number of hits by round $n$. A hit consumes the following step's speculation window — after a correct guess the next call is already cached, so no new speculation is launched there — which yields the two-term recursion

Write $T_{\mathrm{seq}}, M_{\mathrm{seq}}$ for the latency and token cost of strictly sequential execution, and $T_{\mathrm{spec}}, M_{\mathrm{spec}}$ for their speculative counterparts. In the depth regime latencies are taken deterministic: $a$ for a real call, $b < a$ for a speculative one.

Target

The goal theorem is the finite-horizon latency ratio for breadth-focused speculation (Proposition 1), with $p(k)$ abbreviated $p_k$:

The supporting targets, ordered as the analysis builds them:

1. the closed form $S_n = \frac{p}{1+p}n + \frac{p^2}{(1+p)^2}\bigl(1 - (-p)^n\bigr)$ solving the recursion;
2. the per-hit saving $\mathbb{E}[(B-A)^+] = \frac{\alpha}{\beta(\alpha+\beta)}$ for independent $A \sim \mathrm{Exp}(\alpha)$, $B \sim \mathrm{Exp}(\beta)$;
3. the $T \to \infty$ limit $1 - \frac{p_k}{1+p_k}\cdot\frac{\alpha}{\alpha+\beta}$, and the resulting 50% ceiling: the latency reduction is strictly below $\tfrac12$ for every $p_k \le 1$;
4. the cost counterpart (Theorem 4), finite-horizon and in the limit, with $\tilde k$ the number of distinct actions across the $k$ branches;
5. the depth-focused time and cost identities (Theorem 6), whose latency coefficient is $p$ rather than $\frac{p}{1+p}$ — raising the speedup ceiling from $\tfrac12$ to $1$;
6. the structure of confidence-aware selective speculation (Theorem 3 and Corollary 5): with sorted per-branch confidences, the marginal hit-probability gain is non-increasing, so the optimal breadth is the greedy threshold rule "add a branch while $\Delta^\star \delta q(m) \ge c$".

Significance

The analysis is what turns speculation from a trick into a tunable system. Proposition 1 and Theorem 4 are governed by the same quantity $p_k$, so a practitioner who can estimate hit probability can choose $k$ offline against a latency/cost budget rather than by trial. The 50% ceiling is a genuine negative result — it says breadth alone cannot do better, and motivates the depth regime, where the ceiling becomes 1. Theorem 3 explains why confidence-based branch selection is cheap in practice: the whole dynamic program collapses to one scalar continuation value, so a runtime system sorts confidences and adds branches greedily in $O(k)$ per step.

The paper's proofs are pen-and-paper and, as far as we are aware, none of these results has a machine-checked proof. Three parts reward formalization specifically. The recursion's closed form is derived by a characteristic-equation argument with a particular solution that collides with the homogeneous part — routine but error-prone. The per-hit saving is an honest two-dimensional integral over independent exponentials. And Theorem 6's cost expression is stated in the paper with a floor function and then immediately replaced by an approximation, so formalizing it forces a decision about which claim is actually being asserted (see Formalization scope).

Difficulty

The obvious first move on the recursion — guess a constant particular solution — fails, because $r = 1$ is a root of the characteristic polynomial $r^2 - (1-p)r - p$ and a constant trial collides with the homogeneous family; the particular solution is linear in $n$, and the $\frac{p^2}{(1+p)^2}$ coefficient comes out of matching both initial conditions, not one.

The interesting hypothesis is the one the recursion's shape encodes and the prose states only in passing: a hit at round $t$ removes the speculation window at round $t+1$. Drop it and the recursion becomes one-term and the answer changes.

For the per-hit saving, the difficulty is analytic rather than algebraic: the inner antiderivative of $(b-a)\alpha e^{-\alpha a}$ must be handled, and the outer integral runs over an unbounded interval, so integrability has to be established rather than assumed.

The asymptotic statements need the oscillating term $(-p_k)^{T-1}$ controlled uniformly — it is bounded, not vanishing termwise in an obvious way — before the $\tfrac1T$ prefactor can be taken to zero.

Formalization scope

Everything is over $\mathbb{R}$. The model lives in one definition bundle, Def_SpecActions_model, in namespace SpecActions; the mission's Lean names match the prose symbols ($S_n$ is hits, $p(k)$ is phit, $\tilde k$ is kt).

The model is formalized at the level the paper's own proofs use: $\mathbb{E}[T]$ and $\mathbb{E}[M]$ are defined by the expressions Appendix A derives for them (specTime, specCost, and their depth analogues), and the theorems assert the algebraic and asymptotic identities relating those quantities. Deriving those expressions from a measure-theoretic model of the execution trace is deliberately not in scope — with one exception: milestone 2 states the per-hit saving as a genuine iterated integral against the exponential densities, so the one probabilistic step the paper actually computes is formalized as an integral rather than assumed.

Conventions a solver should know before starting:

• Statements are quantified over $\alpha, \beta > 0$ and $0 \le p_k \le 1$; the standing assumption $\beta < \alpha$ is not imposed, since none of the identities need it.
• Finite-horizon statements carry $1 \le T$, and $T-1$ is natural-number subtraction — the $T = 0$ case is excluded rather than silently truncated.
• hits takes $p_k$ (the per-step hit probability $p(k)$), not the per-branch $p$; phit relates the two, and Thm_SpecActions_phit_bounds supplies the $0 \le p(k) \le 1$ range facts the other statements assume.
• Theorem 6's cost is stated as the exact identity, not the paper's approximation. The paper gives an exact expression involving $\lfloor a/b \rfloor$ and then an $\approx$ form with $\frac{a}{2b} - \frac12$; these coincide only when $a/b$ is an integer. The milestone asserts the exact floor version, which is what the proof establishes.
• The 50% ceiling is stated as the strict bound $\frac{p_k}{1+p_k}\cdot\frac{\alpha}{\alpha+\beta} < \frac12$, which holds for all admissible parameters; the paper's "upper bound of 50%, occurring when $p=1$ and $\alpha = \infty$" describes an unattained supremum.
• Theorem 3's dynamic program is formalized as the two facts that carry its content — diminishing marginal returns, and optimality of the greedy threshold breadth — rather than as a Bellman recursion over a mode process, which would require a full MDP development.

Reusable beyond this mission: the two-term linear recursion solved in milestone 1, and the $\mathbb{E}[(B-A)^+]$ computation for independent exponentials, which is a standard fact absent from Mathlib. Contributions extending the model toward an actual measure on execution traces — deriving specTime rather than defining it — are welcome as follow-on work.

Selected references

• Naimeng Ye, Arnav Ahuja, Georgios Liargkovas, Yunan Lu, Kostis Kaffes, Tianyi Peng. Speculative Actions: A Lossless Framework for Faster Agentic Systems. ICLR 2026. arXiv:2510.04371 — Proposition 1 (p. 4), Appendix A (pp. 13–14), Theorem 3 (p. 10), Theorem 4 (p. 19), Corollary 5 (p. 22), Theorem 6 (p. 23).
• Yaniv Leviathan, Matan Kalman, Yossi Matias. Fast Inference from Transformers via Speculative Decoding. ICML 2023. arXiv:2211.17192 — the speculate-verify pattern at token level.
• Wenyue Hua, Mengting Wan, Shashank Vadrevu, Ryan Nadel, Yongfeng Zhang, Chi Wang. Interactive Speculative Planning. 2024. arXiv:2410.00079 — depth-oriented speculation on a single planning branch.
• Yilin Guan et al. Dynamic Speculative Agent Planning. 2025. arXiv:2509.01920 — online RL for choosing speculation depth under a cost-latency trade-off.
• Robert M. Tomasulo. An Efficient Algorithm for Exploiting Multiple Arithmetic Units. IBM Journal of Research and Development, 1967. DOI:10.1147/rd.111.0025 — speculative execution in hardware.

Beneath Exp3, Exp4 and their relatives lies one algorithm: minimize past losses plus a convex regularizer. Chapters 26–28 of Lattimore–Szepesvári develop this unifying view. For a Legendre potential $F$ with Bregman divergence $D_F$, both mirror descent and follow-the-regularised-leader satisfy the master bound $R_n(a) \le \frac{F(a) - F(a_1)}{\eta} + \frac{1}{\eta}\sum_t D_F(a_t, \tilde a_{t+1})$; the negentropy potential on the simplex recovers Exp3 exactly. The goal theorem is the payoff for adversarial linear bandits: FTRL on the unit ball with the self-concordant-flavoured potential $F(a) = -\log(1-\|a\|) - \|a\|$ achieves $R_n \le 2\sqrt{3nd\log n}$ — improving the $\sqrt{d}$ factor over the Exp3-style approach of Chapter 27 and matching the $\Omega(d\sqrt{n})$ lower bound of Mission XI up to logarithms.

Real decisions come with context: a news site chooses an article for a particular user. Competing with the single best arm is then meaningless; the right benchmark is the best mapping from contexts to arms, or more generally the best of $M$ expert policies. Chapter 18 of Lattimore–Szepesvári formalizes this via Exp4 — exponential weighting over experts, fed by the importance-weighted estimator of Mission V. The goal theorem: with learning rate $\eta = \sqrt{2\log(M)/(nk)}$, Exp4 satisfies $R_n \le \sqrt{2nk\log M}$ against the best of $M$ experts. Since $M$ enters only logarithmically, the learner can compete with exponentially large policy classes — the conceptual gateway from bandits to reinforcement learning with function approximation.

Formal verification complete: Theorem 2 and all 12 mission targets are Proved. The main theorem has 0 open proof dependencies, including the zero-multiplier case. These results apply to the formalized independent-box model described below. Paper: arXiv v3. The revised statement formalized here is intended for v4, which will be available on arXiv shortly.

Why study budgeted Pandora’s box?

An experiment can be valuable without being worth its price. In Bayesian optimization, a function evaluation may require training a model, running a simulation, or carrying out a laboratory experiment. The decision is therefore both where to evaluate and how much information to buy. A rule that responds only to relative costs can miss the effect of making every evaluation expensive.

Xie, Astudillo, Frazier, Scully, and Terenin connect this question to Pandora's box, a sequential search problem with independent hidden rewards. Their paper uses the exact solution of this simpler model to motivate an acquisition function for Bayesian optimization. This mission formalizes the author-revised version of the paper's fixed expected-budget result, Theorem 2, including its zero-multiplier boundary. Theorem 1, the classical cost-per-sample result attributed to Weitzman, is supporting material. Paper, Sections 2–3.

Independent boxes and a fixed expected budget

There are finitely many boxes in a nonempty set $X$. Box $i$ contains a real random reward $F_i$ with a known distribution and finite absolute first moment. The rewards are mutually independent. Opening it costs a deterministic amount $c_i>0$ and reveals its reward exactly.

A policy chooses a first box, then uses its observed history and independent randomization to choose an unopened box or stop. It must open at least one box, cannot reopen a box, and cannot resume after stopping. Its terminal reward is the largest observed reward. There is no outside option of zero. Define

For a fixed real budget $B$, feasibility means $C(\pi)\le B$. This is an expected-budget constraint; individual runs may exceed $B$. Write

For $\lambda>0$, the reservation index $g_i(\lambda)$ is the unique finite solution of

At zero multiplier, the formalization defines

This is an explicit extension of the paper's positive-cost definition. A generalized Gittins policy first opens a maximal-index box. It subsequently opens a remaining maximal-index box when its index exceeds the incumbent, and stops when the incumbent exceeds every remaining index. Equality permits either action. Ties between boxes and stopping ties may be randomized. The zero branch is not defined by choosing an arbitrary root of the zero-level improvement equation. Original positive-cost rule, equation (5) and Section 3.2.

Formalization target: Theorem 2 with the zero boundary

Use the budget condition in the author-revised Theorem 2 supplied on 23 October 2026,

There is no upper budget cap: the target includes $B\ge\sum_i c_i$, as well as budgets below that sum.

Define the dual envelope

The main target is the existence of $\lambda\ge0$ minimizing this envelope on $[0,\infty)$ and a generalized Gittins policy $\pi$ satisfying

Thus the policy is optimal for the fixed expected budget. Positive multipliers require exact expected expenditure; zero multipliers permit slack. No strict reward-gap assumption is imposed on this main target, and no continuity, atomlessness, or boundedness assumption is added to the reward laws.

Correction relative to arXiv v3. Theorem 2 in that version requires a strictly positive multiplier under the numerical budget bounds alone. That conclusion has a counterexample: independent rewards $F_1=1$, $F_2\sim\mathrm{Uniform}[0,2]$, unit costs, and $B=3/2$. Opening box 2 and then box 1 exactly when $F_2<1$ achieves $R_{\mathrm{all}}=5/4$ within budget, whereas every positive-multiplier index policy has strictly smaller reward. Lemma 9's inference that every zero-penalty optimum must open all boxes is invalid. The mission targets the author-revised statement, which allows zero multipliers and removes the upper budget bound, rather than the strict-positivity claim in arXiv v3. Source statement and proof, Appendix B.5.

Supporting targets include the reservation root; Theorem 1; exact-budget and complementary-slackness certificates; convexity and attainment of the dual minimum; joint expected-cost and expected-reward limits at positive multipliers; budget matching and Gittins cost extrema; a zero-multiplier policy-limit lemma; and a feasible zero-index policy when zero minimizes the dual. Some milestones isolate or strengthen steps used in the revised proof rather than reproducing a numbered statement verbatim. The corrected positivity result remains as a supporting corollary:

What the verified result establishes

Budgeted Pandora’s box is motivated by a common feature of Bayesian optimization in practice: experiments, simulations, and model evaluations must be chosen within a limited resource budget. It asks how to choose evaluations sequentially to maximize the expected best observed reward, subject here to a bound on expected total cost. An optimal index policy for this problem provides a theoretical benchmark for budget-aware Bayesian optimization.

All 12 linked theorem targets have completed Lean proofs, including the zero-multiplier boundary and the supporting probability, policy-limit, and optimization results. The verified development establishes the author-revised expected-budget theorem and supplies reusable components for finite adaptive search and constrained optimization.

The empirical performance of the acquisition function and the extension to correlated Gaussian-process Bayesian optimization are outside the goal. No global Bayes-optimality claim is made for that correlated extension. Paper, Section 3.3.

Key steps in the formal argument

An optimal penalized policy need not meet a prescribed budget. Atoms and tied indices can make expenditure jump as the multiplier changes. The targets therefore include expected-cost extremes within the index-policy class and randomization that interpolates those costs.

At zero multiplier, finite reservation roots can escape to positive infinity. Arbitrary ties in the resulting extended indices need not respect the budget; the zero-dual-minimum milestone asserts existence of a feasible choice, rather than feasibility of all choices. Finite first moments must suffice throughout.

The convergence target records both expected reward and expected cost, including the slope needed in the envelope argument. The directional derivative of an affine objective uses its linear part; Lemma 13 in the author-revised manuscript includes this correction. Earlier proof in arXiv v3, before the revision, Appendices B.5–B.6.

Formalization scope and conventions

The Lean model uses $n+1$ boxes, integrable Borel probability laws on $\mathbb R$, their product measure, and an independent uniform seed on $[0,1]$. Measurable policies observe ordered padded histories. Finite recursion enforces irreversible stopping and distinct openings. Terminal maxima start at the first observation, not zero.

Expectations are Bochner integrals; their integrability and the boundedness of real policy suprema are established from the model in the verified supporting lemmas. Zero indices use EReal and the infimum of almost-sure upper bounds. Index conditions apply to all valid histories, including histories of probability zero. Negative multipliers are excluded by the generalized-rule definition. The strict-gap predicate is used only for the positivity corollary.

Definitions contain no proof placeholders. All twelve linked targets have status Proved, and the main theorem has zero open proof dependencies. Verification was performed by Prove2Me's server using Lean 4.33.1 with Mathlib revision 0df444a360eaa60ab8c11dca51a86af692955474.

Selected references

The paper is available as arXiv v3. The revised statement formalized here is intended for v4, which will be available on arXiv shortly.

• Qian Xie, Raul Astudillo, Peter I. Frazier, Ziv Scully, Alexander Terenin. Cost-aware Bayesian Optimization via the Pandora's Box Gittins Index. NeurIPS 2024. arXiv v3, 16 January 2025; proceedings.
• Martin L. Weitzman. Optimal Search for the Best Alternative. Econometrica 47(3), 1979, pp. 641–654. DOI.

Random pruning changes more than the expected output

A fine-tuned model can be stored as a pretrained model together with its parameter changes. Pruning these delta parameters reduces the amount of task-specific information to store. The DARE procedure independently deletes each change with probability $p$ and multiplies every surviving change by $1/(1-p)$. This preserves the expected linear-layer output, but a single pruned model can still differ substantially from that expectation.

Deng and coauthors investigate this distinction in DARE the Extreme: Revisiting Delta-Parameter Pruning For Fine-Tuned Models, ICLR 2025. The paper motivates changes to the rescaling rule and to fine-tuning regularization. This mission focuses on its finite-sample mathematical analysis: the relation between random pruning, coefficient energy, and output concentration. Its goal is the Kearns–Saul bound in Appendix E.1, equation (8), PDF p. 30, expressed through the coefficient statistics used in Section 3.2.

The distinction between that appendix result and the printed Theorem 3.1 matters. The mission does not assert the latter's piecewise formula. Its low-pruning branch omits the square root present in equation (8), and its high-pruning branch applies a one-sided refinement to a two-sided event. The exact target below retains the appendix's valid bound across the entire interval $0<p<1$.

A fixed layer and a random mask

Fix one output coordinate of a linear layer, an input vector $x$, and a delta-weight row $\Delta W$. There are $n>0$ input coordinates. Define the deterministic influence coefficients $c_j=\Delta W_jx_j$, their sum $S=\sum_jc_j$, and their energy $Q=\sum_jc_j^2$. Equivalently, the formal statements quantify over every real coefficient vector $c$; choosing $x_j=1$ realizes every such vector in the layer model.

The only randomness is the pruning mask. Write $\omega_j=1$ for a dropped coordinate, with mutually independent $\omega_j\sim\operatorname{Bernoulli}(p)$. A mask has probability

Expectations and event probabilities are the finite weighted sums against $w_p$. A surviving coordinate is rescaled by $1/q$, where $q>0$. The output error, original minus pruned, is

DARE uses $q=1-p$; denote its error by $H$. This is the retention-mask formulation in Section 3.2, equation (2), PDF p. 5, with $\delta_j=1-\omega_j$. The empirical coefficient mean and variance are $\bar c=S/n$ and $\sigma^2=n^{-1}\sum_j(c_j-\bar c)^2$. These statistics describe a fixed vector, not another source of randomness.

Formalization targets

Define the concentration coefficient with its removable singularity filled in:

For every $0<p<1$ and failure probability $0<\gamma<1$, the goal is

This is Appendix E.1, equation (8), PDF p. 30, followed by the unnumbered energy identity on PDF p. 31. Zero coefficients are included: no positive-energy assumption is attached to the goal.

Four supporting milestones state the following results.

1. Coefficient statistics: $Q=n(\bar c^2+\sigma^2)$ for $n>0$, as used in the final algebraic step of Appendix E.1, PDF p. 31.
2. Exact moments: for $0\le p\le1$ and $q>0$, let $b_q=(1-(1-p)/q)S$. Then $\mathbb EH_q=b_q$, $\mathbb E(H_q-b_q)^2=p(1-p)Q/q^2$, and $\mathbb EH_q^2=b_q^2+p(1-p)Q/q^2$. This is a paper-derived extension of the calculations on PDF p. 29 to the general rescaling model introduced in Appendix E.2, PDF p. 31. In particular, DARE has mean zero and mean square $pQ/(1-p)$.
3. Kearns–Saul exponential moment: $0<\Phi(p)\le1/2$ and, for every real $t$,

The analytic input is Berend–Kontorovich, Section 3, Theorem 4, equation (6), PDF pp. 3–4. The bound on $\Phi$ is also stated in the DAREx appendix on PDF p. 30. 4. Exponential output tail: for $Q>0$ and $t>0$,

This is the unnumbered display immediately preceding equation (8), PDF p. 30.

What the result establishes

The target quantifies the error of a randomly selected pruned layer in terms of its actual influence coefficients. It distinguishes preserving an expectation from controlling a realization. The moment identities also expose the bias introduced by choosing a rescaling denominator different from $1-p$.

Formalization supplies a precise probability model and checks every coefficient, sign, and exceptional case. The finite mask model and its normalization already compile locally with proofs. The five milestone and goal statements have been elaborated, but their theorem proofs remain open. Completing this mission would formalize the selected appendix result; it would not establish the paper's experimental accuracy claims, a whole-network guarantee, or an optimal rescaling rule.

Why the tail direction matters

Signed coefficients require exponential-moment control for both positive and negative arguments. The sharper estimate in Berend–Kontorovich, Lemma 5, equation (9), PDF p. 4 has a nonnegative-argument restriction. Using it for an unrestricted absolute tail loses an essential hypothesis.

For example, with $n=1$, $c_1=1$, $p=99/100$, and $\gamma=1/200$, the error is $1$ with probability $99/100$ and $-99$ with probability $1/100$. The printed Theorem 3.1 threshold is $\sqrt{198\log400}<99$, so its failure probability exceeds $\gamma$. This concrete source audit is the reason for selecting equation (8), not a claim that the printed theorem has been formally disproved in Lean.

Formalization scope

The Lean model uses real coefficients indexed by Fin n and Boolean functions for masks. Nonnegative masses and normalization are proved from the product formula; concentration is never assumed in a structure field. Fixed weights and inputs are external data. Random training, dependence between masks, nonlinear activations, structural pruning, and empirical validation are outside this mission.

The main goal requires $n>0$ for the empirical statistics, $0<p<1$ for DARE rescaling, and $0<\gamma<1$ for the confidence level. The moments permit empty coefficient vectors and endpoint probabilities because they use a separate positive $q$. The exponential-tail milestone requires $Q>0$ to avoid division by zero; the main goal includes $Q=0$. The value $\Phi(1/2)=1/2$ is explicit. No theorem relies on Lean's total division or logarithm to supply a missing analytic hypothesis.

Selected references

• Wenlong Deng, Yize Zhao, Vala Vakilian, Minghui Chen, Xiaoxiao Li, Christos Thrampoulidis. DARE the Extreme: Revisiting Delta-Parameter Pruning For Fine-Tuned Models. ICLR 2025. arXiv:2410.09344v2. Section 3.2, PDF p. 5, equation (2), Theorem 3.1; Appendix E.1, PDF pp. 28–31, Theorem E.1 and equations (6)–(8); Appendix E.2, PDF p. 31, initial unnumbered rescaling identity.
• Daniel Berend and Aryeh Kontorovich. On the Concentration of the Missing Mass. Electronic Communications in Probability 18 (2013). arXiv:1210.3248v1. Section 3, PDF pp. 3–4, Theorem 4 and equation (6); Lemma 5 and equation (9) explain the excluded one-sided refinement.

Why initialization matters for transfer learning

Transfer learning starts with a representation learned on an earlier task and adapts it to a new one. Two common choices are linear probing, which changes only the final linear predictor, and fine-tuning, which changes the representation as well. These procedures optimize related training objectives, but their behavior away from the training data can differ. Kumar and coauthors study this distinction through two-layer linear networks, alongside experiments with nonlinear networks. This mission formalizes their perfect-feature LP-FT result, rather than the empirical claims or the general imperfect-feature comparison. See Section 3.4, Proposition 3.7, PDF p. 10.

LP-FT first learns a head by linear probing and then uses that head to initialize full fine-tuning. The perfect-feature setting isolates the effect of head initialization: the representation already contains exactly the features needed to predict the labels, but the head initially need not use them correctly. The mathematical question is whether joint training preserves or loses the representation's ability to predict outside the observed training subspace.

Linear predictors, training data, and OOD loss

An input is a vector $x\in\mathbb R^d$. A feature extractor is a matrix $B\in\mathbb R^{k\times d}$, and a head is a vector $v\in\mathbb R^k$. Together they predict $v^\top Bx$, with effective weight vector $B^\top v$. The fixed matrix $X\in\mathbb R^{n\times d}$ contains the $n$ training inputs as rows. Their span is $S=\operatorname{rowspace}(X)$, with dimension $m$.

The ground truth has orthonormal-row features $B_\star$ and a nonzero head $v_\star$. Write $w_\star=B_\star^\top v_\star$ and $Y=Xw_\star$. Perfect pretrained features mean $B_0=UB_\star$ for an orthogonal matrix $U$. The corresponding aligned head is $u=Uv_\star$. The dimensions satisfy $1\le k\le m$ and $m+k<d$.

The two geometric assumptions require the orthogonal projections from $R_0=\operatorname{rowspace}(B_0)$ into $S$ and into $S^\perp$ to be injective. In this dimension regime these are exactly the positive largest-principal-angle cosine conditions used by the paper. They demand more than two subspaces having some nonorthogonal directions. The Lean definition spells out injectivity of $v\mapsto\Pi_T B_0^\top v$ for each $T\in\{S,S^\perp\}$. See Definition 3.2 and Appendix A.1, PDF pp. 7 and 22-23.

An out-of-distribution law $\mu$ is any probability measure on $\mathbb R^d$ with a finite second moment and positive-definite uncentered second-moment matrix $\Sigma=\mathbb E_\mu[xx^\top]$. Its mean need not be zero. Define

Both training methods use the unnormalized loss $\widehat L(v,B)=\|XB^\top v-Y\|_2^2$. Fine-tuning follows its gradient flow in both parameters; linear probing keeps $B=B_0$. Time is real and nonnegative. These are the paper's equations (3.2)-(3.3), PDF p. 6.

Formalization targets

The goal is Proposition 3.7 in an explicit nonzero-signal regime. For $\sigma>0$, initialize an FT head with independent Gaussian coordinates, $v_0\sim\mathcal N(0,\sigma^2I_k)$. Establish

Linear probing, from any initial head, must converge to $u$. Fine-tuning initialized at its limit must satisfy

The probability-one event applies to all times simultaneously. The goal also asserts existence of the relevant global flows; a conditional claim about a possibly nonexistent trajectory would not suffice. The statement does not assert a numerical error lower bound or a positive time-infimum.

Seven milestones supply the supporting results: global flow existence and FT uniqueness; unchanged features orthogonal to the training span; the balancedness invariant; the second-moment identity for OOD risk; almost-sure Gaussian head misalignment; exact LP recovery; and stationarity after LP initialization. The principal source is Appendices A.2 and A.7, PDF pp. 23-31 and 45-47.

What the result establishes

The result distinguishes two initializations of the same joint-training procedure. In this idealized setting, a head obtained by linear probing gives zero OOD loss throughout subsequent fine-tuning, while a Gaussian head almost surely has positive OOD loss at every finite time. The conclusion concerns population squared prediction error, not classification accuracy or a finite test-set estimate.

The paper establishes the mathematical claim; this mission asks for a Lean proof of the stated model and result. The scope is deliberately limited to perfect pretrained features. It does not claim an LP-FT upper bound for imperfect features, which the authors identify as a further challenge in Section 3.4, PDF p. 10. A completed development would also provide reusable components for finite dimensional gradient flows, factorized linear models, and population risk.

Why the proof needs the training dynamics

The training loss alone does not select a unique effective predictor in an overparameterized problem. Knowing that a predictor fits the observed examples therefore does not determine its OOD loss. Formalization must track the head and feature extractor together, and it must distinguish parameter stationarity from a claim that a derivative happens to vanish at one time. The Gaussian conclusion also requires one event controlling an uncountable set of times; separate probability-one statements for individual times would be weaker.

Formalization scope and conventions

Vectors use Mathlib's finite dimensional real Euclidean spaces. Matrices are represented as continuous linear maps, with Euclidean adjoints and operator norms. The feature update is written explicitly as the Frobenius-gradient equation; it is not a gradient with respect to the operator norm. Differentiability is imposed within $[0,\infty)$, including the right derivative at zero.

The dimensions, nonzero target, positive Gaussian scale, finite second moments, and projection injectivity are explicit. The nonzero target restricts the formalization to the regime of the Gaussian alignment argument in Lemma A.12; $k\le m$ makes the identifiability condition used in Proposition A.20 precise. The random-head law is the scaled standard Gaussian measure. No randomness of the fixed training matrix or independence from an additional data draw is assumed.

The model contains no assumed convergence, invariant, or desired risk bound. Each of those is a theorem obligation. The well-posedness milestone makes explicit an analytic prerequisite of the source's flow notation. The risk milestone uses the identity in (A.29)-(A.32), avoiding the reversed inequality printed in (A.28). The quantitative constant in Theorem 3.3 is outside this mission. Source-aligned proofs and the supporting analysis infrastructure are welcome; changing the learning rule or assuming a milestone inside the model would change the task.

Selected references

• Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang, Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution, ICLR 2022, arXiv:2202.10054v1. Main target: Section 3.4, Proposition 3.7, PDF p. 10, equations (3.10)-(3.11); proof: Appendix A.7, PDF pp. 45-47, Proposition A.20 and (A.208)-(A.218). Supporting invariants: Appendix A.2, PDF p. 24, Lemmas A.3-A.4, equations (A.15)-(A.20). Gaussian alignment: Appendix A.3, PDF pp. 34-35, Lemmas A.11-A.12.

Why the geometry of a minimum needs a parameter convention

A trained neural network is used through its predictions, while its parameters are the coordinates in which training takes place. Distinct parameter vectors can describe exactly the same prediction function. A proposed explanation of generalization based on the shape of the parameter-space loss therefore needs to account for these equivalences. This mission concerns Hessian sharpness: the spectral norm of the matrix of second derivatives of the loss at a minimum. The question is whether that number is intrinsic to the predictor or can change without changing any prediction.

Dinh, Pascanu, Bengio, and Bengio establish that, for a one-hidden-layer rectified network, every sufficiently differentiable critical minimum with nonzero Hessian has equivalent parameterizations of arbitrarily large Hessian sharpness. The formalization target is their Section 4.2, Theorem 4, PDF pp. 5–6. The goal theorem and both supporting milestones now have accepted Lean proofs on Prove2Me. This public research-paper mission is complete.

The immediate historical sequence is:

• 2016–2017: Keskar and collaborators reported numerical evidence relating large-batch training, sharp minima, and a generalization gap. This is empirical context, not an assumption or a theorem to be proved in this mission. ICLR 2017 paper.
• March 2017: Dinh and collaborators released a mathematical analysis of parameter symmetries and several flatness measures. The fixed source for this mission is their May 2017 revision, arXiv version 2, which determines the theorem numbering and PDF page citations. Version history.

Networks, losses, and reciprocal layer scaling

Fix positive integers $d$ and $h$, the input dimension and hidden width. Let $W\in\mathbb R^{d\times h}$ and $v\in\mathbb R^h$ be the network's weights. For an input $x\in\mathbb R^d$, define

This is a bias-free network with one hidden layer, rectified activation, and a scalar linear output. Its parameter vector $\theta=(W,v)$ has $n=dh+h$ coordinates and the Euclidean norm. A function-based loss is a real-valued functional $\ell$ of the entire prediction function, giving $L(\theta)=\ell(f_\theta)$. The Hessian targets assume $L$ is continuous. These conventions come from Sections 2–3, PDF pp. 2–4, Definition 3.

Two parameters are observationally equivalent if their predictions agree on every input. For a positive real number $\alpha$, the layer scaling is

The definition rescales the incoming and outgoing weights in opposite ways. It is the transformation in Section 3, Definition 5, PDF p. 4.

Write $DL(\theta)$ for the first Fréchet derivative and $H_L(\theta)=D(DL)(\theta)$ for the second derivative. The local differentiability condition requires $L$ to be differentiable throughout a neighborhood of $\theta$, with $DL$ differentiable at $\theta$. This states the regularity needed for the paper's Hessian notation explicitly, without requiring global smoothness or continuous second derivatives. A critical local minimum satisfies $DL(\theta)=0$ and has no smaller loss in some neighborhood; it may belong to a continuum of minima.

Symmetry, derivative transformation, and the goal theorem

The first milestone formalizes the scaling consequence of positive ReLU homogeneity in Theorem 1 and Definition 5, PDF p. 4:

and $T_\alpha\theta$ is a local minimum precisely when $\theta$ is. These statements hold for every parameter and every $\alpha>0$; the milestone itself requires no loss regularity.

The second milestone is Theorem 3, Section 4.2, PDF p. 5. At every point satisfying the stated local differentiability condition,

for all parameter directions $u,w$. The same regularity holds at the rescaled point. This is the coordinate-free form of the paper's gradient formula and Hessian congruence with $D_\alpha=\operatorname{diag}(\alpha^{-1}I_{dh},\alpha I_h)$.

The goal is Theorem 4. If $\theta$ is a critical local minimum with the stated differentiability and $H_L(\theta)\ne0$, then

For that same rescaling, predictions and loss agree with the original ones, and the transformed parameter remains a critical local minimum with the required differentiability. The subscript $2$ denotes the Euclidean spectral norm. The source statement and its interpretation are in Section 4.2, PDF pp. 5–6. The relevant displays in all three targets are unnumbered.

What the formalization establishes

The result separates prediction behavior from this particular numerical measure of parameter-space curvature. Pointwise identical predictors have identical function-based evaluation losses wherever those are defined, even though their Hessian sharpness can be made arbitrarily large under the stated conditions. This is the scope of the obstruction: it concerns the unnormalized Euclidean Hessian norm and this rescaling symmetry. It does not assert that training reaches every point on a scaling orbit or supply a numerical bound on test error. Theorem 4 and following discussion, PDF p. 6.

The completed Lean development connects an explicitly computed ReLU prediction function to its actual first and second derivatives. It provides reusable results about invariant losses, transport of local minima, and curvature under linear changes of parameters. All three published theorem statements were proved unchanged and accepted on 26 September 2026. Both milestones are complete, and the goal has no remaining open leaves.

The analytic obligations

ReLU is not differentiable at every activation boundary. The loss-level local regularity must therefore be retained rather than inferred from the network's syntax. A nonzero symmetric matrix alone is also insufficient for the sharpness claim: local minimality supplies an additional sign constraint. Finally, the second derivative is an operator on the entire parameter space; bounds for a single arbitrarily chosen scalar model do not establish the quantified network result. These are obligations of the formal proof, not hypotheses assuming the desired transformation identities.

Formalization note: scope and conventions

Parameters are represented by EuclideanSpace over the disjoint union of first-layer matrix coordinates and output-weight coordinates. This retains the sum-of-squares geometry rather than a product maximum norm. The Hessian is the derivative of the actual Fréchet derivative, represented as a continuous bilinear form. Its operator norm is the spectral norm under the Euclidean/Riesz identification. The regularity condition excludes reliance on default derivative values at nondifferentiable points.

The statement covers every positive input dimension and hidden width, arbitrary function-based losses with the specified regularity, and arbitrary qualifying minima. It imposes no separate nonzero-weight or active-neuron assumption. An input or output weight block may vanish when the loss hypotheses still hold. The nonzero-Hessian requirement is essential. Scaling by zero or a negative number is outside the claim. No probability model, random initialization, sampling measure, or training algorithm is assumed.

The mission focuses on the single-hidden-layer Hessian result. Volume flatness, the multiple-eigenvalue deep-network theorem, and other reparameterizations remain separate results in the paper. The local environment is Lean 4.30.0 with supported Mathlib revision c5ea00351c28e24afc9f0f84379aa41082b1188f.

Selected references

• Laurent Dinh, Razvan Pascanu, Samy Bengio, Yoshua Bengio. Sharp Minima Can Generalize For Deep Nets. ICML 2017, PMLR 70:1019–1028. arXiv:1703.04933v2. Primary anchors: Section 2, PDF p. 2; Section 3, PDF pp. 3–4, Definitions 3–5 and Theorem 1; Section 4.2, PDF pp. 5–6, Theorems 3–4.
• Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, Ping Tak Peter Tang. On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima. ICLR 2017. arXiv:1609.04836v2. Historical context.

Why an initialization kernel matters

A neural network is nonlinear in its parameters, but a small change in those parameters changes its predictions through a Jacobian. The neural tangent kernel is the Gram kernel of that Jacobian: it records which changes in predictions can be produced by common parameter updates. An initialization limit identifies a deterministic object behind this random kernel. It supplies a mathematical starting point for studying wide networks through kernel methods, before addressing the additional question of how the kernel changes during training. Jacot, Gabriel, and Hongler establish this initialization limit in Section 4.1, Theorem 1, PDF p. 5.

The requested result is already a theorem of the paper. The open work here is its formal proof in Lean, including the probability model and the order of limits. The mission is classified as OpenProblem at the request of its proposer; that label does not assert that the underlying mathematical result remains an unsolved research question.

The source appeared in 2018 and was published at NeurIPS 2018; this formalization fixes arXiv version 4, dated February 10, 2020, so that page references and conventions remain stable. Its Appendix A explicitly distinguishes the sequential limit proved there from a possible stronger simultaneous-width limit.

Networks, randomness, and the two kernels

Fix positive integers $d$ and $q$, the input and output dimensions, a hidden-layer count $h\ge0$, a bias scale $\beta>0$, and a Lipschitz function $\sigma:\mathbb R\to\mathbb R$. Write $L=h+1$ for the number of affine layers, with $n_0=d$, $n_L=q$, and positive hidden widths $n_1,\ldots,n_h$. The parameters are all entries of the weight matrices and bias vectors. Every parameter is sampled independently from $\mathcal N(0,1)$.

For an input $x$, let $a^{(0)}(x)=x$ and define

At each hidden layer, $a^{(\ell)}=\sigma(z^{(\ell)})$ coordinatewise. The output is $f_\theta(x)=z^{(L)}(x)$, with no final activation. These are the conventions of Section 2, PDF pp. 2–3. Matrix storage order in Lean uses destination then source; the displayed operation is unchanged.

For output coordinates $k,k'$, set

The sum includes every weight and every bias, as in Section 4, PDF p. 5.

The covariance kernel starts with

Given a centered Gaussian pair $(U,V)$ with covariance matrix obtained by evaluating $\Sigma^{(\ell)}$ on $(x,y)$, define

The deterministic limiting NTK is

These recurrences appear in Section 4.1, Proposition 1 and Theorem 1, PDF p. 5; their displays are unnumbered.

Formalization targets

The goal is Theorem 1, on every fixed finite family $X=(x_1,\ldots,x_N)$ of inputs and for every $\varepsilon>0$:

Here $\delta_{kk'}$ is one when the output coordinates coincide and zero otherwise. The limit takes $n_1\to\infty$ first, then $n_2\to\infty$, through $n_h\to\infty$, precisely as specified in Appendix A, PDF p. 11, and Appendix A.1, PDF pp. 12–13.

Two supporting milestones expose the required mathematical content. First, the recursively defined $\Sigma^{(L)}$ has positive semidefinite Gram matrices on every finite input family and satisfies $\Sigma^{(L)}(x,x)\ge\beta^2$. This is a paper-derived well-definedness obligation for the covariance in Proposition 1. Second, Proposition 1 asserts joint convergence in distribution of $(f_{\theta,k}(x_i))_{i,k}$ to the centered Gaussian vector with covariance $\Sigma^{(L)}(x_i,x_j)\delta_{kk'}$. This expresses the paper's independent output Gaussian processes through all their finite-dimensional distributions.

What completing the mission would establish

The result connects an explicitly parameterized random finite network to a deterministic kernel computed from its activation and depth. All output correlations, bias contributions, and layer normalizations remain visible in that connection. It would provide a checked foundation on which a separate training-stability development could build.

The formal contribution is the passage from finite random Jacobians to the limiting kernel. It does not assume that the empirical NTK already equals its limit. Nor does the goal claim convergence of a training trajectory, positive definiteness on a sphere, an early-stopping guarantee, or a generalization bound; those are separate results and questions in the source.

Where the mathematical work lies

The parameter space changes with the widths. Outputs, hidden activations, and their parameter derivatives are dependent random quantities, so the limit of their products requires more than a scalar law of large numbers. The weak Gaussian limit alone also does not justify substituting an arbitrary discontinuous derivative into expectations. The Lipschitz-only hypothesis is part of the target and must be retained, including for nonsmooth activations. Remark 3, PDF p. 5 identifies almost-everywhere differentiation as the relevant convention.

Formalization scope and conventions

Lean represents parameter coordinates by a finite dependent index carrying a layer, destination neuron, and optional source neuron; the missing source denotes a bias. Initialization is the finite product of Mathlib's standard Gaussian measures. Network outputs are obtained by the displayed recursion, and NTK entries use actual Fréchet derivatives in coordinate directions. Mathlib's derivative is zero where differentiation fails; the proof must establish that the exceptional parameters are null under the stated initialization law.

Centered Gaussian laws use Mathlib's multivariateGaussian, including singular covariance. The covariance-validity milestone must justify its covariance interpretation. No invertibility, distinct-input, smooth-activation, or positive-definite-kernel assumption is added. Positive actual hidden widths are indexed as $w_i+1$ for $w_i\in\mathbb N$, a cofinal reindexing. Depth one, constant activations, repeated or zero inputs, and empty finite families are included. The empty family is harmless because the same theorem quantifies over every nonempty family as well.

The sequential filter puts the last hidden width outermost. For two hidden widths, a probability tolerance is met by first choosing a threshold for $n_2$, then a threshold for $n_1$ that may depend on $n_2$. There is no uniform limit over the entire input space and no simultaneous-width assertion. The development uses Lean 4.30.0 and supported Mathlib revision c5ea00351c28e24afc9f0f84379aa41082b1188f. The model and statements are locally checked; the three theorem proofs remain open. Contributions to Gaussian covariance consistency, finite-dimensional distribution limits, almost-everywhere network differentiation, and the NTK limit are in scope.

Selected references

• Arthur Jacot, Franck Gabriel, and Clément Hongler, Neural Tangent Kernel: Convergence and Generalization in Neural Networks, Advances in Neural Information Processing Systems 31, 2018. arXiv:1806.07572v4. Primary anchors: Section 2, PDF pp. 2–3; Section 4.1, PDF p. 5, Proposition 1, Theorem 1, Remarks 2–3; Appendix A and A.1, PDF pp. 11–13. Relevant displays have no equation numbers.

Why stochastic-gradient convergence matters

Training a machine-learning model often means choosing a vector of parameters to minimize an average loss. Evaluating the full gradient can require processing an entire dataset. A stochastic-gradient method instead updates the parameters using a random direction obtained from a smaller amount of information. Its computational appeal raises a mathematical question: which assumptions on those directions and the stepsizes guarantee progress, and what kind of convergence follows?

This mission formalizes the core convergence theory in Section 4 of Bottou, Curtis, and Nocedal, Optimization Methods for Large-Scale Machine Learning. The results distinguish strongly convex objectives, where expected objective error can be controlled, from general smooth objectives, where the guarantee concerns gradients. They also distinguish constant stepsizes, which leave a noise-dependent error bound, from diminishing stepsizes.

Historical timeline

• 1951: Robbins and Monro introduced stochastic approximation for finding a root using noisy observations. Their work is the historical foundation for the stepsize conditions used here; the original root-finding theorem is not a separate target of this mission. Original paper.
• 2016: Bottou, Curtis, and Nocedal released the first version of their survey, organizing stochastic-gradient theory around smoothness and moment assumptions. arXiv record.
• 2018: The revised survey appeared in SIAM Review. This mission fixes arXiv version 3 for stable theorem numbering and PDF page citations. Published article.

The mathematics targeted here is already proved in the literature; the task is its Lean formalization, not a claim that these convergence results are unresolved research conjectures.

Objective, algorithm, and probability model

Let $F:\mathbb R^d\to\mathbb R$ be differentiable with an $L$-Lipschitz gradient, where $L>0$. On a probability space $(\Omega,\mathcal A,\mathbb P)$, let $\mathcal H_k$ contain the history before step $k$. Starting from a deterministic vector $w_0$, the algorithm uses positive deterministic stepsizes $\alpha_k$ and random directions $g_k$ to update

The state $w_k$ is measurable with respect to $\mathcal H_k$; the direction $g_k$ is measurable with respect to $\mathcal H_{k+1}$ and has a finite second moment. Conditional assertions hold almost surely. This uses the adapted-process formulation expressly permitted in footnote 4, PDF p. 22, rather than requiring independent sample seeds. The local index $k=0$ corresponds to the paper's $k=1$. Algorithm 4.1 and footnote 4.

Write $\mathbb E_k$ for conditioning on $\mathcal H_k$. The moment conditions use constants $\mu_G\ge\mu>0$ and $M,M_V\ge0$:

Define $M_G=M_V+\mu_G^2$. The iterates lie almost surely in an open region on which $F\ge F_{\inf}$ for a real lower bound $F_{\inf}$. These are Assumptions 4.1 and 4.3, PDF pp. 23–24, equations (4.6)–(4.9). Source.

Formalization targets

The goal is Theorem 4.10, Section 4.3, PDF p. 33, equations (4.30a)–(4.30b). Suppose

For $A_K=\sum_{k=0}^{K-1}\alpha_k$, establish both

There is no convexity assumption and no restriction that the initial stepsizes already satisfy a small-step bound. Theorem 4.10.

Four milestones capture the surrounding theory. Lemma 4.4 gives the two successive conditional expected-descent inequalities (4.10a)–(4.10b), PDF pp. 24–25. It is a common input to the convergence results. Theorem 4.6 gives a geometric upper bound for strongly convex objectives with a constant stepsize. Theorem 4.7 gives the corresponding $O(1/k)$ bound for $\alpha_k=\beta/(\gamma+k+1)$. These are parallel strongly convex targets, not prerequisites for the nonconvex goal. Theorem 4.8 gives finite-horizon sum and average squared-gradient bounds for general objectives with a constant stepsize. Section 4.

For example, if $c>0$ is the strong-convexity constant, $d\ge1$, and $0<a\le\mu/(LM_G)$, Theorem 4.6 states, with $B=aLM/(2c\mu)$ and $F_*=\inf_xF(x)$,

The upper bound tends to $B$; this does not assert that the actual expected error tends to $B$. Every milestone retains the paper's constants and all displayed conclusions. Equations (4.13)–(4.14), PDF p. 26.

What a completed formalization provides

The result explains precisely how noise and stepsize interact. The general-objective goal guarantees that the stepsize-weighted expected squared gradients average to zero even when $M>0$. The strongly convex milestones quantify objective error and the effect of initialization. These results concern the stated quantities; they do not assert convergence of iterates, global optimality for a nonconvex objective, or almost-sure convergence. Sections 4.2–4.3.

A completed development would provide reusable Lean results for smooth objective functions, conditional moment bounds, stochastic updates, and expected convergence. The proposed statements are open proof obligations. Compilation establishes that the definitions and statements are well formed, not that their convergence claims have already been proved.

Mathematical and formal difficulties

Finite conditional expectations must be connected to unconditional integrals without relying on total-function defaults. The infinite-horizon result also requires careful handling of a finite initial segment: square summability gives eventually small steps, not a bound at every step. Strong convexity must justify the objective-gap estimates and the properties of the optimum. Treating a descent recurrence as a hypothesis would omit the analytic content that this mission is intended to formalize.

Formalization scope

The model uses real finite-dimensional Euclidean space, Mathlib gradients, filtrations, Bochner conditional expectations, and ordinary real integrals. The probability space is arbitrary; no finite-support or standard-Borel restriction is imposed. Directions have explicit finite second moments, making the finite-expectation convention in the paper visible. Square integrability of iterates and integrability of losses are consequences to establish, not extra model fields.

Strong convexity uses Mathlib's StrongConvexOn, equivalent here to Assumption 4.5's first-order inequality. The optimum is defined as the infimum of the range of $F$, with its finiteness to be derived in the strongly convex branch. That branch requires $d\ge1$: the paper's deduction $c\le L$ implicitly uses a nontrivial space. The nonconvex statements allow $d=0$. Zero noise is allowed. Finite-horizon averages require $K>0$; the value assigned at $K=0$ does not affect an asymptotic limit.

Contributions to the conditional-descent infrastructure and any of the four milestones are welcome. Variance reduction, Newton-type methods, and the remainder of the survey are outside this initial mission.

Selected references

• Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization Methods for Large-Scale Machine Learning. SIAM Review 60(2), 223–311, 2018. arXiv:1606.04838v3; DOI. All page numbers above refer to the 95-page arXiv v3 PDF.
• Herbert Robbins and Sutton Monro. A Stochastic Approximation Method. Annals of Mathematical Statistics 22(3), 400–407, 1951. DOI. Historical background only.

Why control variates matter in federated optimization

Federated optimization trains one model using objectives held by many clients. Performing several local updates between communication rounds saves communication, but clients with different objectives can move in different directions. SCAFFOLD maintains a correction for each client to address this disagreement, including when only some clients participate. Karimireddy et al. establish convergence without a bound on how similar the client objectives are. The relevant result is Theorem III, Section 5, PDF p. 5.

SCAFFOLD: Convergence with Client Sampling concerns that known result's formalization. Its exact targets are finite-round bounds from the Appendix E analysis, covering strongly convex, general convex, and nonconvex objectives. The paper's FedAvg results and its separate quadratic-acceleration theorem are outside this scope.

Setting: local updates and sampled clients

There are $N\ge1$ clients, with differentiable objectives $f_i:\mathbb R^d\to\mathbb R$. The global objective is $f(x)=N^{-1}\sum_i f_i(x)$. Every client gradient is $\beta$-Lipschitz, where $\beta>0$. A stochastic gradient query is conditionally unbiased and has conditional squared-error expectation at most $\sigma^2$, with $\sigma\ge0$. Current queries on different clients are conditionally independent given the full history. These make precise the fresh-oracle interpretation of A4–A5, Appendix B.1, PDF p. 14.

Each of $T\ge1$ rounds selects a uniform subset $A_r$ of $S$ clients, where $1\le S\le N$. Each selected client takes $K\ge1$ local steps. Let $x^r$ be the server model, $c_i^r$ its stored client control variates, and $c^r=N^{-1}\sum_i c_i^r$. With constant local and global step sizes $\eta_l>0$ and $\eta_g\ge1$, a client starts at $y_{i,0}^r=x^r$ and uses

Option II replaces a participating client's control with $K^{-1}\sum_{k=0}^{K-1}g_{i,k}^r$ and retains every other control. These are Algorithm 1, PDF p. 4, and Appendix E, equations (18)–(21), PDF pp. 25–26. Write $h=K\eta_l\eta_g$ for the effective step size.

Formalization targets

Four milestones support a root goal that is the conjunction of the two convergence statements below.

GradientGrowth states, for convex clients and a global minimizer $x^\star$,

This is Appendix B.1, equation (9), PDF p. 14. PerturbedStrongConvexity states, for each $\mu$-strongly convex client and $\mu\ge0$,

as in Appendix C, Lemma 5, PDF p. 17.

ConvexFiniteRoundConvergence allows arbitrary deterministic initial controls $c_i^0$. Assume all clients are $\mu$-strongly convex with $\mu\ge0$, including ordinary convexity when $\mu=0$, and fix a minimizer $x^\star$ of $f$. Define

For $h\le1/(81\beta)$ and $\mu h\le S/(15N)$, the target is

This is a finite-round formulation of Appendix E.1, Lemma 15, PDF p. 30; its $\mu=0$ case is the first bound on PDF p. 31.

NonconvexFiniteRoundConvergence assumes a lower bound $f_\mathrm{lower}$ for $f$ and initializes every client control with $K$ fresh gradients at $x^0$. Put $F_0=f(x^0)-f_\mathrm{lower}$. For $h\le(S/N)^{2/3}/(24\beta)$, establish

This is the explicit finite-round consequence targeted from Appendix E.2, Lemma 19, PDF p. 34, with the initialization specified on PDF pp. 31 and 35. The full-client warm start's communication cost is separate from these $T$ optimization rounds.

What the result establishes

The bounds quantify optimization progress under stochastic gradients, several local steps, and partial participation. They permit arbitrarily different client objectives within the stated smoothness and convexity assumptions. Their output is a weighted random server iterate, represented by its expected loss, or a uniform random server iterate for the squared-gradient guarantee; this follows the output convention in equation (22), PDF p. 26.

The mathematical convergence analysis is published. The open work here is a Lean proof under the explicit stochastic-run model. The proposed statements do not themselves supply a machine-checked convergence proof. A completed development would provide reusable results about sampled finite averages, adaptive gradient queries, and optimization with stored noisy controls.

Where the difficulty lies

Local gradients are evaluated at different client states, and inactive clients retain controls computed in earlier rounds. Consequently, treating the server update as a centralized stochastic-gradient step discards both local disagreement and stale-control error. The challenge is to control these quantities while preserving their dependence on earlier randomness, as reflected in Appendix E.1–E.2, PDF pp. 26–35.

The source also needs careful transcription. The printed Theorem VII nonconvex noise term on PDF p. 25 differs from Lemma 19 in its smoothness and sampling factors, while output indices differ between equation (22) and the finite sum on p. 31. The exact targets above use the proof's finite-round constants and explicit indexing; they do not claim a verbatim formalization of every printed asymptotic rate.

Formalization scope

SCAFFOLD.Space is finite-dimensional real Euclidean space; SCAFFOLD.Problem records differentiable client objectives and smoothness. SCAFFOLD.Run records a standard Borel probability space, full-history filtration, measurable square-integrable iterates and oracle samples, and the concrete updates. Virtual local paths are generated before a fresh uniform size-$S$ subset is sampled, as permitted by the paragraph after equation (22), PDF p. 26.

The conventions include $S=N$, $K=1$, zero noise, and $\mu=0$. The output indices are precisely $0,\ldots,T-1$. Convex initial controls are deterministic; the nonconvex branch uses the specified stochastic warm start. The run definition assumes no descent inequality or convergence conclusion. Contributions should establish the four named propositions and necessary analysis/probability infrastructure while retaining these semantics.

Selected references

• Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh, SCAFFOLD: Stochastic Controlled Averaging for Federated Learning, ICML 2020, PMLR 119:5132–5143; arXiv:1910.06378v4, revised 2021. Primary scope: Theorem III, PDF p. 5; A3–A5 and (9), p. 14; Lemma 5, p. 17; Appendix E, pp. 25–35, especially Lemmas 15 and 19.

Why local training needs a convergence guarantee

Federated optimization studies learning when data and computation are spread across clients. Communicating after every stochastic gradient step can be costly, so clients often take several steps before averaging their models. The difficulty is that different clients can optimize different objective functions. Their models then move apart between communication rounds. A convergence guarantee must account for both stochastic gradient noise and this disagreement. Wang et al. give an explicit analysis of this tradeoff in A Field Guide to Federated Optimization, Section 6.1.

The relevant historical sequence is the introduction of FedAvg by McMahan et al. in 2017, the development of local-SGD convergence analyses reviewed by Wang et al., and the unified illustrative analysis in the 2021 field guide. The present target is that guide's known convex convergence theorem, rather than a new conjecture about arbitrary federated learning. The remaining open task is a Lean proof of the stated result. The mission is categorized as ResearchPaper: it formalizes a specific known result rather than proposing a new mathematical conjecture.

Setting: full participation with uniform weights

There are $M\ge1$ clients and a parameter vector in $\mathbb R^d$. Client $i$ has a differentiable convex objective $F_i$. The global objective is $F(x)=M^{-1}\sum_{i=1}^M F_i(x)$. Every local gradient is $L$-Lipschitz for the same $L>0$. Fix a global minimizer $x^\star$ of $F$ and a deterministic initial model $x_0$, and write $D=\|x_0-x^\star\|$.

Every client participates in every round. Each of $T\ge1$ rounds consists of $\tau\ge1$ local steps with constant learning rate $\eta$. If $x_i^{t,k}$ is client $i$'s state after $k$ local steps of round $t$, its next state is $x_i^{t,k+1}=x_i^{t,k}-\eta g_i^{t,k}$. At the next round all clients restart from the average of the preceding round's terminal states. The shadow iterate is $\bar x^{t,k}=M^{-1}\sum_i x_i^{t,k}$; it is defined even at local steps where clients do not communicate.

On a probability space $(\Omega,\mathcal A,\mathbb P)$, the history before a step contains all past oracle draws. The stochastic gradients are conditionally unbiased, their conditional squared errors have expectation at most $\sigma^2$, and the different clients' current gradients are conditionally independent. The heterogeneity bound is $\|\nabla F_i(x)-\nabla F(x)\|\le\zeta$ for every client and every point, where $\sigma,\zeta\ge0$. These are the assumptions of Section 6.1.1, PDF p. 40, equations (11)–(14), with the history and independence convention used explicitly in Appendix D.1 immediately after equation (27), PDF p. 87.

Formalization targets

The principal goal is Theorem 1, equation (15). For $0<\eta\le1/(4L)$, establish

Two source milestones describe the intermediate results. Lemma 1 bounds the conditional average loss within a round by the decrease of squared distance to the minimizer, plus a noise term and a sum of client disagreements. Lemma 2 bounds each conditional squared disagreement by $18\tau^2\eta^2\zeta^2+4\tau\eta^2\sigma^2$. Both are stated on PDF p. 41, Section 6.1.2, with proofs in Appendix D, PDF pp. 86–88. They remain genuine open proof obligations; the algorithm model does not assume either estimate.

An additional milestone records the same theorem's tuned-step consequence, equations (16)–(17), in the regime $D,\sigma,\zeta>0$. With

the same expected loss is at most

The principal goal includes zero-noise and zero-heterogeneity cases. The extra positivity conditions apply only to the printed tuned-step formula, whose denominators otherwise require separate conventions.

What this establishes

The bound separates an initial-distance term, a noise term improved by the number of clients, and two costs of local updates. It quantifies how local work interacts with stochastic noise and differing client objectives. Its conclusion concerns the average objective gap along the post-update shadow sequence; it does not assert the same bound for every last iterate or for the average of client losses. These distinctions follow directly from the quantity defined in equation (14).

A completed formalization would provide reusable checked components for stochastic optimization: finite client averages, history-conditioned oracle assumptions, per-round potential estimates, and disagreement bounds. The paper supplies the mathematical proof; this draft supplies checked statements and definitions. No convergence proof is claimed by creating or compiling the proposal.

Where the difficulty lies

The average update evaluates each gradient at its own client's state. It therefore does not directly equal a centralized stochastic-gradient step at the shadow iterate. A proof must control that discrepancy quantitatively, preserve the conditioning on the round's starting history, and justify the $1/M$ noise improvement using independent client sampling. Ignoring the sampling relationship can invalidate the advertised bound even for scalar quadratic objectives.

Formalization scope

The model uses finite-dimensional real Euclidean space, including the harmless zero-dimensional case, and a standard Borel probability space. A filtration indexed by $t\tau+k$ records the full past. Local states and gradients carry explicit measurability and finite-second-moment conditions. These probability conventions support actual Bochner and conditional expectations; integrals are not treated as arbitrary total functions without analytic obligations.

The local objectives, global objective, gradients, iterates, and shadow averages are concrete functions. The model assumes neither a drift bound nor a progress bound. Source hypotheses are uniform in the model point, and the minimizer is an actual minimizer of the averaged objective. Unequal weighting, partial client participation, nonconvex objectives, adaptive step sizes, and privacy mechanisms are outside this particular theorem. Contributions should prove the named source lemmas or their necessary analytic infrastructure while retaining these statements.

Selected references

• Jianyu Wang et al., A Field Guide to Federated Optimization, 2021, arXiv:2107.06917v1, Section 6.1.1–6.1.2, PDF pp. 40–41; Appendix D, PDF pp. 86–88.
• H. Brendan McMahan et al., Communication-Efficient Learning of Deep Networks from Decentralized Data, AISTATS 2017, arXiv:1602.05629, the FedAvg algorithm cited by the field guide. This is historical context, not an additional target.

Motivation

Transformer sequence models modulate attention scores by a function of the relative position between a query and a key: the attention logit between token $m$ and token $n$ is scaled by a fixed profile $f(p_m - p_n)$. Realizing this modulation the naive way means evaluating $f$ once per pair $(m,n)$ and materializing an $L\times L$ adjustment over the whole sequence — quadratic in the sequence length $L$. Rotary Position Embedding (RoPE) Su et al. 2021 avoids this entirely: it rotates the query at position $p_m$ and the key at position $p_n$ independently, each by an angle depending only on its own position, so that the pairwise quantity $f(p_m-p_n)$ falls out of the dot product of the two separately rotated vectors — $f$ is never evaluated pairwise, and no $L\times L$ matrix is ever built. This is what lets RoPE stay a linear, per-token preprocessing step compatible with efficient (sub-quadratic) attention implementations, rather than an $O(L^2)$ modulation. The catch is that RoPE's specific log-linear frequency schedule bakes in one particular profile: a monotone, decaying $f$. That schedule is a poor fit whenever the correlation structure of the data is not monotonically decaying with distance — the leading example being periodicity: in sequential recommendation, interactions separated by exactly one day or one week are more correlated than interactions separated by, say, half a day, and a decaying profile cannot express that "attention comes back" at the period.

Chen, Ainslie, Choromanski et al., ClockRoPE: Random Fourier Rotations for Temporal Routine Modeling (arXiv:2607.26369), ask a more general question first: which attention-modulation profiles $f$ can be realized at all by this same per-token, pairwise-iteration-free rotation trick — rotate each token once, on its own, and let the pairwise profile emerge from the dot product — and by what rotation-frequency schedule? Their answer is a random-features construction — sample the rotation frequencies from the kernel's own Fourier transform, rather than fixing them log-linearly — that realizes any continuous, normalized, positive-definite profile in expectation, with a quantified concentration rate, all while keeping the exact same per-token rotate-then-dot-product computation RoPE already uses. ClockRoPE is the periodic instance of this general theory, later deployed in a production-scale generative-retrieval system.

Setting

Fix an embedding dimension $d = 2n$ and group a vector $v \in \mathbb{R}^d$ into $n$ consecutive feature pairs $v^{(j)} = (v_{2j}, v_{2j+1}) \in \mathbb{R}^2$ for $j = 0, \dots, n-1$. For an angle $\theta$, let

be the $2\times 2$ (Givens) rotation matrix. A real kernel $f : \mathbb{R} \to \mathbb{R}$ is positive definite if for every finite family of points $x_1,\dots,x_N \in \mathbb{R}$ and complex coefficients $c_1,\dots,c_N$, $\sum_{i,j} \overline{c_i} c_j f(x_i - x_j)$ has nonnegative real part; it is normalized if $f(0) = 1$. When $f$ is also continuous and Lebesgue-integrable, its Fourier transform