Stochastic preconditioning: designing the reformulation parameters (B, D) to minimize the condition number zeta(A, B, D) in stochastic reformulations of linear systems

In the stochastic reformulation framework of Richtárik and Takáč (SIAM J. Matrix Anal. Appl., 2020), a consistent linear system Ax=b is replaced by a stochastic problem governed by two user-chosen parameters: a positive definite matrix B defining the geometry and a distribution D over random sketching matrices. The linear convergence rates of the basic, minibatch, and accelerated methods for the reformulation are all controlled by the condition number zeta(A,B,D) = lambda_max/lambda_min^+ of the matrix W = B^{-1/2} E[Z] B^{-1/2}, where Z is built from the sketched system. Of the open directions recorded when this framework was published, the missing accelerated rate for the L2 (mean-square) error and the extension to matrix inversion and accelerated quasi-Newton updates were subsequently resolved by later work; the surviving central open problem is stochastic preconditioning: the principled design or optimization of the pair (B, D) — for instance within parametric families of comparable per-iteration cost — so as to minimize or provably reduce zeta(A,B,D) for broad classes of system matrices.

Background

Solving a large consistent linear system $Ax=b$ — one that possesses at least one solution — by randomized methods is a central topic at the interface of numerical linear algebra and stochastic optimization. A unifying viewpoint is the sketch-and-project framework of Gower and Richtárik (SIAM J. Matrix Anal. Appl., 2015): at each iteration one draws a random sketching matrix $S$ from a fixed distribution $\mathcal{D}$ and moves to the point closest to the current iterate, in the norm $\|x\|_B^2 = x^\top B x$ defined by a symmetric positive definite matrix $B$, that satisfies the randomly sketched equation $S^\top A x = S^\top b$. Classical instances include the randomized Kaczmarz method (pick a random row of $A$ and project onto the corresponding hyperplane) and randomized coordinate descent (which updates a random subset of variables).

Richtárik and Takáč (2020) recast this construction as a family of stochastic reformulations of $Ax=b$: equivalent stochastic optimization, stochastic linear system, stochastic fixed point, and probabilistic intersection problems, all parameterized exactly by the pair $(B,\mathcal{D})$, and all having the same solution set as the original system under an exactness condition that the authors characterize. They analyze three algorithms for the reformulation: the basic method (interpretable simultaneously as stochastic gradient descent, stochastic Newton, stochastic proximal point, and stochastic projection with a fixed stepsize), a parallel (minibatch) method, and an accelerated method. All converge linearly, and all rates are governed by the spectrum of a single matrix

$$W = B^{-1/2}\,\mathbb{E}[Z]\,B^{-1/2}, \qquad Z = A^\top S\,(S^\top A B^{-1} A^\top S)^{\dagger} S^\top A,$$

where $\dagger$ denotes the Moore–Penrose pseudoinverse and the expectation is over $S\sim\mathcal{D}$. Writing $\lambda_{\max}$ and $\lambda^+_{\min}$ for the largest and smallest positive eigenvalues of $W$, the iteration complexities (up to logarithmic factors) are: $1/\lambda^+_{\min}$ for the L2 (mean-square) error $\mathbb{E}\|x_k - x_*\|_B^2$ of the basic method with unit stepsize; $\zeta := \lambda_{\max}/\lambda^+_{\min}$ for the expected iterate of the basic method with optimal stepsize; $\sqrt{\zeta}$ for the expected iterate of the accelerated method; and the minibatch size $\tau$ interpolates the parallel method from $1/\lambda^+_{\min}$ toward $\zeta$ (note $\lambda_{\max} \le 1$) as $\tau$ grows.

The pair $(B,\mathcal{D})$ therefore plays a role formally analogous to a preconditioner in classical iterative solvers: the rate of every method in the family is controlled by the condition number $\zeta(A,B,\mathcal{D}) = \|W\|\,\|W^\dagger\|$ of the induced matrix $W$. Richtárik and Takáč call the problem of choosing $(B,\mathcal{D})$ so that this condition number is small stochastic preconditioning, and they deliberately leave it open. Known special cases illustrate the phenomenon: for the randomized Kaczmarz method with importance sampling ($B=I$, $S=e_i$ with probability $\|A_{i:}\|_2^2/\|A\|_F^2$), $\zeta$ equals the condition number of $A^\top A$; for coordinate descent, the $\zeta$-minimizing sampling probabilities over the probability simplex are known (Richtárik and Takáč, Optim. Lett., 2016); and if one could sketch with the eigenvectors of $A$ (taking $B=A$), $\zeta$ would become independent of the conditioning of $A$ — but computing eigenvectors already solves the original problem, which exposes the core tension: the information that makes $\zeta$ small is precisely the information the iterative method is meant to avoid computing. Stochastic preconditioning is also distinct from randomized preconditioning in the Blendenpik/LSRN sense, where a random projection is used once, up front, to build a classical preconditioner for a subsequent least-squares solver, rather than to shape the per-iteration distribution of the iterative method itself.

Of the open directions recorded around this framework — besides the design problem, these were the missing accelerated rate for the L2 error, practical estimation of the spectral quantities $\lambda_{\max}$ and $\lambda^+_{\min}$, extensions to matrix inversion and randomized quasi-Newton updates, and the practical minibatch/wall-clock tradeoff — the later literature resolved the L2-acceleration question and the quasi-Newton extension (see Ref. [3]), leaving the design of $(B,\mathcal{D})$ as the surviving central open problem.

Problem Statement

In the stochastic reformulation framework, the consistent system $Ax=b$ (one possessing at least one solution) is parameterized by the pair $(B,\mathcal{D})$: a symmetric positive definite matrix $B$ defining the norm geometry, and a distribution $\mathcal{D}$ over random sketching matrices $S$. Every convergence rate in the framework is governed by the condition number

$$\zeta(A,B,\mathcal{D}) = \|W\|\,\|W^\dagger\| = \frac{\lambda_{\max}}{\lambda^+_{\min}}, \qquad W = B^{-1/2}\,\mathbb{E}[Z]\,B^{-1/2}, \qquad Z = A^\top S\,(S^\top A B^{-1} A^\top S)^{\dagger} S^\top A,$$

where $\lambda_{\max}$ and $\lambda^+_{\min}$ are the largest and smallest positive eigenvalues of $W$, the expectation is over $S\sim\mathcal{D}$, and $\dagger$ denotes the Moore–Penrose pseudoinverse.

The open problem of stochastic preconditioning, as coined and left open by Richtárik and Takáč, is: for a given system matrix $A$, or a structured class of system matrices, design the pair $(B,\mathcal{D})$ so as to minimize, or provably substantially reduce, $\zeta(A,B,\mathcal{D})$, under the constraint that the per-iteration cost of the resulting method remains comparable to that of standard choices. Two scientific directions for this problem, mirroring the source's own framing, are:

1. A constructive design: an efficiently computable rule for producing $(B,\mathcal{D})$ from $A$, or from cheaply accessible statistics of $A$, together with a certified bound on $\zeta(A,B,\mathcal{D})$ that improves on the bound attained by a stated standard choice — ideally a bound independent of the condition number of $A$ — and an analysis showing that the per-iteration cost stays comparable to a standard iteration.
2. An optimal characterization: for a stated parametric family of comparable-cost pairs that is not one for which the $\zeta$-minimizing member is already known, identification of the $\zeta$-minimizing member together with the optimal value, and a matching lower bound establishing optimality within the family.

The question is the general one across structured matrix classes rather than a single fixed instance; the source explicitly anticipates that different stochastic preconditioning strategies will be needed for different classes of problems. The companion questions recorded alongside this problem in the source — the accelerated L2 convergence rate and the extension to matrix inversion and randomized quasi-Newton updates — were resolved by subsequent work and are not part of this problem statement.

References

1. Peter Richtárik, Martin Takáč. Stochastic Reformulations of Linear Systems: Algorithms and Convergence Theory. SIAM Journal on Matrix Analysis and Applications 41(2):487–524, 2020. DOI: 10.1137/18M1179249. https://doi.org/10.1137/18M1179249 (preprint: arXiv:1706.01108, https://arxiv.org/abs/1706.01108)
2. Robert M. Gower, Peter Richtárik. Randomized Iterative Methods for Linear Systems. SIAM Journal on Matrix Analysis and Applications 36(4):1660–1690, 2015. DOI: 10.1137/15M1025487. https://doi.org/10.1137/15M1025487
3. Robert M. Gower, Filip Hanzely, Peter Richtárik, Sebastian U. Stich. Accelerated Stochastic Matrix Inversion: General Theory and Speeding up BFGS Rules for Faster Second-Order Optimization. Advances in Neural Information Processing Systems 31, 2018. arXiv:1802.04079. https://arxiv.org/abs/1802.04079
4. Dmitry Kovalev, Eduard Gorbunov, Elnur Gasanov, Peter Richtárik. Stochastic Spectral and Conjugate Descent Methods. Advances in Neural Information Processing Systems 31, 2018. arXiv:1802.03703. https://arxiv.org/abs/1802.03703
5. Robert M. Gower, Denali Molitor, Jacob Moorman, Deanna Needell. On Adaptive Sketch-and-Project for Solving Linear Systems. SIAM Journal on Matrix Analysis and Applications 42(2):954–989, 2021. DOI: 10.1137/19M1285846. https://doi.org/10.1137/19M1285846 (preprint: arXiv:1909.03604, https://arxiv.org/abs/1909.03604)
6. Peter Richtárik, Martin Takáč. On optimal probabilities in stochastic coordinate descent methods. Optimization Letters 10(6):1233–1243, 2016. DOI: 10.1007/s11590-015-0916-1. https://doi.org/1007/s11590-015-0916-1
7. Nicolas Loizou, Peter Richtárik. Momentum and stochastic momentum for stochastic gradient, Newton, proximal point and subspace descent methods. Computational Optimization and Applications 77(3):653–710, 2020. DOI: 10.1007/s10589-020-00220-z. https://doi.org/10.1007/s10589-020-00220-z
8. Yun Zeng, Deren Han, Yansheng Su, Jiaxin Xie. On Adaptive Stochastic Heavy Ball Momentum for Solving Linear Systems. SIAM Journal on Matrix Analysis and Applications 45(3):1259–1286, 2024. DOI: 10.1137/23M1575883. https://doi.org/10.1137/23M1575883 (preprint: arXiv:2305.05482, https://arxiv.org/abs/2305.05482)
9. Michał Dereziński, Daniel LeJeune, Deanna Needell, Elizaveta Rebrova. Fine-grained Analysis and Faster Algorithms for Iteratively Solving Linear Systems. arXiv:2405.05818, 2024. https://arxiv.org/abs/2405.05818
10. Johannes J. Brust, Michael A. Saunders. Structured Sketching for Linear Systems. arXiv:2407.00746, 2024. https://arxiv.org/abs/2407.00746