Search arXiv⌕ Search

arXiv subjects

Peter Richtárik

Publications and source records attributed to Peter Richtárik.

At least 19 recordsLinked to original sources

Parallel AI Agent Collaboration as Graph-Structured Distributed Block Optimization

Parallel AI agents can produce useful local edits yet degrade a shared solution when their edits are combined. We model this coordination problem as discrete block optimization with an exactly evaluable objective and a known dependency graph. Independent sets certify compatible edits whose verified gains, measured against the same current solution, add exactly. A coordinator constructs a safe compatible candidate and can also evaluate interacting edits to capture favorable interactions. We prove that an optimal fractional coloring yields a merger that preserves a $1/χ_f(G)$ fraction of the total verified local improvement, where $χ_f(G)$ is the graph's fractional chromatic number. This guarantee is tight among mergers that select a subset of the same proposals. Under conditional local-improvement guarantees and a block error-bound condition, repeated collaboration converges geometrically in the objective gap. The rate separates task structure, agent allocation, and computation cost, yielding a linear program for agent assignment and an explicit time-cost trade-off. Code-repair experiments show that verified merging resolves conflicts between individually useful edits that cause naive merging to fail. On larger conflict-free tasks, parallel agents reduce modeled completion time at higher input-token-equivalent cost. On three public road networks, fractional integration retains 922 units of verified graph-cut improvement across 90 paired runs. For route-search agents sharing exclusive access to vehicle routes, fractional scheduling reduces measured wall time by a median of 15% across twelve paired runs, with route-cost changes varying across runs.

cs.MA↗

A Quadratic Lens on Muon's Polar Update: Benefits, Pitfalls, and the Role of Newton-Schulz

Muon updates matrix-valued parameters along approximately orthogonalized directions and often trains large language models faster than AdamW, yet it remains unclear when and why orthogonalization helps. We study PolarGD, the orthogonalized update at the core of Muon, on matrix quadratics, where optimizers can be analyzed in closed form, and compare update directions under line search to remove the effect of learning-rate tuning. Orthogonalization pays off when the gradient is far from low rank and its largest singular directions see high curvature. PolarGD reaches the minimizer in one step whenever the rows of the error are orthogonal with equal norms, for any Hessian, whereas gradient descent needs all Hessian eigenvalues to be equal; and unlike sign descent, PolarGD is unaffected by rotations of the problem. Over many steps, PolarGD acts as gradient descent preconditioned by the current iterate, so its speed is governed by the conditioning of the iterate rather than of the Hessian. This explains fast regimes as well as a slow one, caused by tiny singular values of the iterate, in which PolarGD is much slower than gradient descent. Finally, the few Newton-Schulz iterations used in Muon treat these tiny singular values as gradient descent does, so fewer iterations can converge faster than exact orthogonalization; in a transformer trained with Muon, about a quarter of the singular values of the attention matrices fall in this regime.

cs.LG↗

Local Work, Global Progress: A Convergence Theory of AI Agent Collaboration

Teams of large language model (LLM) agents are increasingly applied to complex tasks, but their collaboration protocols are mostly designed empirically, without convergence guarantees. We formalize collaborative problem solving as block optimization: each agent improves one part of a shared solution, guided by its local objective plus a correction for how its changes affect the other parts. A coordinator checks every proposal against the global objective and merges proposals that do not interfere with each other. We prove linear convergence up to an error floor that vanishes when the corrections are exact. Interfering proposals can also be merged once their combination has been checked; the rate then depends on how well the proposals actually combine, not on worst-case interactions. With GPT-OSS-120B agents on robot routing, a three-agent team finds collision-free plans in every run, better than those of a single agent that controls all robots. On a larger routing task, a six-agent team completes its search about twice as fast as the single agent (in median). Merging the best checked combination of proposals, rather than only proposals that do not interfere, cuts the team's search time by about 30%. With constants measured on the runs, our convergence bound holds in every round and follows the observed progress. We also test the theory on query planning in SQLite and on a synthetic task.

cs.MA↗

A Nesterov-Accelerated Primal-Dual Splitting Algorithm for Convex Nonsmooth Optimization

We investigate the integration of Nesterov-type acceleration into primal-dual methods for structured convex optimization. While proximal splitting algorithms efficiently handle composite problems of the form $\min_x f(x) + g(x) + h(K x)$, accelerating their convergence with respect to the smooth term $f$ is notoriously challenging, due to the rotational dynamics in the primal-dual space. Our guiding principle is that Nesterov acceleration exploits the strong convexity of the dual problem. We first revisit accelerated proximal gradient descent and give a unified analysis of its convex and strongly convex regimes, including the weak convergence of its iterates. We then propose the Accelerated Proximal Alternating Predictor-Corrector algorithm (APAPC) for the setting where $g(x) = (μ_g/2)\|x\|^2$. APAPC is fully split: it only calls $\nabla f$, $K$, $K^*$ and the proximity operator of $h$, and never solves a linear system. From a single Lyapunov inequality, we derive optimal $O(1/t^2)$ convergence rates, as well as accelerated linear rates when $μ_g > 0$, in three regimes where the dual problem is strongly convex: when $h$ is smooth, when $K^*$ is bounded below, and for linearly constrained problems. In the last two regimes, the $O(1/t^2)$ rates without strong convexity of the primal problem are, to our knowledge, the first for fully split algorithms. Finally, leveraging recent results on accelerated gradient descent, we establish the weak convergence of the primal iterates and the strong convergence of the dual iterates to a saddle point.

math.OC↗

QK-Wanda: Coupling Queries and Keys for Unstructured Pruning

Wanda (Sun et al., 2024) prunes large language models by scoring weights independently within each linear projection, although queries and keys interact through dot products. We introduce QK-Wanda, which scores query and key weights by their individual deletion costs under an unmasked pre-RoPE reconstruction objective. It augments Wanda scores with information from the opposite projection (keys for query weights, and queries for key weights), allowing both projections to share a pruning budget. Its closed-form scores require no gradients or weight updates; full pruning takes 1.3% longer than Wanda on A100 and 3.1% longer on H200 with the calibration used in our main experiments. We evaluate QK-only pruning across 15 models from TinyLlama, Llama 2, Llama 3, and Qwen2.5, spanning 0.5B-72B parameters. Relative to Wanda, QK-Wanda reduces QK reconstruction error by an average of 60% at 50% sparsity and 45% at 80%. Downstream gains depend on the model. At 80% sparsity on Llama 2 70B, WikiText-2 and C4 perplexity decrease by 20.3% and 13.5%, while mean zero-shot accuracy rises by 5.94 percentage points. Qwen2.5-72B also improves, but Llama-3.1-70B has substantially higher perplexity despite lower reconstruction error. These results show both the promise of coupled pruning criteria and the limits of local reconstruction as a predictor of model quality.

cs.LG↗

Broximal Gradient Descent: A Projection-Free Sister of Projected Gradient Descent

We propose Broximal Gradient Descent (BroxGD), a projection-free sister method to projected gradient descent for constrained optimization. Its forward--backward construction replaces the proximal backward operation by the broximal operation of Gruntkowska et al. (2025). Instead of projecting, each step minimizes a linear function over the intersection of the constraint set $\mathcal{X}$ and a ball $\mathbb{B}(x_k,t_k)$ centered at the current iterate $x_k$, of suitable radius $t_k>0$: \[ x_{k+1}\in\arg\min_{z\in\mathcal{X}\cap\mathbb{B}(x_k,t_k)}\langle\nabla f(x_k),z\rangle. \] We develop a comprehensive convergence theory spanning a wide range of optimization regimes and radius rules. We expect BroxGD to find many applications and inspire numerous extensions, much like projected gradient descent. Our contribution is theoretical; potential applications and toy experiments illustrate the method and suggest directions for future work.

math.OC↗

A Sharper Theory of Ball-Proximal Optimization: Convergence and Radius Selection

We study the exact Euclidean ball-proximal point method for proper, closed, convex functions, where each iteration minimizes the objective over a ball centered at the current point. Retaining the objective gap in the decrease of squared distance yields sharper bounds on objective values, stationarity, and the symmetric Bregman distance to a minimizer. For constant radius $t>0$ and initial distance $D_0>0$ to the solution set, the objective gap after $K$ iterations is at most its initial value multiplied by $\exp(-2K^2t^2/D_0^2)$. We characterize convergence for arbitrary positive radius sequences. If their sum diverges, the method reaches a minimizer in finitely many iterations whenever one exists; otherwise, the objective values converge to the infimum and the iterates escape every bounded set. If the radii are summable, the iterates converge to a possibly nonoptimal point. Self-contraction gives finite trajectory length for bounded iterates and a constant-radius termination bound $O_d(1+D_0/t)$, whose implicit constant depends only on the dimension. A polyhedral family shows that the dimension-independent quadratic bound remains asymptotically sharp. We also identify a minimum successful geometric decay factor and analyze adaptive radius rules based on subgradients or objective gaps, epigraph reformulation, and relaxed updates. Together, these results strengthen the foundations and convergence guarantees of the method without assuming smoothness or strong convexity.

math.OC↗

Local LMO is Secretly a Projection Method!

The local linear minimization oracle (Ferris and Zavriev, 1996; arXiv:2605.08850), or Local LMO, solves constrained convex problems without having to compute a projection: it minimizes a linear model over the intersection of the feasible set with a ball around the current iterate. We show that, whenever the ball radius does not exceed the Polyak radius, the Local LMO step is the Euclidean projection of the current iterate onto the intersection of the feasible set with a half-space that separates the iterate from the solution set; this projection is nonetheless computable by a linear oracle alone. We demonstrate that Local LMO belongs to a broader family of projection methods which may be indexed by the depth of the localizing half-space. For objectives with $\vartheta$-Hölder continuous gradient, every method from this family whose half-space lies sufficiently deep drives the best of its first $K$ iterates to optimality at the universal rate $\mathcal{O}(K^{-(1+\vartheta)/2})$, matching the non-accelerated universal gradient method of Nesterov (2015). Run at the Polyak radius, Local LMO attains the same rate when the constrained optima are also unconstrained ($\|\nabla f(x_\star)\| = 0$), and the rate $\mathcal{O}(K^{-1/(2-\vartheta)})$ otherwise.

math.OC↗

Optimizers for Diffusion Models: A Controlled Benchmark

Discrete diffusion models now match autoregressive language models on several benchmarks, while the question of how best to train them has received far less attention: the optimizer is inherited from one paper to the next and never compared. New optimizers, meanwhile, are validated almost exclusively on autoregressive pretraining, a different objective on a different loss surface. We present a controlled optimizer benchmark across four diffusion formulations, to our knowledge the first for discrete diffusion: seven optimizers (AdamW, Lion, Muon, SOAP, MARS, MARS-M, Schedule-Free) on masked diffusion (text8), uniform diffusion (QM9, and LM1B through the Gaussian duality) and Gaussian diffusion on images (CelebA-64), each on a task with published reference values. Every optimizer receives the same search protocol, and every winner is retrained at the full budget with three seeds. AdamW is a strong default but not always the right choice: it is beaten by a resolved margin on two of the four tasks, and the winner changes with the formulation, so the optimizer deserves the same care as the rest of the training recipe. Notably, methods validated on autoregressive language model pretraining transfer well: Muon, MARS-M and SOAP each beat the tuned AdamW on at least one diffusion formulation. The benchmark, all runs and every figure are reproducible end to end from the released code at https://github.com/armanbolatov/diffusion-baselines.

cs.LG↗

Super-Tuning: From Activation-Aware Pruning to Sparse Fine-Tuning

Large language models (LLMs) remain expensive to fine-tune because full-parameter updates require substantial memory, compute, and per-task storage. We study whether saliency signals originally developed for pruning can be reused to choose where a model should adapt. We propose Super, a sparse parameter-efficient fine-tuning (PEFT) method that fixes a small trainable support using a Wanda-style activation-weighted magnitude score [Sun et al., 2023] computed from a calibration pass. We then introduce Supra, a hybrid adapter that combines this sparse update with LoRA while preserving a matched trainable-parameter budget through a simple budget-splitting rule. In single-seed Math17K arithmetic experiments on Llama-3.2-1B and Meta-Llama-3-8B, the best Super/Supra variants achieve the highest average accuracy among the tested schedule-selected adapter configurations. We also include a PaFi-style magnitude-only support as a closest training-free sparse baseline and find that low-score supports under both magnitude and Wanda-style orderings can be effective. These results suggest that simple pruning-inspired orderings can provide useful fixed sparse supports for PEFT, especially when combined with low-rank adapters.

cs.LG↗

Convergence Analysis of the ProbAbilistic Gradient Estimator Algorithm for Weakly Convex Finite-Sum Optimization

The ProbAbilistic Gradient Estimator algorithm (PAGE), a stochastic algorithm introduced by Li et al. in 2021, was designed to find stationary points for the average of smooth nonconvex functions. In this work, we study PAGE within the broad framework of $τ$-weakly convex functions, providing a continuous interpolation between the general nonconvex $L$-smooth regime ($τ=L$) and the convex regime ($τ=0$). We establish new convergence rates for PAGE, showing that its complexity improves as $τ$ decreases.

math.OC↗

Convergence Analysis of Muon-type Methods with Inexact LMO in the Degenerate Case

Muon-type methods have demonstrated potentially superior performance over Adam and its variants, and have shown hyperparameter transferability across model sizes when specific norms are chosen for the LMO in deep architectures. However, while the LMO is solved approximately via iterative algorithms in practice, most convergence analyses consider the ideal case where the search direction is the exact solution to the LMO. Recently, the inexact Muon update was analyzed by Shulgin et al. [2025], which reveals a fundamental coupling between the inexactness and the optimal step size and momentum. However, the convergence is guaranteed for the non-degenerate case only, i.e., the smallest positive singular value of the rescaled momentum is assumed to be bounded below by some positive constant when the spectral norm is used. In this work, we investigate Muon-type methods with inexact LMO in the degenerate case, where the smallest positive singular value of the rescaled momentum can approach zero, for the general non-convex case and the star-convex case with weight decay. Novel assumptions are proposed to address the challenges posed by inexact LMO in such degenerate scenarios, and convergence rates are established under the layer-wise $(L^0, L^1)$-smooth assumption for both cases.

math.OC↗

SILAGE: Memory-Efficient, Full-Gradient-Free Nonconvex Optimization for Nested Finite Sums

Empirical risk minimization on massive datasets naturally exhibits a nested double finite-sum structure, where $N=nm$ total samples are logically or physically partitioned into $n$ blocks of size $m$ (e.g., in pooled data silos, out-of-core learning, or deliberate stratification). While variance-reduced methods achieve optimal oracle complexities for nonconvex objectives, they suffer from severe scaling bottlenecks in this centralized regime. Recursive estimators, such as PAGE, require periodic global full-gradient refreshes over all $nm$ samples, which are computationally expensive. Conversely, single-loop methods, such as SILVER, avoid such refreshes but require an impractical $\mathcal{O}(nm)$ memory footprint to store a control variate for every sample. In this paper, we propose SILAGE, a variance-reduced algorithm that addresses this trade-off. By actively exploiting the double-sum structure, SILAGE eliminates periodic global full-gradient refreshes over all $nm$ components (evaluating at most one local group gradient per iteration) while requiring only $\mathcal{O}(n)$ memory. Furthermore, we provide a tight convergence analysis that avoids pessimistic worst-case Lipschitz constants. Instead, SILAGE's complexity natively adapts to the underlying data geometry via nested functional similarities: across-group ($δ_1$) and within-group ($δ_2$) heterogeneity. Our results improve existing state-of-the-art bounds in several practically relevant regimes.

cs.LG↗

Stochastic Sign Descent Methods: New Algorithms and Better Theory

Various gradient compression schemes have been proposed to mitigate the communication cost in distributed training of large scale machine learning models. Sign-based methods, such as signSGD, have recently been gaining popularity because of their simple compression rule and connection to adaptive gradient methods, like ADAM. In this paper, we analyze sign-based methods for non-convex optimization in three key settings: (i) standard single node, (ii) parallel with shared data and (iii) distributed with partitioned data. For single machine case, we generalize the previous analysis of signSGD relying on intuitive bounds on success probabilities and allowing even biased estimators. Furthermore, we extend the analysis to parallel setting within a parameter server framework, where exponentially fast noise reduction is guaranteed with respect to number of nodes, maintaining $1$-bit compression in both directions and using small mini-batch sizes. Next, we identify a fundamental issue with signSGD to converge in distributed environment. To resolve this issue, we propose a new sign-based method, {\em Stochastic Sign Descent with Momentum (SSDM)}, which converges under standard bounded variance assumption with the optimal asymptotic rate. We validate several aspects of our theoretical findings with numerical experiments.

math.OC↗

TAMUNA: Doubly Accelerated Distributed Optimization under Partial Participation

In distributed optimization and federated learning, slow and costly communication between parallel devices and the central server constitutes the primary bottleneck. To alleviate this burden, two strategies have emerged: 1) local training (LT), which reduces communication frequency by performing multiple local computations between rounds, and 2) compression (CC), which consists of transmitting lower-dimensional, compact representations. Recent theoretical advances have successfully combined LT and CC to achieve doubly-accelerated communication rates, with respect to both condition number and model dimension. However, these methods have a major drawback: they require full client participation and break down when idle clients miss communication triggers. We introduce TAMUNA, the first algorithm to successfully intertwine LT, CC, and partial participation. By decoupling primal model updates from dual control variates, TAMUNA overcomes the architectural deadlock of prior methods. In the strongly convex setting, TAMUNA converges linearly to the exact solution, establishing a new state of the art by exhibiting doubly-accelerated convergence, while supporting arbitrary levels of client participation.

cs.LG↗

Demystifying Pipeline Parallelism: First Theory for PipeDream

Training modern machine learning models increasingly requires computation to be distributed across many accelerators. Data parallelism remains the default choice and is often paired with tensor-parallel sharding, but model parallelism becomes unavoidable once parameters, activations, or optimizer states no longer fit on a single device. This paper studies pipeline model parallelism through the lens of PipeDream (PD) (Harlap et al., 2018). Our first contribution is theoretical: we introduce Randomized PipeDream (RPD), a stale block-SGD abstraction that yields, to our knowledge, the first clean nonconvex convergence guarantee for a PD-style method. Our second contribution is a scaling diagnosis: we prove that the delay induced by steady-state PD grows as $S^2 - S/2 + O(1)$ for $S$ stages, so the stale-read contribution in the convergence theorem scales as $Θ(γ^2 S^4)$, equivalently as $Θ(S^4/K)$ in the tuned-rate form. Our third contribution is a comparison with LocalSGD, whose periodic model averaging trades weight staleness for synchronization bubbles. In our reported simulated-time experiments, the better-performing method depends on the objective: PD performs better on the quadratic objective and on a small language-modeling training-loss task, while for logistic regression LocalSGD becomes superior as the number of stages increases.

cs.LG↗

A Unified Primal-Dual Recipe for Accelerating Three-Operator Splitting Methods

Composite optimization problems, formulated as the minimization of three functions, are ubiquitous in large-scale machine learning and signal processing. While state-of-the-art splitting methods such as Condat-Vũ (CV) [Condat, 2013, Vũ, 2013], Primal-Dual Davis-Yin (PDDY) [Salim et al., 2022b], and Primal-Dual Twice Reflected (PDTR) [Malitsky and Tam, 2026] are highly versatile, they inherently exhibit non-accelerated convergence rates. Existing accelerated primal-dual splitting results either focus on special structures like one of the functions being zero, or linearly constrained problems, or smooth regimes, or directly use Nesterov-type momentum. Our contribution is a unified Bregman primal-dual framework that yields four variants and a common Lyapunov analysis. By applying the Chambolle-Pock algorithm [Chambolle and Pock, 2011] to primal-dual reformulations, we systematically derive four novel accelerated algorithms: Accelerated Condat-Vũ (ACV-I and ACV-II) and Accelerated Primal-Dual Twice Reflected (APDTR-I and APDTR-II). Through a simplified Lyapunov-based analysis, we establish iteration complexities for both smooth and nonsmooth cases, successfully removing the restrictive assumptions restrictions required by prior works.

math.OC↗

LOSCAR-SGD: Local SGD with Communication-Computation Overlap and Delay-Corrected Sparse Model Averaging

Communication is a major bottleneck in distributed learning, especially in large-scale settings and in federated learning environments with slow links. Three standard ways to reduce this cost are communication compression, local training, and communication-computation overlap. Methods that combine these ingredients are used in practice and have been found to be effective for large-scale training, but there is little theory for methods that combine all three. We study a heterogeneous-compute setting in which different workers may take different numbers of local steps, and we propose LOSCAR-SGD, a Local SGD method that communicates only a sparse subset of model coordinates and continues optimizing while communication is in flight. A key ingredient is a delay-corrected merge rule that incorporates delayed synchronized information without discarding the progress made during the overlap phase. We give convergence guarantees for smooth non-convex objectives and show how sparsity, overlap, and worker heterogeneity affect the rate. To the best of our knowledge, this is the first theory for this combination of ingredients. Experiments further show that communication-computation overlap reduces training time and that the delay-corrected merge outperforms naive overwriting.

cs.LG↗