# Minimal error momentum Bregman-Kaczmarz

Dirk Lorenz, Maximilian Winkler  Dirk A. Lorenz Institute of Analysis and
Algebra, TU Braunschweig, d.lorenz@tu-braunschweig Maximilian Winkler Insitute
of Analysis and Algebra, TU Braunschweig, maximilian.winkler@tu-
braunschweig.de

(November 2022)

###### Abstract

The Bregman-Kaczmarz method is an iterative method which can solve strongly
convex problems with linear constraints and uses only one or a selected number
of rows of the system matrix in each iteration, thereby making it amenable for
large-scale systems. To speed up convergence, we investigate acceleration by
heavy ball momentum in the so-called dual update. Heavy ball acceleration of
the Kaczmarz method with constant parameters has turned out to be difficult to
analyze, in particular no accelerated convergence for the
$\mathcal{L}_{2}$-error of the iterates has been proven to the best of our
knowledge. Here we propose a way to adaptively choose the momentum parameter
by a minimal-error principle similar to a recently proposed method for the
standard randomized Kaczmarz method. The momentum parameter can be chosen to
exactly minimize the error in the next iterate or to minimize a relaxed
version of the minimal error principle. The former choice leads to a
theoretically optimal step while the latter is cheaper to compute. We prove
improved convergence results compared to the non-accelerated method. Numerical
experiments show that the proposed methods can accelerate convergence in
practice, also for matrices which arise from applications such as
computational tomography.

##  1 Introduction

We consider the problem

$\displaystyle\min\varphi(x)\qquad\text{s.t.
}\mathbf{A}x=b,\qquad\mathbf{A}\in\mathbb{R}^{m\times n},\ b\in\mathbb{R}^{m}$
(1)

with a $\sigma$-strongly convex function
$\varphi\colon\mathbb{R}^{n}\to\mathbb{R}$. To invoke stochastic optimization,
we make use of the stochastic reformulation of the linear system given by

$\displaystyle\min f(x):=\mathbb{E}[f_{S}(x)],$ (2)

where the expectation is taken over so-called _sketching matrices_
$S\sim\mathcal{D}$ which stem from a distribution $\mathcal{D}$ on the set of
real-valued matrices with $m$ rows, and

$\displaystyle f_{S}(x):=\frac{1}{2}\|S^{T}(\mathbf{A}x-b)\|_{2}^{2}.$

The corresponding stochastic “sketched Bregman method” for problem (1) with
stochastic representation (2) is given by the update

$\displaystyle x_{k+1}^{*}$ $\displaystyle=x_{k}^{*}-\alpha_{k}\nabla
f_{S_{k}}(x_{k}),$ $\displaystyle x_{k+1}$
$\displaystyle=\nabla\varphi^{*}(x_{k+1}^{*}).$ (3)

Here, $\alpha_{k}>0$ is a sequence of step sizes, $S_{k}$ is drawn from the
distribution $\mathcal{D}$ and $\varphi^{*}\colon\mathbb{R}^{n}\to\mathbb{R}$
the convex conjugate function, which is everywhere finite and differentiable
with $\sigma^{-1}$-Lipschitz continuous gradient. Method (3) generalizes the
_block Bregman-Kaczmarz_ method, initially named _block sparse Kaczmarz_ [13,
18, 20], which iterates

$\displaystyle x_{k+1}^{*}$
$\displaystyle=x_{k}^{*}-\alpha_{k}\mathbf{A}_{(i_{k})}^{T}(\mathbf{A}_{(i_{k})}x_{k}-b_{(i_{k})}),$
$\displaystyle x_{k+1}$ $\displaystyle=\nabla\varphi^{*}(x_{k+1}^{*}).$ (4)

Here, the matrix $\mathbf{A}$ and the vector $b$ are partitioned as

$\displaystyle\mathbf{A}=\begin{pmatrix}\mathbf{A}_{(1)}\\\ \vdots\\\
\mathbf{A}_{(c)}\end{pmatrix},\qquad b=\begin{pmatrix}b_{(1)}\\\ \vdots\\\
b_{(c)}\end{pmatrix}.$

Note that, if $\mathbf{A}_{(i_{k})}$ consists of rows $r,r+1,...,s$, we
recover the update by setting

$\displaystyle
S_{k}=\begin{pmatrix}\ldots&e_{r}&e_{r+1}&\ldots&e_{s}&\ldots\end{pmatrix}$

in (3), with the corresponding unit vectors. The block Bregman-Kaczmarz method
(4) uses only a selected number of rows of $\mathbf{A}$ and if $x_{0}^{*}=0$
and the step size is chosen as
$\alpha_{k}=\sigma\|\mathbf{A}_{S_{i_{k}}}\|_{F}^{-2}$, it is known to
converge linearly in expectation to the solution of problem (1), see [20]. For
instance, with the $1$-strongly convex function
$\varphi(x)=\lambda\|x\|_{1}+\frac{1}{2}\|x\|_{2}^{2}$, the map
$\nabla\varphi^{*}$ is the soft-shrinkage function

$\displaystyle
x\mapsto\mathcal{S}_{\lambda}(x)=\mathrm{sign}(x)\cdot\max(|x|-\lambda,0)$

and the method converges to the solution of the _regularized basis pursuit
problem_

$\displaystyle\min\varphi(x)=\lambda\|x\|_{1}+\frac{1}{2}\|x\|_{2}^{2}\qquad\text{s.t.
}\mathbf{A}x=b,$

which is sparse for $\lambda>0$ large enough, see [4, 13]. For the choice
$\varphi(x)=\frac{1}{2}\|x\|_{2}^{2}$, we recover the _block Kaczmarz method_

$\displaystyle
x_{k+1}=x_{k}-\alpha_{k}\mathbf{A}_{(i_{k})}^{T}(\mathbf{A}_{(i_{k})}x_{k}-b_{(i_{k})}).$

###  1.1 Contribution and outline

As a first result we show in Section 3 that the iterates $x_{k}$ of (3)
converge linearly in expectation to the solution of (1), if the $S_{k}$ are
independent samples from an appropriate distribution $\mathcal{D}$ and
$\alpha_{k}$ is a suitable sequence of step sizes.

Next, in Section 4 we enrich method (3) by heavy ball momentum following the
_minimal error_ strategy. That is, we consider a parameterized update

$\displaystyle x_{k+1}^{*}(\alpha,\beta)$
$\displaystyle=x_{k}^{*}-\alpha\nabla f_{S_{k}}(x_{k})+\beta d_{k}^{*},$
$\displaystyle x_{k+1}(\alpha,\beta)$
$\displaystyle=\nabla\varphi^{*}(x_{k+1}^{*}(\alpha,\beta)),$

and actually would like to find $\alpha_{k}=\alpha$ and $\beta_{k}=\beta$ such
that the Bregman distance

$\displaystyle
D_{\varphi}^{x_{k+1}^{*}(\alpha,\beta)}(x_{k+1}(\alpha,\beta),\hat{x})=\varphi(\hat{x})-\varphi(x_{k+1}(\alpha,\beta))-\langle
x_{k+1}^{*}(\alpha,\beta),\hat{x}-x_{k+1}(\alpha,\beta)\rangle$

to the solution $\hat{x}$ of (1) is minimized over $\alpha$ and $\beta$. This
idea is inspired by the _exact-step sparse_ Kaczmarz method [13], which for
the special case $S_{k}=e_{i_{k}}$ and $\beta=0$ selects the step size
$\alpha_{k}$ with the least Bregman distance to the solution $\hat{x}$ and has
been demonstrated to significantly accelerate convergence in examples.
However, we will argue that this problem can be computationally hard to solve.
Therefore, in Section 4.1 we look at an easier problem, where we fix the step
size $\alpha=\alpha_{k}$, which can be stochastic and dependent on expressions
known at iteration $k$, and optimize just over $\beta$.

While our approach at first glance seems to require knowledge of the solution
$\hat{x}$, we can circumvent this problem by introducing an auxiliary variable
$s_{k}\in\mathbb{R}$, which is recursively updated together with
$x_{k},x_{k}^{*}$ and in each iteration equals
$\langle\hat{x},x_{k}^{*}-x_{k-1}^{*}\rangle$. It results that the momentum
parameter $\beta_{k}$ is determined by the one-dimensional convex optimization
problem

$\displaystyle\beta_{k}\in\operatorname*{\arg\\!\min}_{\beta\in\mathbb{R}}\varphi^{*}(x_{k}^{*}-\alpha_{k}\nabla
f_{S_{k}}(x_{k})+\beta\big{(}x_{k}^{*}-x_{k-1}^{*})\big{)}-\beta s_{k}$ (5)

and the next $s_{k}$ can be found by setting
$s_{k+1}=\beta_{k}s_{k}-\alpha_{k}b_{i_{k}}$. We propose the corresponding
heavy ball update as _randomized Bregman-Kaczmarz with exact momentum_ (BK-
EM).

In addition, in Section 4.2, since problem (5) is still costly to solve in
general, we propose a relaxed approach by minimizing a quadratic upper bound
of the Bregman distance
$D_{\varphi}^{x_{k+1}^{*}(\alpha,\beta)}(x_{k+1}(\alpha,\beta),\hat{x})$
jointly over $\alpha$ and $\beta$. This approach gives “relaxed exact” step
sizes $\alpha_{k}$ and momentum parameters $\beta_{k}$ and we introduce this
method as _Bregman-Kaczmarz method with relaxed exact momentum_ (BK-REM).

Finally, Section 5 shows numerical experiments to illustrate the speedup in
number of iterations and computational time.

###  1.2 Related work

In this article, we use the general stochastic reformulation of the linear
system through sketching matrices as in [24]. The framework of sketching
matrices was first proposed by [6].

The standard error to be considered in convergence analysis for methods of
randomized Kaczmarz - type and its accelerations is the
$\mathcal{L}_{2}$-error $\mathbb{E}[\|x_{k}-\hat{x}\|_{2}^{2}]$, where
$\hat{x}$ is the solution of (1) and the expectation is taken over all
possible combinations of sketching matrices. The seminal paper by Strohmer and
Vershynin [21] proved the rate

$\displaystyle\mathbb{E}[\|x_{k}-\hat{x}\|_{2}^{2}]\leq\big{(}1-\frac{\sigma_{\min}(\mathbf{A})^{2}}{\|\mathbf{A}\|_{F}^{2}}\big{)}^{k}\|x_{0}-\hat{x}\|_{2}^{2}.$

for the iterates $x_{k}$ of the Kaczmarz method, where $S_{k}=\\{e_{i_{k}}\\}$
and the $i$-th row $a_{i}$ is selected with probability
$\|a_{i}\|_{2}^{2}/|\mathbf{A}\|_{F}^{2}$. For both the Kaczmarz and the
Bregman-Kaczmarz method, is has been demonstrated that update averaging with
appropriate weights gives a better convergence rate [14, 22]. Accelerated
rates have been also established for enhanced multistep versions [8, 17] and
accelerations of the Kaczmarz method of Nesterov type [5, 10, 11], see also
[16] for an extension to systems of linear inequalities $\mathbf{A}x\leq b$.
However, strict accelerated rates for the $\mathcal{L}_{2}$-error have not
been shown for heavy ball momentum for the Kaczmarz method with a fixed
momentum parameter. To the best of our knowledge, a convergence speedup has
only been established for the quantity
$\|\mathbb{E}[x_{k}-\hat{x}]\|_{2}^{2}$, see [12, Theorem 4], which is the
error of the expected iterates and different from the above
$\mathcal{L}_{2}$-error. Moreover, linear convergence of the $\mathcal{L}_{2}$
error has only been guaranteed for relatively restricted choices of the
parameter, see [9, 12, 15]. In [2], accelerated rates have been shown for
heavy ball acceleration of minibatch Kaczmarz with sufficiently large batch
sizes. Very recently, acceleration of Tseng type has been studied for the
Bregman-Kaczmarz method via its interpretation as a dual coordinate descent
method [23]. Also recently, in [24], heavy ball acceleration with an adaptive
momentum parameter has been derived for the Kaczmarz method. There, the
authors propose to use the momentum parameter which moves the update closest
to the exact solution, which is computable due to the linearity of the system.
We call this idea the minimum error approach. It was shown that the method
from [24] converges at least with the Strohmer-Vershynin rate. In our work, we
transfer this approach to the Bregman-Kaczmarz method and obtain _exact
momentum_ (SRK-EM) and _relaxed exact momentum_ (SRK-REM) in the minimum error
sense in Bregman distance.

##  2 Notation and Preliminaries

###  2.1 Notation

We set $\mathbb{N}=\\{0,1,...\\}$. By $\mathcal{R}(B)$, $I_{n}$ and $e_{i}$,
we denote the range of a matrix $B$, the identity matrix of dimension $n$ and
the $i$-th unit vector, where its dimension is clear from the context. The
linear space generated by $S\subset\mathbb{R}^{n}$ is written as $\langle
S\rangle$. The orthogonal projection onto a linear space
$V\subset\mathbb{R}^{n}$ is denoted by $P_{V}\colon\mathbb{R}^{n}\to V$. If
$V=\langle\\{x\\}\rangle$, we use the short-hand notation $P_{V}=P_{x}$.

###  2.2 The stochastic reformulation

In this section, we make assumptions on the distribution $\mathcal{D}$ of the
sketching matrices which ensure exactness of the stochastic reformulation (2)
of the linear system $\mathbf{A}x=b$ and give useful lemmas which deal with
the sketching.

######  Assumption 1.

The matrix $\mathbb{E}_{S\sim\mathcal{D}}[SS^{T}]$ has only finite entries and
is positive definite.

Under assumption (1), the reformulation is exact, as the following Lemma
states.

######  Lemma 2.1.

[24, Lemma 2.2] Let Assumption 1 hold. Then, the minimizers of (2) are exactly
the solutions of the linear system $\mathbf{A}x=b$.

A central object for the convergence analysis will be the matrix

$\displaystyle
M:=\mathbb{E}_{S\sim\mathcal{D}}\big{[}\frac{SS^{T}}{\|\mathbf{A}^{T}S\|_{2}^{2}}\big{]}.$
(6)

######  Assumption 2.

The matrix $M$ in (6) has only finite entries.

Note that Assumption 2 requires in particular that $\mathbf{A}^{T}S$ is
nonzero $\mathcal{D}$-almost surely. We discuss the definition of $M$ and
Assumptions 1 and 2 in two examples.

######  Example 2.2.

  1. (i)

Single row sketching. If $S=e_{i}$ is sampled with probability $p_{i}$ such as
in the randomized Kaczmarz method, we compute

$\displaystyle\mathbb{E}[SS^{T}]=\sum_{i=1}^{m}p_{i}e_{i}e_{i}^{T}=\begin{pmatrix}p_{1}&&&\\\
&p_{2}&&\\\ &&\ddots&\\\ &&&p_{m}\end{pmatrix}$

and

$\displaystyle
M=\sum_{i=1}^{m}p_{i}\frac{e_{i}e_{i}^{T}}{\|a_{i}\|_{2}^{2}}=\begin{pmatrix}\frac{p_{1}}{\|a_{1}\|_{2}^{2}}&&&\\\
&\frac{p_{2}}{\|a_{2}\|_{2}^{2}}&&\\\ &&\ddots&\\\
&&&\frac{p_{m}}{\|a_{m}\|_{2}^{2}}\end{pmatrix}.$

Hence, Assumption 1 is equivalent to $p_{i}>0$ for $i=1,...,m$ and Assumption
2 is fulfilled exactly if it holds $\|a_{i}\|_{2}>0$ and $p_{i}>0$ for
$i=1,...,m$.

  2. (ii)

Block sketching. If
$S=\begin{pmatrix}\ldots&e_{r_{i}}&e_{r_{i}+1}&\ldots&e_{s_{i}}&\ldots\end{pmatrix}$
is sampled with probability $p_{i}$ such as in the case of block Kaczmarz,
Assumption 1 is equivalent to $p_{i}>0$ for all $i$ and Assumption 2 is
equivalent to the condition $\|A_{(i)}\|_{F}>0$ and $p_{i}>0$ for all $i$.

The proof of the next lemma is along the lines of the one of [24, Lemma 2.5].

######  Lemma 2.3.

Let Assumption 1 hold. Then, the matrix $M$ is positive definite.

###### Proof.

Since the matrix $L:=\mathbb{E}[SS^{T}]$ is positive definite, there exists
$\epsilon>0$ such that for all $x\in\mathbb{R}^{n}$ with $\|x\|_{2}=1$ it
holds that $x^{T}Lx>\epsilon$. For $k\in\mathbb{N}$ we set

$\Lambda_{k,k+1}:=\\{S:k\leq\|\mathbf{A}^{T}S\|_{2}^{2}<k+1\\}.$

By Levi’s theorem on monotone convergence, we have that

$\displaystyle x^{T}Lx=\int{x^{T}SS^{T}x}\
\mathrm{d}\mathcal{D}(S)=\sum_{k=0}^{\infty}\int_{\Lambda_{k,k+1}}x^{T}SS^{T}x\
\mathrm{d}\mathcal{D}(S).$

Since the left-hand side is finite and $x^{T}SS^{T}x\geq 0$ for all $S$, there
exists $k_{0}\in\mathbb{N}$ such that

$\displaystyle\sum_{k=k_{0}}^{\infty}\int_{\Lambda_{k,k+1}}x^{T}SS^{T}x\
\mathrm{d}\mathcal{D}(S)<\frac{\epsilon}{2},$

which means that

$\displaystyle\sum_{k=0}^{k_{0}-1}\int_{\Lambda_{k,k+1}}x^{T}SS^{T}x\
\mathrm{d}\mathcal{D}(S)=x^{T}Lx-\sum_{k=k_{0}}^{\infty}\int_{\Lambda_{k,k+1}}x^{T}SS^{T}x\
\mathrm{d}\mathcal{D}(S)>\frac{\epsilon}{2}.$

Hence, using Levi’s theorem again we conclude that

$\displaystyle x^{T}Mx$
$\displaystyle=\sum_{k=0}^{\infty}\int_{\Lambda_{k,k+1}}\frac{x^{T}SS^{T}x}{\|\mathbf{A}^{T}S\|_{2}^{2}}\
\mathrm{d}\mathcal{D}(S)$
$\displaystyle\geq\sum_{k=0}^{k_{0}-1}\int_{\Lambda_{k,k+1}}\frac{x^{T}SS^{T}x}{\|\mathbf{A}^{T}S\|_{2}^{2}}\
\mathrm{d}\mathcal{D}(S)$
$\displaystyle\geq\frac{1}{k_{0}}\sum_{k=0}^{k_{0}-1}\int_{\Lambda_{k,k+1}}x^{T}SS^{T}x\
\mathrm{d}\mathcal{D}(S)$ $\displaystyle>\frac{\epsilon}{2k_{0}}.$

∎

The following lemma will be also useful.

######  Lemma 2.4.

[24, Lemma 2.3] Assume that the linear system $Ax=b$ is consistent. Then for
any real-valued matrix $S$ with $m$ rows and any vector
$\tilde{x}\in\mathbb{R}^{n}$ it holds that
$\mathbf{A}^{T}SS^{T}(\mathbf{A}\tilde{x}-b)\neq 0$ if and only if
$S^{T}(\mathbf{A}\tilde{x}-b)\neq 0$.

###  2.3 Convex analysis

We recall some concepts and properties of convex functions [1].

Let $\varphi\colon\mathbb{R}^{n}\to\mathbb{R}$ be convex. Since $\varphi$ is
convex and finite everywhere, it is continuous and the set

$\displaystyle\partial\varphi(x)=\\{x^{*}\in\mathbb{R}^{n}:\varphi(y)\geq\varphi(x)+\langle
x^{*},y-x\rangle\ \text{for all }y\in\mathbb{R}^{n}\\},$

is nonempty, convex and compact for every $x\in\mathbb{R}^{n}$.

Throughout this article, we further assume that $\varphi$ is $\sigma$-strongly
convex for some $\sigma>0$, which means that for all $x,y\in\mathbb{R}^{n}$
and $x^{*}\in\partial\varphi(x)$ we have that

$\displaystyle\varphi(y)\geq\varphi(x)+\langle
x^{*},y-x\rangle+\frac{\sigma}{2}\|y-x\|_{2}^{2}.$

The _Bregman distance_ between $x$ and $y$ with respect to
$x^{*}\in\partial\varphi(x)$ is defined as

$\displaystyle D_{\varphi}^{x^{*}}(x,y)=\varphi(y)-\varphi(x)-\langle
x^{*},y-x\rangle.$ (7)

The _convex conjugate_ of $\varphi$ is defined as

$\displaystyle\varphi^{*}(x^{*})=\sup_{x\in\mathbb{R}^{n}}\langle
x^{*},x\rangle-\varphi(x),\qquad x^{*}\in\mathbb{R}^{n}.$

One can prove that the $\sigma$-strong convexity of $\varphi$ implies that the
function $\varphi^{*}$ is finite everywhere and differentiable with
$\sigma^{-1}$-Lipschitz continuous gradient. Moreover, for all
$x,x^{*}\in\mathbb{R}^{n}$ it holds that $x^{*}\in\partial\varphi(x)$ if and
only if $x=\nabla\varphi^{*}(x^{*})$. In this case, the _Fenchel equality_

$\displaystyle\varphi^{*}(x^{*})=\langle x^{*},x\rangle-\varphi(x)$

holds and the definition of the Bregman distance (7) can be rewritten as

$\displaystyle D_{\varphi}^{x^{*}}(x,y)=\varphi^{*}(x^{*})-\langle
x^{*},y\rangle+\varphi(y).$ (8)

Finally, we can lower bound the Bregman distance between
$x,y\in\mathbb{R}^{n}$ for all $x^{*}\in\partial\varphi(x)$ by

$\displaystyle D_{\varphi}^{x^{*}}(x,y)\geq\frac{\sigma}{2}\|x-y\|_{2}^{2}.$
(9)

##  3 The sketched Bregman-Kaczmarz method

In this section, we study convergence of the sketched Bregman-Kaczmarz method
for solving problem (1) given by

$\displaystyle\min\varphi(x)\qquad\text{s.t.
}\mathbf{A}x=b,\qquad\mathbf{A}\in\mathbb{R}^{m\times n},\
b\in\mathbb{R}^{m}.$

We suppose that the system $\mathbf{A}x=b$ is consistent, i.e. has a solution,
and $b\neq 0$. By strong convexity of $\varphi$, there is exactly one solution
$\hat{x}$ to problem (1). For a distribution $\mathcal{D}$ on the set of real-
valued matrices with $m$ rows, we consider the stochastic reformulation (2)
given by

$\displaystyle\min f(x):=\mathbb{E}[f_{S}(x)],\qquad
f_{S}(x)=\frac{1}{2}\|S^{T}(\mathbf{A}x-b)\|_{2}^{2}.$

The stochastic reformulation is exact under Assumption 1 due to Lemma 2.1. The
gradient of $f_{S}$ can be easily calculated as

$\displaystyle\nabla f_{S}(x)=\mathbf{A}^{T}SS^{T}(\mathbf{A}x-b)$

and the sketched Bregman-Kaczmarz method is given by Algorithm 1.

1:Input: $x_{0}^{*}=0\in\mathbb{R}^{n}$, step sizes $\alpha_{k}>0$ and a
distribution $\mathcal{D}$ on the set of matrices with $m$ rows

2:Initialization: $x_{0}=\nabla\varphi^{*}(x_{0}^{*})$

3:for $k=0,1,...$ do

4: sample a random sketching matrix $S_{k}\sim{\mathcal{D}}$ (independent of
$S_{0},...,S_{k-1}$)

5: $x_{k+1}^{*}=x_{k}^{*}-\alpha_{k}\nabla
f_{S_{k}}(x_{k})=x_{k}^{*}-\alpha_{k}\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)$

6: $x_{k+1}=\nabla\varphi^{*}(x_{k+1}^{*})$

Algorithm 1 Basic method: Randomized Bregman-Kaczmarz method (BK) with
sketching matrices

Note that, if $x_{0}^{*}$ is chosen from $\mathcal{R}(\mathbf{A}^{T})$, by
induction it holds $x_{k}^{*}\in\mathcal{R}(\mathbf{A}^{T})$ at every
iteration $k$. Also, by subdifferential inversion we have
$x_{k}^{*}\in\partial\varphi(x_{k})$. As in [13] we refer to $x_{k}$,
$x_{k}^{*}$ as the primal variable and the dual variable, respectively.

We now prove convergence of the $x_{k}$ in Algorithm 1 to the solution
$\hat{x}$ for two different step size sequences.

To this end, we need to assume an error bound which relates the least-squares
residual $\|\mathbf{A}x-b\|_{2}^{2}$ to the Bregman distance to the exact
solution $\hat{x}$ for iterates of Algorithm 1.

######  Assumption 3.

Let $\hat{x}$ be the solution of problem (1). There exists $\gamma>0$ such
that for all $x\in\mathbb{R}^{n}$ with
$\partial\varphi(x)\cap\mathcal{R}(\mathbf{A}^{T})\neq\emptyset$ and all
$x^{*}\in\partial\varphi(x)\cap\mathcal{R}(\mathbf{A}^{T})$ we have

$\displaystyle
D_{\varphi}^{x^{*}}(x,\hat{x})\leq\gamma\cdot\|\mathbf{A}x-b\|_{2}^{2}.$

For the function $\varphi(x)=\lambda\|x\|_{1}+\frac{1}{2}\|x\|_{2}^{2}$, it
holds that $\mathrm{dom}\ \varphi=\mathbb{R}^{n}$ and Assumption 3 holds with
an explicit constant $\gamma$ depending on $A$ and $\hat{x}$, see [19, Lemma
3.1]. See also [20, Theorem 3.9] for a sufficient condition for Assumption 3
to hold. The next lemma will be used to prove convergence of Algorithm 1 for
suitable step sizes.

######  Lemma 3.1.

Let Assumptions 1\- 2 be fulfilled. Then, the iterates $x_{k}$, $x_{k}^{*}$ of
Algorithm 1 fulfill that

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})-\alpha_{k}\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}+\frac{\alpha_{k}^{2}}{2\sigma}\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}.$

###### Proof.

Since $\varphi$ is $\sigma$-strongly convex, $\varphi^{*}$ is $\sigma^{-1}$
smooth. Hence, using (8) and the descent lemma [1, Theorem 18.15(iii)] give
that

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})$
$\displaystyle=\varphi^{*}(x_{k+1}^{*})-\langle
x_{k+1}^{*},\hat{x}\rangle+\varphi(\hat{x})$
$\displaystyle\leq\varphi^{*}(x_{k}^{*})+\langle\nabla\varphi^{*}(x_{k}^{*}),x_{k+1}^{*}-x_{k}^{*}\rangle+\frac{1}{2\sigma}\|x_{k+1}^{*}-x_{k}^{*}\|_{2}^{2}-\langle
x_{k+1}^{*},\hat{x}\rangle+\varphi(\hat{x}).$

By inserting the update and using that $\nabla\varphi(x_{k}^{*})=x_{k}$, we
obtain

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})$
$\displaystyle\leq\varphi^{*}(x_{k}^{*})+\langle
x_{k},-\alpha_{k}\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\rangle+\frac{\alpha_{k}^{2}}{2\sigma}\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}$
$\displaystyle\hskip 14.22636pt-\langle
x_{k}^{*},\hat{x}\rangle+\alpha_{k}\langle\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b),\hat{x}\rangle+\varphi(\hat{x})$
$\displaystyle=D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})-\alpha_{k}\langle
x_{k}-\hat{x},\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\rangle+\frac{\alpha_{k}^{2}}{2\sigma}\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}$
$\displaystyle=D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})-\alpha_{k}\langle
S_{k}^{T}\mathbf{A}(x_{k}-\hat{x}),S_{k}^{T}(\mathbf{A}x_{k}-b)\rangle+\frac{\alpha_{k}^{2}}{2\sigma}\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}$

and the assertion follows by using that $\mathbf{A}\hat{x}=b$. ∎

######  Theorem 3.2.

Let Assumptions 1\- 3 be fulfilled and let $x_{k}$, $x_{k}^{*}$ be the
iterates of Algorithm 1. If at each iteration the step sizes are chose by

$\displaystyle\alpha_{k}=\begin{cases}\frac{\sigma}{\|\mathbf{A}^{T}S_{k}\|_{2}^{2}},&\mathbf{A}^{T}S_{k}\neq
0,\\\ 0,&\text{otherwise},\end{cases}$ (10)

or

$\displaystyle\alpha_{k}=\begin{cases}\frac{\sigma\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}},&S_{k}^{T}(\mathbf{A}x_{k}-b)\neq
0,\\\ 0,&\text{otherwise},\end{cases}$ (11)

is chosen, it holds that

$\displaystyle\mathbb{E}\big{[}D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})\big{]}\leq\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}\mathbb{E}\big{[}D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})\big{]}.$
(12)

In particular, the primal iterates $x_{k}$ converge to $\hat{x}$ in
$\mathcal{L}_{2}$-sense with the linear rate

$\displaystyle\mathbb{E}[\|x_{k}-\hat{x}\|_{2}^{2}]\leq\frac{2\varphi(\hat{x})}{\sigma}\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}^{k}.$
(13)

###### Proof.

We insert the corresponding step sizes into Lemma 3.1. For step size (10),
again using Lemma 3.1, this yields

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})$ $\displaystyle\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})-\frac{\sigma\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\mathbf{A}^{T}S_{k}\|_{2}^{2}}+\frac{\sigma\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{2\|\mathbf{A}^{T}S_{k}\|_{2}^{4}}$
$\displaystyle\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})-\frac{\sigma\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{2\|\mathbf{A}^{T}S_{k}\|_{2}^{2}}.$

For step size (11), we estimate

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})$ $\displaystyle\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})-\frac{\sigma\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{4}}{2\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}$
$\displaystyle\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})-\frac{\sigma\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{2\|\mathbf{A}^{T}S_{k}\|_{2}^{2}}.$
(14)

Since the $S_{k}$ are sampled independently, we can lower bound

$\displaystyle\mathbb{E}\Big{[}\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\mathbf{A}^{T}S_{k}\|_{2}^{2}}\mid
S_{0},...S_{k-1}\Big{]}$
$\displaystyle=\mathbb{E}\Big{[}\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\mathbf{A}^{T}S_{k}\|_{2}^{2}}\mid
S_{0},...S_{k-1}\Big{]}$
$\displaystyle=\bigl{\langle}\mathbf{A}x_{k}-b,\mathbb{E}\Big{[}\frac{S_{k}S_{k}^{T}}{\|\mathbf{A}^{T}S_{k}\|_{2}^{2}}\Big{]}(\mathbf{A}x_{k}-b)\bigr{\rangle}$
$\displaystyle=\langle\mathbf{A}x_{k}-b,M(\mathbf{A}x_{k}-b)\rangle$
$\displaystyle\geq\lambda_{\min}(M)\|\mathbf{A}x_{k}-b\|_{2}^{2}$
$\displaystyle\geq\frac{\lambda_{\min}(M)}{\gamma}D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x}),$

where the last step is due to Assumption 3. By the law of total expectation,
we conclude

$\displaystyle\mathbb{E}\big{[}D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})\big{]}\leq\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}\mathbb{E}\big{[}D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})\big{]}$

and hence inductively,

$\displaystyle\mathbb{E}[D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})]\leq\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}^{k}\
D_{\varphi}^{x_{0}^{*}}(x_{0},\hat{x})=\varphi(\hat{x})\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}^{k},$

where the last equality is due to (8) and the fact that $x_{0}^{*}=0$. Finally
by (9), strong convexity of $\varphi$ implies (13). ∎

Note that step size (11) is always defined by Lemma 2.4. The first step size
(10) generalizes the non-adaptive step size for block Kaczmarz [20], the
second one (11) the adaptive step size studied by [24].

######  Example 3.3.

  1. (i)

Single row sketching. In case that the sketching matrices are chosen as unit
vectors $S_{k}=e_{i_{k}}$, Algorithm 1 coincides with the Bregman-Kaczmarz
method with single row sketching [13], and both step sizes (10) and (11) equal
$\alpha_{k}=\frac{\sigma}{\|a_{i_{k}}\|_{2}^{2}}$. Indeed, in this case we
have that

$\displaystyle\nabla f_{S_{k}}(x_{k})$
$\displaystyle=\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)=(\langle
a_{i_{k}},x_{k}\rangle-b_{i_{k}})a_{i_{k}},$ $\displaystyle
S_{k}^{T}\mathbf{A}x_{k}$ $\displaystyle=\langle a_{i_{k}},x_{k}\rangle,$
$\displaystyle S_{k}^{T}b$ $\displaystyle=b_{i_{k}}.$ (15)

If the row $a_{i}$ is sampled with the probability
$\|a_{i}\|_{2}^{2}/\|A\|_{F}^{2}$ suggested by Strohmer and Vershynin, we have
that $\lambda_{\min}(M)=\frac{1}{\|A\|_{F}^{2}}$ and hence, Theorem 3.2
recovers the convergence rate from [11] for
$\varphi(x)=\lambda\|x\|_{1}+\frac{1}{2}\|x\|_{2}^{2}$.

  2. (ii)

Block sketching. Similarly, if
$S_{k}=\begin{pmatrix}\ldots&e_{r_{i_{k}}}&e_{r_{i_{k}}+1}&\ldots&e_{s_{i_{k}}}&\ldots\end{pmatrix}$,
we calculate that

$\displaystyle\nabla f_{S_{k}}(x_{k})$
$\displaystyle=A_{(i_{k})}^{T}(A_{(i_{k})}x_{k}-b_{(i_{k})}),$ $\displaystyle
S_{k}^{T}\mathbf{A}x_{k}$ $\displaystyle=A_{(i_{k})}x_{k}-b_{(i_{k})},$
$\displaystyle S_{k}^{T}b$ $\displaystyle=b_{(i_{k})}.$ (16)

If $S_{k}$ is sampled with probability $\|A_{(i)}\|_{F}^{2}/\|A\|_{F}^{2}$, we
have that $\lambda_{\min}(M)=\frac{1}{\|A\|_{F}^{2}}$ and the contraction
factor in the convergence rate (13) is

$\displaystyle
1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}=1-\frac{\sigma}{2\gamma\|A\|_{F}^{2}},$

which to the best of our knowledge is novel in the literature.

##  4 Adaptive Heavy Ball acceleration

We consider the parameterized update

$\displaystyle x_{k+1}^{*}(\alpha,\beta)$
$\displaystyle=x_{k}^{*}-\alpha\nabla
f_{S_{k}}(x_{k})+\beta(x_{k}^{*}-x_{k-1}^{*}),$ $\displaystyle
x_{k+1}(\alpha,\beta)$ $\displaystyle=\nabla\varphi^{*}(x_{k+1}^{*}),$ (17)

that is, we incorporate heavy ball acceleration into the dual update.

In this section, we want to derive expressions for $\alpha=\alpha_{k}$ and
$\beta=\beta_{k}$ which ensure convergence of the update. For notational
purposes, we set

$\displaystyle y_{k}^{*}=x_{k}^{*}-\alpha_{k}\nabla f_{S_{k}}(x_{k}),\qquad
y_{k}=\nabla\varphi^{*}(y_{k}^{*}),\qquad d_{k}^{*}=x_{k}^{*}-x_{k-1}^{*}.$
(18)

By (8), the Bregman distance of the primal variable to the solution $\hat{x}$
of problem (1) can be calculated as

$\displaystyle
D_{\varphi}^{x_{k+1}^{*}(\alpha,\beta)}(x_{k+1}(\alpha,\beta),\hat{x})$
$\displaystyle=\varphi^{*}(x_{k}^{*}-\alpha\nabla f_{S_{k}}(x_{k})+\beta
d_{k}^{*})-\langle x_{k}^{*}-\alpha\nabla f_{S_{k}}(x_{k})+\beta
d_{k}^{*},\hat{x}\rangle+\varphi(\hat{x})$
$\displaystyle=\varphi^{*}(x_{k}^{*}-\alpha\nabla f_{S_{k}}(x_{k})+\beta
d_{k}^{*})-\langle x_{k}^{*},\hat{x}\rangle+\alpha\langle
S_{k}^{T}(\mathbf{A}x_{k}-b),S_{k}^{T}b\rangle$ $\displaystyle\hskip
14.22636pt+\beta\langle d_{k}^{*},\hat{x}\rangle+\varphi(\hat{x}).$ (19)

Hence, minimizing (19) over $\alpha$ and $\beta$ would give us momentum with
minimal errors measured in Bregman distance, i.e. our update would fulfill the
property

$\displaystyle(\alpha,\beta)\in\operatorname*{\arg\\!\min}_{\alpha,\beta\in\mathbb{R}}D_{\varphi}^{x_{k+1}^{*}(\alpha,\beta)}\big{(}x_{k+1}(\alpha,\beta),\hat{x}\big{)}.$

However, there are two problems occurring here. First, problem (19) depends on
the exact solution $\hat{x}$, which is not known. Second, $\varphi^{*}$ is an
arbitrary convex function which Lipschitz continuous gradient and hence, we
expect minimization over two variables to be computationally expensive. For
instance, the function $\varphi(x)=\lambda\|x\|_{1}+\frac{1}{2}\|x\|_{2}^{2}$
has the conjugate
$\varphi^{*}(x)=\frac{1}{2}\|\mathcal{S}_{\lambda}(x)\|_{2}^{2}$, which does
not admit a closed form solution for minimizing (19) over $\alpha$ and
$\beta$. We will resolve the first problem of the $\hat{x}$-dependence
completely by introducing a new scalar variable $s_{k}$, which is updated
simultaneously with $x_{k}$, $x_{k}^{*}$ and in each iteration fulfills that
$s_{k}=\langle x_{k}^{*}-x_{k-1}^{*},\hat{x}\rangle$. In Section 4.1, as a
first approach to address the second problem, we optimize only over $\beta$
for a fixed step size $\alpha_{k}$ chosen as to Algorithm 1. For the function
$\varphi(x)=\lambda\|x\|_{1}+\frac{1}{2}\|x\|_{2}^{2}$, this can be done with
reasonable effort by a sorting-based algorithm, exploiting the monotonicity of
the partial derivative $\tfrac{\partial\varphi^{*}}{\partial\beta}$ [13].

Therefore, as a second approach, in Section 4.2 we optimize over an upper
bound of (19), obtained by the descent lemma, which is quadratic in $\alpha$
and $\beta$ and hence allows for easy joint minimization in both variables.

###  4.1 Minimum error steps

In this section, we assume that the step size sequence $\alpha_{k}$ is
fixed,111Here, $\alpha_{k}$ can be stochastic and adaptive, that is, dependent
on expressions which are known in iteration $k$. We will prove convergence of
the method in the case that $\alpha_{k}$ is chosen according to (10) or (11),
but the method can be formulated for any sequence $\alpha_{k}$ which can be
implemented. and for a momentum parameter $\beta>0$ we set

$\displaystyle x_{k+1}^{*}(\beta)$ $\displaystyle=x_{k}^{*}-\alpha_{k}\nabla
f_{S_{i_{k}}}(x_{k})+\beta(x_{k}^{*}-x_{k-1}^{*}),$ $\displaystyle
x_{k+1}(\beta)$ $\displaystyle=\nabla\varphi^{*}(x_{k+1}^{*}(\beta)).$

Here, we want to pursue the idea of choosing $\beta$ with a minimal error,
that is, such that the primal update $x_{k+1}$ is closest to the solution
$\hat{x}$ in Bregman distance. In formulas, we seek for $\beta=\beta_{k}$ with

$\displaystyle\beta_{k}\in\operatorname*{\arg\\!\min}_{\beta\in\mathbb{R}}D_{\varphi}^{x_{k+1}^{*}(\beta)}\big{(}x_{k+1}(\beta),\hat{x}\big{)}.$
(20)

To this end, using (8), we rewrite the right-hand side (20) as

$\displaystyle
D_{\varphi}^{x_{k+1}^{*}(\beta)}\big{(}x_{k+1}(\beta),\hat{x}\big{)}$
$\displaystyle=\varphi^{*}(x_{k+1}^{*}(\beta))-\langle
x_{k+1}^{*}(\beta),\hat{x}\rangle+\varphi(\hat{x})$
$\displaystyle=\varphi^{*}(y_{k}^{*}+\beta d_{k}^{*})-\langle
y_{k}^{*},\hat{x}\rangle-\beta\langle
d_{k}^{*},\hat{x}\rangle+\varphi(\hat{x}),$

where $y_{k}^{*}$ and $d_{k}^{*}$ are defined as in (18). Consequently, the
parameter $\beta$ is given as the solution to the nonsmooth convex
optimization problem

$\displaystyle\beta_{k}\in\operatorname*{\arg\\!\min}_{\beta\in\mathbb{R}}\varphi^{*}(y_{k}^{*}+\beta
d_{k}^{*})-\beta\langle d_{k}^{*},\hat{x}\rangle.$ (21)

At this point, problem (21) depends on the solution $\hat{x}$. However, for

$\displaystyle s_{k}:=\langle d_{k}^{*},\hat{x}\rangle,$ (22)

we have $s_{0}=0$, if we initialize $x_{-1}^{*}=x_{0}^{*}$, and we have the
recursion

$\displaystyle s_{k+1}$ $\displaystyle=\langle
x_{k+1}^{*}-x_{k}^{*},\hat{x}\rangle$
$\displaystyle=\langle-\alpha_{k}\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b),\hat{x}\rangle+\langle\beta_{k}d_{k}^{*},\hat{x}\rangle$
$\displaystyle=-\alpha_{k}\langle
S_{k}^{T}(\mathbf{A}x_{k}-b),S_{k}^{T}b\rangle+\beta_{k}s_{k}.$ (23)

That is, we can set up a method which updates $x_{k}^{*},x_{k}$ and $s_{k}$
with the update (23) and determine $\beta_{k}$ by

$\displaystyle\beta_{k}\in\operatorname*{\arg\\!\min}_{\beta\in\mathbb{R}}\varphi^{*}(y_{k}^{*}+\beta
d_{k}^{*})-\beta s_{k}.$ (24)

This is summarized in Algorithm 2.

1:Input: $x_{0}^{*}=0\in\mathbb{R}^{n}$ and a distribution $\mathcal{D}$ on
the set of matrices with $m$ rows

2:Initialization: $x_{-1}^{*}=x_{0}^{*}$, $x_{0}=\nabla\varphi^{*}(x_{0}^{*})$
and $s_{0}=0$

3:for $k=0,1,...$ do

4: sample a random sketching matrix $S_{k}\sim{\mathcal{D}}$ (independent of
$S_{0},...,S_{k-1}$)

5: choose a step size $\alpha_{k}>0$

6: $y_{k}^{*}=x_{k}^{*}-\alpha_{k}\nabla
f_{S_{k}}(x_{k})=x_{k}^{*}-\alpha_{k}\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)$

7: $d_{k}^{*}=x_{k}^{*}-x_{k-1}^{*}$

8: if $d_{k}^{*}\neq 0$ then

9:
$\beta_{k}\in\operatorname*{\arg\\!\min}_{\beta\in\mathbb{R}}\varphi^{*}(y_{k}^{*}+\beta
d_{k}^{*})-\beta s_{k}$

10: else

11: $\beta_{k}=0$

12: $x_{k+1}^{*}=y_{k}^{*}+\beta_{k}d_{k}^{*}$

13: $x_{k+1}=\nabla\varphi^{*}(x_{k+1}^{*})$

14: $s_{k+1}=-\alpha_{k}\langle
S_{k}^{T}(\mathbf{A}x_{k}-b),S_{k}^{T}b\rangle+\beta_{k}s_{k}.$

Algorithm 2 Randomized Bregman-Kaczmarz method with exact minimal error
momentum (BK-EM)

######  Remark 4.1.

By the derivation of Algorithm 2, in iteration $k$ it holds

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})\leq
D_{\varphi}^{x_{k+1}^{*}(0)}(x_{k+1}(0),\hat{x}),$

where $x_{k+1}^{*}(0)$, $x_{k+1}(0)$ are the updates which Algorithm 1 would
compute. This means that each step of Algorithm 2 makes at least as much
progress as Algorithm 1 towards the solution $\hat{x}$ measured in the Bregman
distance induced by $\varphi$. We can not conclude at this point that
Algorithm 2 converges faster than Algorithm 1. Nonetheless, in Theorem 4.3 we
will give an upper bound for the Bregman distances towards the solution which
is as least as good as the one for Algorithm 1 from Theorem 3.2.

######  Example 4.2 (Algorithm 2 for single row sketching).

We revisit the case of the Bregman-Kaczmarz method with single row sketching
(Example 3.3), that is, with $S_{k}=e_{i_{k}}$ and
$\alpha_{k}=\frac{\sigma}{\|a_{i_{k}}\|_{2}^{2}}$. Then in view of the
equations from (15), the steps of Algorithm 2 are given as follows:

$\displaystyle d_{k}^{*}$ $\displaystyle=x_{k}^{*}-x_{k-1}^{*}$ $\displaystyle
y_{k}^{*}$ $\displaystyle=x_{k}^{*}-\frac{\sigma(\langle
a_{i_{k}},x_{k}\rangle-b_{i_{k}})}{\|a_{i_{k}}\|_{2}^{2}}a_{i_{k}}$
$\displaystyle\beta_{k}^{*}$ $\displaystyle\
\begin{cases}\in\operatorname*{\arg\\!\min}_{\beta\in\mathbb{R}}\varphi^{*}(y_{k}^{*}+\beta
d_{k}^{*})-\beta s_{k},&d_{k}^{*}\neq 0,\\\ =0,&\text{otherwise}\end{cases}$
$\displaystyle x_{k+1}^{*}$ $\displaystyle=y_{k}^{*}+\beta_{k}d_{k}^{*},$
$\displaystyle x_{k+1}$ $\displaystyle=\nabla\varphi^{*}(x_{k+1}^{*}),$
$\displaystyle s_{k+1}$ $\displaystyle=-\frac{\sigma(\langle
a_{i_{k}},x_{k}\rangle-
b_{i_{k}})}{\|a_{i_{k}}\|_{2}^{2}}b_{i_{k}}+\beta_{k}s_{k}.$

We now prove convergence of Algorithm 1.

######  Theorem 4.3.

Let Assumptions 1\- 3 be fulfilled. Let $x_{k}$, $x_{k}^{*}$ be the iterates
of Algorithm 2 and $d_{k}^{*}$, $y_{k}^{*}$ and $y_{k}$ defined according to
(18). Then it holds that

$\displaystyle\mathbb{E}\big{[}D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})\big{]}\leq\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}\mathbb{E}\big{[}D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})\big{]}-\frac{\sigma}{2}\mathbb{E}\big{[}\frac{\langle
y_{k}-\hat{x},d_{k}^{*}\rangle^{2}}{\|d_{k}^{*}\|_{2}^{2}}\big{]}.$ (25)

In particular, if at each iteration $\alpha_{k}$ is chosen according to either
(10) or (11), $x_{k}$ converges in expectation to $\hat{x}$ in
$\mathcal{L}_{2}$-sense with the linear rate (13), that is,

$\displaystyle\mathbb{E}[\|x_{k}-\hat{x}\|_{2}^{2}]\leq\frac{2\varphi(\hat{x})}{\sigma}\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}^{k}.$

###### Proof.

For all $\beta\in\mathbb{R}$, by (8) and the descent lemma [1, Theorem
18.15(iii)], we have

$\displaystyle D_{\varphi}^{x_{k+1}^{*}(\beta)}(x_{k+1}(\beta),\hat{x})$
$\displaystyle=\varphi^{*}(y_{k}^{*}+\beta d_{k}^{*})-\langle y_{k}^{*}+\beta
d_{k}^{*},\hat{x}\rangle+\varphi(\hat{x})$
$\displaystyle\leq\varphi^{*}(y_{k}*)+\beta\langle
y_{k},d_{k}^{*}\rangle+\frac{\beta^{2}}{2\sigma}\|d_{k}^{*}\|_{2}^{2}-\langle
y_{k}^{*}+\beta d_{k}^{*},,\hat{x}\rangle+\varphi(\hat{x})$
$\displaystyle=D_{\varphi}^{y_{k}^{*}}(y_{k},\hat{x})+\beta\langle
y_{k},d_{k}^{*}\rangle+\frac{\beta^{2}}{2\sigma}\|d_{k}^{*}\|_{2}^{2}-\beta\langle
d_{k}^{*},\hat{x}\rangle$
$\displaystyle=D_{\varphi}^{y_{k}^{*}}(y_{k},\hat{x})+\beta\langle
y_{k}-\hat{x},d_{k}^{*}\rangle+\frac{\beta^{2}}{2\sigma}\|d_{k}^{*}\|_{2}^{2}.$

Minimizing this upper bound over $\beta$ gives the minimizer

$\displaystyle\hat{\beta}=\frac{\sigma\langle\hat{x}-y_{k},d_{k}^{*}\rangle}{\|d_{k}^{*}\|_{2}^{2}}.$

Hence, using the minimizing property (20) of $\beta_{k}$ we conclude that

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})$
$\displaystyle=D_{\varphi}^{x_{k+1}^{*}(\beta_{k})}(x_{k+1}(\beta_{k}),\hat{x})$
$\displaystyle\leq
D_{\varphi}^{x_{k+1}^{*}(\hat{\beta})}(x_{k+1}(\hat{\beta}),\hat{x})$
$\displaystyle=D_{\varphi}^{y_{k}^{*}}(y_{k},\hat{x})-\frac{\sigma\langle
y_{k}-\hat{x},d_{k}^{*}\rangle^{2}}{2\|d_{k}^{*}\|_{2}^{2}}.$

Taking expectation gives that

$\displaystyle\mathbb{E}\big{[}D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})\big{]}\leq\mathbb{E}\big{[}D_{\varphi}^{y_{k}^{*}}(y_{k},\hat{x})\big{]}-\frac{\sigma}{2}\mathbb{E}\big{[}\frac{\langle
y_{k}-\hat{x},d_{k}^{*}\rangle^{2}}{2\|d_{k}^{*}\|_{2}^{2}}\big{]}.$

Finally, we note that $y_{k}$, $y_{k}^{*}$ equal $x_{k+1}$, $x_{k+1}^{*}$ from
Algorithm 1, given that both updates start from $x_{k}$, $x_{k}^{*}$, and
hence the assertion follows by Theorem 3.2. ∎

We see that the convergence estimate (25) for the primal iterates of Algorithm
2 is at least as good as the estimate (12) for the primal iterates of
Algorithm 1. For interpretation of the rates in case of single row sketching
and block sketching, we refer to Example 3.3.

###  4.2 Relaxed minimum error steps

Algorithm 2 suffers from the fact that determining $\beta_{k}$ as the solution
of the problem (24) is costly. Using the descent lemma [1, Theorem 18.15(iii)]
again, we can estimate (19) from above by

$\displaystyle
D_{\varphi}^{x_{k+1}^{*}(\alpha,\beta)}(x_{k+1}(\alpha,\beta),\hat{x})$
$\displaystyle\leq\varphi^{*}(x_{k}^{*})+\langle x_{k},-\alpha\nabla
f_{S_{k}}(x_{k})+\beta d_{k}^{*}\rangle+\frac{1}{2\sigma}\|-\alpha\nabla
f_{S_{k}}(x_{k})+\beta d_{k}^{*}\|_{2}^{2}$ $\displaystyle\hskip
14.22636pt-\langle x_{k}^{*},\hat{x}\rangle+\alpha\langle\nabla
f_{S_{k}}(x_{k}),\hat{x}\rangle-\beta\langle
d_{k}^{*},\hat{x}\rangle+\varphi(\hat{x})$
$\displaystyle=D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})+\langle
x_{k}-\hat{x},-\alpha\nabla f_{S_{k}}(x_{k})+\beta
d_{k}^{*}\rangle+\frac{1}{2\sigma}\|-\alpha\nabla f_{S_{k}}(x_{k})+\beta
d_{k}^{*}\|_{2}^{2}.$ (26)

We note that (26) is quadratic in $\alpha$ and $\beta$ and therefore easy to
minimize.

Taking derivative of (26) with respect to $\alpha$ yields

$\displaystyle 0$ $\displaystyle=-\langle x_{k}-\hat{x},\nabla
f_{S_{k}}(x_{k})\rangle+\frac{1}{\sigma}\langle-\alpha_{k}\nabla
f_{S_{k}}(x_{k})+\beta_{k}d_{k}^{*},-\nabla f_{S_{k}}(x_{k})\rangle$
$\displaystyle=-\langle x_{k}-\hat{x},\nabla
f_{S_{k}}(x_{k})\rangle+\frac{1}{\sigma}\big{(}\alpha_{k}\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}-\beta_{k}\langle d_{k}^{*},\nabla
f_{S_{k}}(x_{k})\rangle\big{)}$

and by rearranging we obtain the equation

$\displaystyle\alpha_{k}\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}-\beta_{k}\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle=\sigma\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}.$
(27)

Taking derivative of (26) with respect to $\beta$ yields

$\displaystyle 0$ $\displaystyle=\langle
x_{k}-\hat{x},d_{k}^{*}\rangle+\frac{1}{\sigma}\langle-\alpha_{k}\nabla
f_{S_{k}}(x_{k})+\beta_{k}d_{k}^{*},d_{k}^{*}\rangle$ $\displaystyle=\langle
x_{k}-\hat{x},d_{k}^{*}\rangle+\frac{1}{\sigma}\big{(}-\alpha_{k}\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle+\beta_{k}\|d_{k}^{*}\|_{2}^{2}\big{)},$

which gives the second equation

$\displaystyle-\alpha_{k}\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle+\beta_{k}\|d_{k}^{*}\|_{2}^{2}=\sigma\langle\hat{x}-x_{k},d_{k}^{*}\rangle.$
(28)

With the same derivation as for Algorithm 3, we can again replace
$\langle\hat{x},d_{k}^{*}\rangle$ by $s_{k}$, where we update $s_{k}$
according to (23).

We can now solve the system of linear equations (27)-(28) for $\alpha$ and
$\beta$. The system is invertible, if

$\displaystyle\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}>0,$

which is equivalent to the condition that the search directions $\nabla
f_{S_{k}}(x_{k})$ and $d_{k}^{*}$ are linearly independent. In case of linear
dependence, we can set $\beta_{k}=0$ and using (27), we obtain back
$\alpha_{k}$ from (11). In case of linear independence, the optimal solution
$(\alpha_{k},\beta_{k})$ is given by

$\displaystyle\alpha_{k}$
$\displaystyle=\sigma\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}+\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle(s_{k}-\langle
x_{k},d_{k}^{*}\rangle)}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}}$ $\displaystyle\beta_{k}$
$\displaystyle=\sigma\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle+\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}(s_{k}-\langle x_{k},d_{k}^{*}\rangle)}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}}.$

We collect the steps in Algorithm 3.

1:Input: $x_{0}^{*}=0\in\mathbb{R}^{n}$ and a distribution $\mathcal{D}$ on
the set of matrices with $m$ rows

2:Initialization: $x_{-1}^{*}=x_{0}^{*}$, $x_{0}=\nabla\varphi^{*}(x_{0}^{*})$
and $s_{0}=0$

3:sample a random sketching matrix $S_{k}\sim{\mathcal{D}}$ (independent of
$S_{0},...,S_{k-1}$)

4:for $k=0,1,...$ do

5: $d_{k}^{*}=x_{k}^{*}-x_{k-1}^{*}$

6: if $\|\nabla f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}>0$ then

7:
$\alpha_{k}=\sigma\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}+\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle(s_{k}-\langle
x_{k},d_{k}^{*}\rangle)}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}}$

8:
$\beta_{k}=\sigma\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle+\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}(s_{k}-\langle x_{k},d_{k}^{*}\rangle)}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}}$

9: else

10: $\alpha_{k}=\sigma\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}}$

11: $\beta_{k}=0$

12: $x_{k+1}^{*}=x_{k}^{*}-\alpha_{k}\nabla
f_{S_{k}}(x_{k})+\beta_{k}d_{k}^{*}$

13: $x_{k+1}=\nabla\varphi^{*}(x_{k+1}^{*})$

14: $s_{k+1}=-\alpha_{k}\langle
S_{k}^{T}(\mathbf{A}x_{k}-b),S_{k}^{T}b\rangle+\beta_{k}s_{k}.$

Algorithm 3 Randomized Bregman-Kaczmarz method with relaxed minimal error
momentum (BK-REM)

######  Example 4.4 (Algorithm 3 in case of single row sketching).

We revisit again the case of the Bregman-Kaczmarz method with single row
sketching (Example 3.3), that is, with $S_{k}=e_{i_{k}}$ and
$\alpha_{k}=\frac{\sigma}{\|a_{i_{k}}\|_{2}^{2}}$. In view of the equations
from (15), the steps of Algorithm 3 are given as follows:

$\displaystyle d_{k}^{*}$ $\displaystyle=x_{k}^{*}-x_{k-1}^{*},$
$\displaystyle t_{k}$ $\displaystyle=\begin{cases}\sigma\frac{(\langle
a_{i_{k}},x_{k}\rangle-b_{i_{k}})\|d_{k}^{*}\|_{2}^{2}+\langle
a_{i_{k}},d_{k}^{*}\rangle(s_{k}-\langle
x_{k},d_{k}^{*}\rangle)}{\|a_{i_{k}}\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle
a_{i_{k}},d_{k}^{*}\rangle^{2}},&\|a_{i_{k}}\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle
a_{i_{k}},d_{k}^{*}\rangle^{2}>0,\\\ \sigma\frac{(\langle
a_{i_{k}},x_{k}\rangle-
b_{i_{k}})}{\|a_{i_{k}}\|_{2}^{2}},&\text{otherwise},\end{cases}$
$\displaystyle\beta_{k}$ $\displaystyle=\begin{cases}\sigma\frac{(\langle
a_{i_{k}},x_{k}\rangle-b_{i_{k}})\langle
a_{i_{k}},d_{k}^{*}\rangle+\|a_{i_{k}}\|_{2}^{2}(s_{k}-\langle
x_{k},d_{k}^{*}\rangle)}{\|a_{i_{k}}\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle
a_{i_{k}},d_{k}^{*}\rangle^{2}},&\|a_{i_{k}}\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle
a_{i_{k}},d_{k}^{*}\rangle^{2}>0,\\\ \sigma\frac{\langle
a_{i_{k}},x_{k}\rangle-
b_{i_{k}}}{\|a_{i_{k}}\|_{2}^{2}},&\text{otherwise},\end{cases}$
$\displaystyle x_{k+1}^{*}$
$\displaystyle=x_{k}^{*}-t_{k}a_{i_{k}}+\beta_{k}d_{k}^{*},$ $\displaystyle
x_{k+1}$ $\displaystyle=\nabla\varphi^{*}(x_{k+1}^{*}),$ $\displaystyle
s_{k+1}$ $\displaystyle=-b_{i_{k}}t_{k}+\beta_{k}s_{k}.$

Note that we have $\alpha_{k}(\langle a_{i_{k}},x_{k}\rangle-b_{i_{k}})=t_{k}$
and hence $\alpha_{k}\nabla f_{S_{k}}(x_{k})=t_{k}a_{i_{k}}.$

We now investigate convergence of Algorithm 3. The following lemma will be
useful.

######  Lemma 4.5.

Let $x,y,z\in\mathbb{R}^{n}$ such that $(y,z)$ is linearly independent. Then
it holds that

$\displaystyle\|(I-P_{\langle\\{y,z\\}\rangle})x\|_{2}^{2}=\|(I-P_{z})\circ(I-P_{y})x\|_{2}^{2}-\frac{\langle(I-P_{y})x,\
P_{z}\rangle^{2}}{\|(I-P_{z})y\|_{2}^{2}}.$

###### Proof.

For

$u=(I-P_{\langle\\{y,z\\}\rangle})x,\qquad
u^{\prime}=(I-P_{z})\circ(I-P_{y})x$

and any $v\in\langle\\{y,z\\}\rangle^{\perp}$ we note that

$\displaystyle\langle u,y\rangle=\langle u,z\rangle=\langle
u^{\prime},z\rangle=0$

and

$\displaystyle\langle u^{\prime},v\rangle=\langle x,v\rangle=\langle
u,v\rangle.$

Hence, we have that

$\displaystyle
P_{\langle\\{y,z\\}\rangle^{\perp}}u^{\prime}=P_{\langle\\{y,z\\}\rangle^{\perp}}x=P_{\langle\\{y,z\\}\rangle^{\perp}}u=u,$

which shows that

$\displaystyle\|u\|_{2}^{2}=\|P_{\langle\\{y,z\\}\rangle^{\perp}}u^{\prime}\|_{2}^{2}=\|u^{\prime}\|_{2}^{2}-\|P_{\langle\\{y,z\\}\rangle}u^{\prime}\|_{2}^{2}.$

Choosing the orthonormal basis $\big{(}z/\|z\|_{2},\
(I-P_{z})y/\|(I-P_{z})y)\|_{2}\big{)}$ of the space $\langle\\{y,z\\}\rangle$,
we further conclude that

$\displaystyle\|u\|_{2}^{2}$
$\displaystyle=\|u^{\prime}\|_{2}^{2}-\frac{\langle
u^{\prime},z\rangle^{2}}{\|z\|_{2}^{2}}-\frac{\langle(I-P_{z})(I-P_{y})x,\
(I-P_{z})y\rangle^{2}}{\|(I-P_{z})y\|_{2}^{2}}$
$\displaystyle=\|u^{\prime}\|_{2}^{2}-\frac{\langle(I-P_{y})x,\
(I-P_{z})y\rangle^{2}}{\|(I-P_{z})y\|_{2}^{2}}$
$\displaystyle=\|u^{\prime}\|_{2}^{2}-\frac{\langle(I-P_{y})x,\
P_{z}y\rangle^{2}}{\|(I-P_{z})y\|_{2}^{2}}.$

∎

######  Theorem 4.6.

Let Assumptions 1\- 3 be fulfilled. Let $x_{k}$, $x_{k}^{*}$ be the iterates
of Algorithm 3 and

$\displaystyle\tilde{y}_{k}=x_{k}-\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}\mathbf{A}^{T}S_{k}S_{k}^{T}(\mathbf{A}x_{k}-b).$
(29)

Then it holds that

$\displaystyle\mathbb{E}\big{[}D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})\big{]}$
$\displaystyle\leq\big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\big{)}\mathbb{E}\big{[}D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})\big{]}-\frac{\sigma}{2}\mathbb{E}\big{[}\frac{\langle\tilde{y}_{k},d_{k}^{*}\rangle^{2}}{\|d_{k}^{*}\|_{2}^{2}}\big{]}$
$\displaystyle\hskip
14.22636pt-\frac{\sigma}{2}\mathbb{E}\big{[}\frac{\langle\tilde{y}_{k},d_{k}^{*}\rangle^{2}\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}}{\|d_{k}^{*}\|_{2}^{2}\big{(}\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}\big{)}}\big{]}.$ (30)

In particular, $x_{k}$ converges in expectation to $\hat{x}$ in
$\mathcal{L}_{2}$-sense with the linear rate (13), that is,

$\displaystyle\mathbb{E}[\|x_{k}-\hat{x}\|_{2}^{2}]\leq\frac{2\varphi(\hat{x})}{\sigma}\Big{(}1-\frac{\sigma\lambda_{\min}(M)}{2\gamma}\Big{)}^{k}.$

###### Proof.

By (26) and the minimizing property of $\alpha_{k}$ and $\beta_{k}$ we
conclude that

$\displaystyle D_{\varphi}^{x_{k+1}^{*}}(x_{k+1},\hat{x})$ $\displaystyle\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})+\langle x_{k}-\hat{x},-\alpha\nabla
f_{S_{k}}(x_{k})+\beta d_{k}^{*}\rangle+\frac{1}{2\sigma}\|-\alpha\nabla
f_{S_{k}}(x_{k})+\beta d_{k}^{*}\|_{2}^{2}$
$\displaystyle=D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})+\frac{\sigma}{2}\big{(}\|x_{k}-\hat{x}+\frac{1}{\sigma}(-\alpha_{k}\nabla
f_{S_{k}}(x_{k})+\beta_{k}d_{k}^{*})\|_{2}^{2}-\|x_{k}-\hat{x}\|_{2}^{2}\big{)}$
$\displaystyle\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})+\frac{\sigma}{2}\inf_{\alpha,\beta\in\mathbb{R}}\big{(}\|x_{k}-\hat{x}+(-\alpha\nabla
f_{S_{k}}(x_{k})+\beta d_{k}^{*})\|_{2}^{2}-\|x_{k}-\hat{x}\|_{2}^{2}\big{)}$
$\displaystyle\leq
D_{\varphi}^{x_{k}^{*}}(x_{k},\hat{x})+\frac{\sigma}{2}\big{(}\|(I-P_{\langle\\{\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\\}\rangle})(x_{k}-\hat{x})\|_{2}^{2}-\|x_{k}-\hat{x}\|_{2}^{2}\big{)}.$
(31)

We can therefore apply Lemma 4.5 with $x_{k}-\hat{x}$ instead of $x$, $\nabla
f_{S_{k}}(x_{k})$ instead of $y$ and $d_{k}^{*}$ instead of $z$. We calculate
the expressions on the right-hand side in Lemma 4.5 as

$\displaystyle(I-P_{y})x$ $\displaystyle=x_{k}-\hat{x}-\frac{\langle
x_{k}-\hat{x},\nabla f_{S_{k}}(x_{k})\rangle}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}}\nabla f_{S_{k}}(x_{k})$
$\displaystyle=x_{k}-\hat{x}-\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}}\nabla f_{S_{k}}(x_{k}),$
$\displaystyle\|(I-P_{z})\circ(I-P_{y})x\|_{2}^{2}$
$\displaystyle=\|\tilde{y}_{k}-\hat{x}\|_{2}^{2}-\frac{\langle\tilde{y}_{k},d_{k}^{*}\rangle^{2}}{\|d_{k}^{*}\|_{2}^{2}}$
$\displaystyle=\|x_{k}-\hat{x}\|_{2}^{2}-\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{4}}{\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}}-\frac{\langle\tilde{y}_{k},d_{k}^{*}\rangle^{2}}{\|d_{k}^{*}\|_{2}^{2}}$

and as in the proof of Theorem 3.2 we can upper bound

$\displaystyle\|(I-P_{z})\circ(I-P_{y})x\|_{2}^{2}\leq\|x_{k}-\hat{x}\|_{2}^{2}-\frac{\|S_{k}^{T}(\mathbf{A}x_{k}-b)\|_{2}^{2}}{\|\mathbf{A}^{T}S_{k}\|_{2}^{2}}-\frac{\langle\tilde{y}_{k},d_{k}^{*}\rangle^{2}}{\|d_{k}^{*}\|_{2}^{2}}.$

Finally, we compute the right term on the right-hand side in Lemma 4.5 as

$\displaystyle\frac{\langle(I-P_{y})x,\
P_{z}y\rangle^{2}}{\|(I-P_{z})y\|_{2}^{2}}$
$\displaystyle=\frac{\langle\tilde{y}_{k},P_{d_{k}^{*}}\nabla
f_{S_{k}}(x_{k})\rangle^{2}}{\|\nabla f_{S_{k}}(x_{k})-P_{d_{k}^{*}}\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}}$
$\displaystyle=\frac{\langle\tilde{y}_{k},\frac{\langle d_{k}^{*},\nabla
f_{S_{k}}(x_{k})\rangle}{\|d_{k}^{*}\|_{2}^{2}}d_{k}^{*}\rangle^{2}}{\|\nabla
f_{S_{k}}(x_{k})-\frac{\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle}{\|d_{k}^{*}\|_{2}^{2}}d_{k}^{*}\|_{2}^{2}}$
$\displaystyle=\frac{\langle\tilde{y}_{k},d_{k}^{*}\rangle^{2}\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}}{\|d_{k}^{*}\|_{2}^{2}\big{(}\|\nabla
f_{S_{k}}(x_{k})\|_{2}^{2}\|d_{k}^{*}\|_{2}^{2}-\langle\nabla
f_{S_{k}}(x_{k}),d_{k}^{*}\rangle^{2}\big{)}}.$

After inserting all expressions into (31), the assertion follows as in the
proof of Theorem 3.2. ∎

As for Algorithm 2, we see that the convergence estimate (30) for the primal
iterates of Algorithm 3 is at least as good as the bound (12) for Algorithm 1.
Note that we can not directly relate the estimate (30) to the estimate (12)
for Algorithm 2, as $y_{k}$ and $\tilde{y}_{k}$ are different in general. For
interpretation of the rates in case of single row sketching and block
sketching, we refer to Example 3.3.

##  5 Numerical experiments

In this section, we study the computational behaviour of our two proposed
accelerations for the case of the sparse Kaczmarz method, that is, we choose
$\varphi(x)=\lambda\|x\|_{1}+\frac{1}{2}\|x\|_{2}^{2}$ and single row
sketching $S_{k}=e_{i_{k}}$. All experiments are conducted in MATLAB R2022b on
a macbook with 1,2 GHz Quad-Core Intel Core i7 processor and 16 GB memory. The
code to reproduce our results can be found at
https://github.com/MaxiWk/Minimal-error-momentum-Bregman-Kaczmarz.

In experiment (I), we evaluate the performance of the methods on artificial
Gaussian systems. The data was generated by sampling a matrix
$\mathbf{A}\in\mathbb{R}^{m\times n}$ with entries from the standard normal
distribution and a vector $\hat{x}\in\mathbb{R}^{n}$ with $10$ nonzero
entries, also from the standard normal distribution, at random positions
uniformly distributed over $\\{1,...,n\\}$. The vector $b$ was set to
$\mathbf{A}\hat{x}$ to ensure consistency of the system. In all examples, we
sampled the $i$-th row of the matrix $\mathbf{A}$ with probability
$p_{i}=\|a_{i}\|_{2}^{2}/\|\mathbf{A}\|_{F}^{2}$. In Figure 1, we report the
decay of the relative residual $\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$ of the
vanilla sparse Kaczmarz method (SRK), the exact-step sparse Kaczmarz method
(ESRK) and our two proposed accelerations, namely exact momentum with step
size $\alpha_{k}$ from vanilla sparse Kaczmarz (SRK-EM), and relaxed exact
momentum (SRK-REM). First, by comparing the scaling of the horizontal axis in
both subfigures, we observe that all three modifications of SRK accelerate the
method for $\lambda=5$ significantly, and much more than for $\lambda=0.1$.
Next, we compare performance of the methods in the right subfigure. The method
with relaxed momentum (SRK-REM) gives better acceleration than the (SRK-EM)
method, and in the median it is comparable to the acceleration achieved by the
ESRK method, which has more costly iterations. It is also visible that the
ESRK residual decays with a large variance over the random instances, which is
already known from [19]. This appears to be less the case for the SRK-REM
method. From Table 1 we see that the SRK-REM method needs by far the least
computation time to achieve a relative residual of $10^{-6}$. We also note
that it did not pay off to introduce exact momentum with the sparse Kaczmarz
step size (SRK-EM), compared to using the exact sparse Kaczmarz step size
without momentum (ESRK).

$50{,}000$$10^{-12}$$10^{-8}$$10^{-4}$$10^{0}$$k$$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$ | $500$$1{,}000$$10^{-12}$$10^{-8}$$10^{-4}$$10^{0}$$k$$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$
---|---

SRKESRKSRK-EMSRK-REM

Figure 1: Experiment (I) with normal distributed matrix $\mathbf{A}$, $(m,n)=(200,500)$, $\hat{x}$ with $10$ normal distributed nonzero entries, $50$ random repeats. Left: $\lambda=0.1$, right: $\lambda=5$. Thick line shows median over all trials, light area is between min and max, darker area indicates 25th and 75th quantile. In SRK-EM, in case that $\|d_{k}^{*}\|_{2}>\texttt{eps}$, we performed SRK steps. | CPU time (min) | CPU time (mean) | CPU time (max)
---|---|---|---
SRK | 0.050 | * | *
ESRK | 0.043 | 0.206 | 1.901
SRK-EM | 0.108 | 0.311 | 1.592
SRK-REM | 0.003 | 0.011 | 0.102
Table 1: CPU time (s) until a relative residual
$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}<10^{-6}$ is achieved in the problem from
Experiment (I), cf. Figure 1, right subfigure ($\lambda=5$). Here, “*”
indicates that the relative residual was not achieved after less than $10^{5}$
iterations. Columns display minimum, average and maximum time over $50$ random
repeats. SRK: sparse Kaczmarz, ESRK: exact-step sparse Kaczmarz, SRK-EM:
sparse Kaczmarz with exact momentum, SRK-REM: sparse Kaczmarz with relaxed
exact momentum.

In experiment (II), we test the methods on matrices from the
_SuiteSparseCollection_ [3]. To obtain a consistent system, here we chose
$\hat{x}$ with $10\%$ nonzero entries, standard normally distributed, and set
$b=\mathbf{A}\hat{x}$ with the respective matrix $\mathbf{A}$. The results are
given in Figure 2 and Figure 3. We observe different convergence behaviour in
several examples. In the examples in Figure 2, the SRK-REM method converges
fastest initially w.r.t. runtime, and is outperformed by the SRK-EM method for
small residuals, which then converges faster than all other methods. In the
examples in Figure 3, the SRK method is not visibly accelerated by our methods
even over iterations. In the last example, all accelerations except the SRK-
REM method perform worse than vanilla SRK.

$500$$10^{-12}$$10^{-8}$$10^{-4}$$10^{0}$$k$$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$$2\cdot
10^{-2}$$10^{-12}$$10^{-8}$$10^{-4}$$10^{0}$time
(s)$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$$10{,}000$$10^{-12}$$10^{-8}$$10^{-4}$$10^{0}$$k$$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$$0.1$$10^{-12}$$10^{-8}$$10^{-4}$$10^{0}$time
(s)$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$

SRKESRKSRK-EMSRK-REM

Figure 2: Experiment (II) with matrices from the SuiteSparseCollection,
relative residuals over iterations (left) and runtime (right), $\lambda=1$.
First row of figures: ash85 ($85\times 85$), second row: well1033 ($1033\times
320$). Thick line shows median over all trials, light area is between min and
max, darker area indicates 25th and 75th quantile. In SRK-EM, in case that
$\|d_{k}^{*}\|_{2}>10^{-6}$, we executed the step form SRK.

Finally, we study an academic CT example (experiment (III)) using a fan beam
tomography matrix from the _AIR tools_ package [7]. We choose $\hat{x}$ and
$b$ as in experiment (II) and set $\lambda=1$. The results are given in Figure
4. We can see that all modifications of the vanilla sparse Kaczmarz method
speed up convergence in terms of iterations. However, with respect to
computation time the proposed SRK-REM method is the only of our accelerations
which actually gives acceleration.

We also comment on an important detail for implementation. If
$\|d_{k}^{*}\|_{2}$ becomes small, the condition of the optimization problem
in the SRK-EM method can worsen drastically, which then leads to instabilities
and high oscillating errors. We therefore replace the if-condition in
Algorithm 2 by the condition $\|d_{k}^{*}\|_{2}>\texttt{tol}$ in practice. The
choice of tol is critical and depends on the concrete problem instances. For
too small values of tol, the method may oscillate, while for too large values
of tol, convergence may be slowed down. In experiment (I), we found that
$\texttt{tol}=\texttt{eps}$ is a good choice, where $\texttt{eps}\approx
2\cdot 10^{-16}$ is the MATLAB precision value. For $\texttt{tol}=10^{-6}$,
convergence was slowed down so much that in $10^{5}$ iterations a relative
residual of $10^{-6}$ was not achieved, which was not the case for any other
acceleration in the experiment (Table 1). In experiment (II) the choice
$\texttt{tol}=\texttt{eps}$ lead to high oscillation in the errors and we
mostly obtained good results with $\texttt{tol}=10^{-6}$.

$5\cdot
10^{5}$$10^{-3}$$10^{-2}$$10^{-1}$$k$$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$$10$$10^{-3}$$10^{-2}$$10^{-1}$time
(s)$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$$1\cdot
10^{6}$$10^{-3}$$10^{-2}$$10^{-1}$$10^{0}$$k$$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$$100$$10^{-3}$$10^{-2}$$10^{-1}$$10^{0}$time
(s)$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$

SRKESRKSRK-EMSRK-REM

Figure 3: Experiment (II) with matrices from the SuiteSparseCollection,
relative residuals over iterations (left) and runtime (right), $\lambda=1$.
First row of figures: ill1033 ($1033\times 320$), second row: landmark
($71952\times 2704$). Thick line shows median over all trials, light area is
between min and max, darker area indicates 25th and 75th quantile. In SRK-EM,
in case that $\|d_{k}^{*}\|_{2}>10^{-6}$, we executed the step form SRK.

$2\cdot 10^{5}$$10^{-8}$$10^{-6}$$10^{-4}$$10^{-2}$$10^{0}$$k$$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$ | $10$$10^{-8}$$10^{-6}$$10^{-4}$$10^{-2}$$10^{0}$time (s)$\|\mathbf{A}x_{k}-b\|_{2}/\|b\|_{2}$
---|---

SRKESRKSRK-EMSRK-REM

Figure 4: Experiment (III) with an academic CT matrix $\mathbf{A}$ arising
from fan beam tomography, $(m,n)=(3780,900)$, $\hat{x}$ with $10\%$ normal
distributed nonzero entries, $\lambda=1$. In SRK-EM, in case that
$\|d_{k}^{*}\|_{2}>10^{-6}$, we executed the step form SRK.

## References

  * [1] H. H. Bauschke and P. L. Combettes.  Convex analysis and monotone operator theory in Hilbert spaces.  CMS Books in Mathematics/Ouvrages de Mathématiques de la SMC. Springer, Cham, second edition, 2017.  With a foreword by Hédy Attouch.
  * [2] R. Bollapragada, T. Chen, and R. Ward.  On the fast convergence of minibatch heavy ball momentum.  arXiv preprint arXiv:2206.07553, 2022.
  * [3] T. A. Davis and Y. Hu.  The University of Florida sparse matrix collection.  ACM Transactions on Mathematical Software (TOMS), 38(1):1–25, 2011\.
  * [4] M. P. Friedlander and P. Tseng.  Exact regularization of convex programs.  SIAM Journal on Optimization, 18(4):1326–1350, 2008.
  * [5] R. Gower, F. Hanzely, P. Richtárik, and S. 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.
  * [6] R. M. Gower and P. Richtárik.  Randomized iterative methods for linear systems.  SIAM Journal on Matrix Analysis and Applications, 36(4):1660–1690, 2015.
  * [7] P. C. Hansen and M. Saxild-Hansen.  AIR tools—a MATLAB package of algebraic iterative reconstruction methods.  Journal of Computational and Applied Mathematics, 236(8):2167–2178, 2012.
  * [8] S. He, Z. Wang, and Q.-L. Dong.  Inertial randomized Kaczmarz algorithms for solving coherent linear systems.  arXiv preprint arXiv:2306.08185, 2023.
  * [9] B. Jarman, Y. Yaniv, and D. Needell.  Online signal recovery via heavy ball Kaczmarz.  arXiv preprint arXiv:2211.06391, 2022.
  * [10] J. Liu and S. Wright.  An accelerated randomized Kaczmarz algorithm.  Mathematics of Computation, 85(297):153–178, 2016.
  * [11] N. Loizou, M. Rabbat, and P. Richtárik.  Provably accelerated randomized gossip algorithms.  In ICASSP 2019-2019 IEEE international conference on acoustics, speech and signal processing (ICASSP), pages 7505–7509. IEEE, 2019\.
  * [12] N. Loizou and P. 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\.
  * [13] D. A. Lorenz, F. Schöpfer, and S. Wenger.  The linearized Bregman method via split feasibility problems: analysis and generalizations.  SIAM Journal on Imaging Sciences, 7(2):1237–1262, 2014.
  * [14] J. D. Moorman, T. K. Tu, D. Molitor, and D. Needell.  Randomized Kaczmarz with averaging.  BIT Numerical Mathematics, 61:337–359, 2021.
  * [15] M. S. Morshed, S. Ahmad, et al.  Stochastic steepest descent methods for linear systems: Greedy sampling & momentum.  arXiv preprint arXiv:2012.13087, 2020.
  * [16] M. S. Morshed, M. S. Islam, and M. Noor-E-Alam.  Accelerated sampling Kaczmarz Motzkin algorithm for the linear feasibility problem.  Journal of Global Optimization, 77:361–382, 2020.
  * [17] D. Needell and R. Ward.  Two-subspace projection method for coherent overdetermined systems.  Journal of Fourier Analysis and Applications, 19(2):256–269, 2013\.
  * [18] S. Petra.  Randomized sparse block Kaczmarz as randomized dual block-coordinate descent.  Analele ştiinţifice ale Universităţii” Ovidius” Constanţa. Seria Matematică, 23(3):129–149, 2015.
  * [19] F. Schöpfer and D. A. Lorenz.  Linear convergence of the randomized sparse Kaczmarz method.  Mathematical Programming, 173(1):509–536, 2019.
  * [20] F. Schöpfer, D. A. Lorenz, L. Tondji, and M. Winkler.  Extended randomized Kaczmarz method for sparse least squares and impulsive noise problems.  Linear Algebra and its Applications, 652:132–154, 2022.
  * [21] T. Strohmer and R. Vershynin.  A randomized Kaczmarz algorithm with exponential convergence.  Journal of Fourier Analysis and Applications, 15(2):262, 2009.
  * [22] L. Tondji and D. A. Lorenz.  Faster randomized block sparse Kaczmarz by averaging.  Numerical Algorithms, pages 1–35, 2022.
  * [23] L. Tondji, D. A. Lorenz, and I. Necoara.  An accelerated randomized Bregman-Kaczmarz method for strongly convex linearly constraint optimization.  In 2023 European Control Conference (ECC), pages 1–6. IEEE, 2023\.
  * [24] Y. Zeng, D. Han, Y. Su, and J. Xie.  On adaptive stochastic heavy ball momentum for solving linear systems.  arXiv preprint arXiv:2305.05482, 2023.


AtP*+GD | $\frac{L+1}{2}DH^{2}$ | $\frac{L+1}{2}(2D+T+1)KH^{2}$ | $VNH$ | $\frac{L+1}{2}DH$ | $\frac{L+1}{2}(2D+T+1)KH$ | $ND$
Neurons | $DH^{2}$ | $2DH^{2}$ | $VNH$ | $VNH$ | $2KNH$ | $N^{2}$
MLPfix | $DH^{2}$ | $2DH^{2}$ | $VNH$ | $VNH$ | $2KNH$ | $N^{2}$
AtP* | $DH^{2}$ | $(T+1)DH^{2}$ | $VNH$ | $VNH$ | $(T+1)KNH$ | $N^{2}$
AtP*+GD | $\frac{L+1}{2}DH^{2}$ | $\frac{L+1}{2}(T+1)DH^{2}$ | $VNH$ | $\frac{L+1}{2}VNH$ | $\frac{(L+1)(T+1)}{2}KNH$ | $N^{2}$
ATP* (long) | $DH^{2}$ | $(2D+T+1)KH^{2}$ | $VNH$ | $VNH$ | $(T+1)KNH$ | $N^{2}$
AtP*+GD | $\frac{L+1}{2}DH^{2}$ | $\frac{L+1}{2}(2D+T+1)KH^{2}$ | $VNH$ | $\frac{L+1}{2}VNH$ | $\frac{(L+1)(T+1)}{2}KNH$ | $N^{2}$
Table 2: Per-token per-layer-pair total quadratic cost of each kind of
between-layers edge, across edge-AtP variants. For brevity, we omit the layer-
pair $\binom{L}{2}$ factor that would otherwise be in every cell, and use
$D:=d_{\mathrm{resid}},H:=\text{\\# heads per
layer},K:=d_{\mathrm{key}},V:=d_{\mathrm{value}},N:=d_{\mathrm{neurons}}$.
Figure 16: A comparison of edge-AtP variants across model sizes and prompt
lengths. AtP* here is defined to include QKfix and MLPfix, but not GradDrops.
The costs vary across several orders of magnitude for each setting.
In the setting with full-MLP nodes, MLPfix carries substantial cost for short
prompts, but barely matters for long prompts.
In the neuron-nodes setting, MLPfix is costless. But GradDrops in that setting
continues to impose a large cost; even though it doesn’t affect
MLP$\rightarrow$MLP edges, it does affect MLP$\rightarrow$Q,K edges, which
come out dominating the cost with QKfix.

##  Appendix D Distribution of true effects

In Figure 17(d), we show the distribution of $c(n)$ across models and
distributions.

Figure 17: Distribution of true effects across models and prompt pair
distributions

AttentionNodes

NeuronNodes

a.i

a.ii

(a) Pythia-410M

b.i

b.ii

(b) Pythia-1B

c.i

c.ii

(c) Pythia-2.8B

d.i

d.ii

(d) Pythia-12B