Greedy Learning to Optimize with Convergence Guarantees

## Abstract

Abstract Learning to optimize is an approach that leverages training data to accelerate the solution of optimization problems. Many approaches use unrolling to parametrize the update step and learn optimal parameters. Although L2O has shown empirical advantages over classical optimization algorithms, memory restrictions often greatly limit the unroll length and learned algorithms usually do not provide convergence guarantees. In contrast, we introduce a novel method employing a greedy strategy that learns iteration-specific parameters by minimizing the function value at the next iteration. This enables training over significantly more iterations while maintaining constant memory usage. We parameterize the update such that parameter learning corresponds to solving a convex optimization problem at each iteration. In particular, we explore preconditioned gradient descent with multiple parametrizations including a novel convolutional preconditioner. With our learned algorithm, convergence in the training set is proven even when the preconditioner is neither symmetric nor positive definite. Convergence on a class of unseen functions is also obtained, ensuring robust performance and generalization beyond the training data. We test our learned algorithms on two inverse problems, image deblurring and Computed Tomography, on which learned convolutional preconditioner demonstrates improved empirical performance over classical optimization algorithms such as Nesterov’s Accelerated Gradient Method and the quasi-Newton method L-BFGS.

## 1 Introduction

We consider the optimization problem

$$ $\min_{x}f(x),$ (1) $$

with the assumption that $f:\mathcal{X}\to\mathbb{R}$ is convex, $L$-smooth and bounded below, where $\mathcal{X}$ is a Hilbert space. Classic optimization methods are built in a theoretically justified manner, with guarantees on their performance and convergence properties. For example, Nesterov’s Accelerated Gradient Method (NAG) (Nesterov, ) accelerates classical first-order algorithms using momentum. However, practitioners often concentrate on problems within a much smaller class. For example, in reconstructing images from blurred observations $y$ generated by a blurring operator $A$, one might minimize a function from the class:

$$ $\mathcal{F}=\left\{f:\mathcal{X}\to\mathbb{R}:f(x)=\frac{1}{2}\|Ax-y\|^{2}+R(x ),y\in\mathcal{Y}\right\},$ (2) $$

where $R:\mathcal{X}\to\mathbb{R}$ is a chosen regularizer. Learning to optimize (L2O) uses data to learn how to minimize functions $f\in\mathcal{F}$ in a small number of iterations. Typically, the solution at each iteration $t$ is updated by a parametrized function $G_{\theta}:\mathcal{X}\times\mathcal{X}\to\mathcal{X}$ (i.e. the update rule) as dependent on parameters $\theta_{t}$ at iteration $t$ as

$$ $x_{t+1}=x_{t}-G_{\theta_{t}}(x_{t},\nabla f(x_{t})).$ (3) $$

Unrolling algorithms e.g. Monga et al. () directly parametrize the update step as a neural network, often taking the previous iterates of the solution updates and the gradients as input arguments to the neural network. For some $T>0$, the parameters $\theta=(\theta_{0},\cdots,\theta_{T-1})$ can be learned to minimise the loss

$$ $L(\theta)=\mathop{\mathbb{E}}_{f\in\mathcal{F}}\left[\sum_{t=1}^{T}f(x_{t}) \right].$ (4) $$

Learned optimization algorithms often lack convergence guarantees, including many that use RNNs (Andrychowicz et al., ; Metz et al., ) or Reinforcement learning (Li & Malik, ). Liu et al. () consider methods of the form $x_{t+1}=x_{t}-G_{t}\nabla f(x_{t})+b_{t}$, for $f:\mathbb{R}^{n}\to\mathbb{R}$, a diagonal matrix $G_{t}\in\mathbb{R}^{n\times n}$, and a vector $b_{t}\in\mathbb{R}^{n}$. The $G_{t}$ and $b_{t}$ are constructed using the outputs of neural networks. However, their method does not guarantee convergence to a minimizer.

Other approaches achieve provable convergence, which can be enforced with safeguarding (Heaton et al., ), or constructing convergent algorithms by learning parameters within a provably convergent set (Banert et al., ; ). Tan et al. (; ) learn mirror maps using input-convex neural networks within the mirror descent optimization algorithm such that the algorithm is provably convergent. Lastly, Sucker et al. () and Sambharya & Stellato () consider applying the PAC-Bayes framework to L2O.

Unlike NAG, Newton’s method accelerates convergence by applying the inverse Hessian to the gradient, which can be costly in practice. Quasi-Newton methods like BFGS (Nocedal & Wright, ) approximate the Hessian, and L-BFGS (Liu & Nocedal, ) is used when BFGS is too memory-intensive. Similarly, we aim to accelerate the optimization by learning a preconditioner $G_{t}$ in the update $x_{t+1}=x_{t}-G_{t}\nabla f(x_{t})$.

### 1.1 Contributions

Our paper contributes in the following ways:

- •
A novel approach to L2O that learns parameters at each iteration sequentially, using a greedy approach by minimizing the function value at the next iteration. This enables training over significantly more iterations while maintaining constant memory usage: Section 3.
- •
Convergence in the training set is proven even when the preconditioner is neither symmetric nor positive definite: Section 4.
- •
Convergence is proven on a class of unseen functions under certain conditions: Section 4.
- •
Learning parameters is a convex optimization problem for ’linear parametrizations’ of $G_{t}$, enabling training that is significantly faster, with closed-form solutions for least-squares functions: Section 5.
- •
A novel parametrization of $G_{t}$ as a convolution operator. At iteration $t$ we learn a convolutional kernel $\kappa_{t}$ such that $G_{t}x=\kappa_{t}\ast x$. This parametrization is shown to outperform Nesterov’s Accelerated Gradient and L-BFGS on test data: Section 6.

In Section 6, we validate our learned algorithms on two inverse problems: image deblurring and Computed Tomography (CT). Inverse problems represent a crucial class of optimization problems that appear in important fields such as medical imaging and machine learning. Many such problems have an associated forward operator which is highly ill-conditioned, making them an ideal test for optimization algorithms.

## 2 Notation

We require the following notation. Let $\mathcal{X}$ be a Hilbert space with its corresponding field given by $\mathbb{R}$ and norm $\|\cdot\|$.
A function $f:\mathcal{X}\rightarrow\mathbb{R}$ is convex if for all $x,y\in\mathcal{X}$ and for all $\alpha\in[0,1]$:

$$ $f(\alpha x+(1-\alpha)y)\leq\alpha f(x)+(1-\alpha)f(y).$ (5) $$

A function $f:\mathcal{X}\rightarrow\mathbb{R}$ is L-smooth with parameter $L>0$ if its gradient is Lipschitz continuous, i.e., if for all $x,y\in\mathcal{X}$, the following inequality holds:

$$ $\|\nabla f(x)-\nabla f(y)\|\leq L\|x-y\|.$ (6) $$

A function $f:\mathcal{X}\rightarrow\mathbb{R}$ is bounded below if there exists some $M\in\mathbb{R}$ such that $f(x)\geq M$ for all $x\in\mathcal{X}$. We say that $f\in\mathcal{F}_{L}$ if $f$ is convex, $L$-smooth, and bounded below.
We assume that the Hilbert space $\mathcal{X}$ has dimension $\operatorname{dim}(\mathcal{X})=n$ and, therefore, admits a finite orthonormal basis $\{e_{1},\cdots,e_{n}\}$. For $x\in\mathcal{X}$ and $j\in\{1,\cdots,n\}$, define $[x]_{j}:=\langle x,e_{j}\rangle$. For $x,y\in\mathcal{X}$, define the pointwise product $x\odot y$ by

$$ $[x\odot y]_{j}=[x]_{j}[y]_{j}.$ (7) $$

For Hilbert spaces $\mathcal{X}$ and $\mathcal{Y}$, denote the space of linear operators from $\mathcal{X}$ to $\mathcal{Y}$ by $\mathcal{L}(\mathcal{X},\mathcal{Y})$. If $\mathcal{Y}=\mathcal{X}$, we write $\mathcal{L}(\mathcal{X})$. For example, if $\mathcal{X}=\mathbb{R}^{n}$, $\mathcal{L}(\mathcal{X})$ is the space of $n\times n$ matrices. Denote the adjoint of $A\in\mathcal{L}(\mathcal{X},\mathcal{Y})$ by $A^{*}$, meaning that for $x\in\mathcal{X},y\in\mathcal{Y}$, $\langle Ax,y\rangle=\langle x,A^{*}y\rangle$. Denote by $I\in\mathcal{L}(\mathcal{X})$ the identity operator: $I(x)=x$ for all $x\in\mathcal{X}$.

## 3 Greedy learning to optimize of preconditioned gradient descent

This section introduces the proposed method: greedy learning to optimize. Firstly, we introduce how we parametrize the optimization algorithm as preconditioned gradient descent. Next, we detail our training data and define a loss function with which we learn parameters. We then provide an algorithm of how parameters are learned sequentially using a greedy approach. Lastly, we show how our learned algorithm is applied to unseen functions.

At each iteration $t\in\{0,1,2,\cdots\}$, we parametrize the linear operator $G_{t}$ using a Hilbert space $\Theta$ and learn parameters $\theta_{t}\in\Theta$ in the update

$$ $x_{t+1}=x_{t}-G_{\theta_{t}}\nabla f(x_{t}).$ (8) $$

The following propositions show that it is possible to obtain convergence after just one iteration of the update (8). Firstly, we show that it is possible to even when $G$ is a pointwise operator, i.e. $Gx:=p\odot x$ for some $p\in\mathcal{X}$.

###### Proposition 1 .

Assume that $f:\mathcal{X}\to\mathbb{R}$ is convex, continuously differentiable, and has a global minimum, and take any initial point $x_{0}\in\mathcal{X}$. Then there exists $p\in\mathcal{X}$ such that, $x_{0}-p\odot\nabla f(x_{0})\in\arg\min_{x}f(x)$.

While the pointwise parametrization obtains convergence after one iteration for one function, for an arbitrary linear operator $G\in\mathcal{L}(\mathcal{X})$, under certain conditions, one can obtain convergence after one iteration for multiple functions.

###### Proposition 2 .

For $k\in\{1,\cdots,N\}$, assume that $f_{k}:\mathcal{X}\to\mathbb{R}$ is convex, continuously differentiable, and has a global minimum, with any initial point $x_{k}^{0}\in\mathcal{X}$. Assume that the set of gradients $\{\nabla f_{1}(x_{1}^{0}),\cdots,\nabla f_{N}(x_{N}^{0})\}$ is linearly independent. Then if $N\leq n$, there exists an operator $P\in\mathcal{L}(\mathcal{X})$ such that $x_{k}^{0}-P\nabla f_{k}(x_{k}^{0})\in\arg\min_{x}f_{k}(x)$, for all $k\in\{1,\cdots,N\}$.

Propositions 1 and 2 motivate learning $\theta_{t}$ by considering the function values only at the next iteration. In order to learn the parameters $\theta_{t}$ for $t\in\{0,1,2,\cdots\}$, we use a training dataset of functions $\mathcal{T}:=\{f_{1},\cdots,f_{N}\}$, with $f_{k}\in\mathcal{F}_{L_{k}}$ for $k\in\{1,\cdots,N\}$, with corresponding initial points $\mathcal{X}_{0}:=\{x_{1}^{0},\cdots,x_{N}^{0}\}$.

We consider learning parameters using a regularizer $R:\Theta\to\mathbb{R}$ so that undesirable properties are penalized. At iteration $t$, we solve the optimization problem

$$ $\displaystyle\theta_{t}\in\arg\min_{\theta}\left\{g_{t,\lambda_{t}}(\theta):= \frac{1}{N}\sum_{k=1}^{N}f_{k}(x_{k}^{t}-G_{\theta}\nabla f_{k}(x_{k}^{t}))+ \lambda_{t}R(\theta)\right\},$ (9) $$

for some regularization parameter $\lambda_{t}\geq 0$, which is used to balance the importance of the regularizer. Such a strategy is greedy, as learning refers to tuning the parameters $\theta_{t}$ considering only the function values at the next iteration, $f_{k}(x_{k}^{t+1})$. The sequential training procedure for parameter learning is detailed in Algorithm 1. Unlike unrolling, where memory requirements scale linearly with the number of training iterations, greedy training is a memory-friendly method, with memory that is constant with increasing training iterations.

Figure: Algorithm 1 Training algorithm for greedy parameter learning in preconditioned gradient descent

Suppose that training is terminated after iteration $T$, having learned the parameters $\theta_{0},\cdots,\theta_{T-1}$. To minimise an unseen function $f$ with initial point $x_{0}$, we propose Algorithm 2.

Figure: Algorithm 2 Learned algorithm applied to a new function $f$

###### Proposition 1 .

###### Proposition 2 .

## 4 Convergence results

This section contains convergence results for our learned Algorithm 2. Firstly, convergence is obtained on training functions as $T\to\infty$, without the need for the learned operators $G_{\theta_{t}}$ to have properties such as being symmetric or positive definite. Following this, we show convergence results and rates for a class of unseen functions if $\lambda_{t}$ is asymptotically non-vanishing. Before we present the convergence results, we require the following definitions, the first of which provides a condition for which the update rule (8) generalizes gradient descent (GD): $x_{t+1}=x_{t}-\alpha_{t}\nabla f(x_{t})$ for $\alpha_{t}>0$.

###### Definition 1 .

We say that the family $(G_{\theta})$ is GGD (generalizes gradient descent) if for all $\alpha>0$, there exist parameters $\theta$ such that

$$ $G_{\theta}=\alpha I.$ (10) $$

Parametrizations of $(G_{\theta})$ that satisfy the GGD property are shown in section 5. Let $\tau=1/{L_{\text{train}}}$, where $L_{\text{train}}=\max\{L_{1},\cdots,L_{N}\}$ is the largest smoothness coefficient in the training data set. This choice of step size in gradient descent ensures convergence for all functions $f_{k}\in\mathcal{T}$. From this point forward, we assume $(G_{\theta})$ is GGD, meaning there exists some $\tilde{\theta}$ such that $G_{\tilde{\theta}}=\tau I$. Furthermore, the GGD property can be leveraged to establish provable convergence for a set of unseen functions by introducing a penalty when the parameters deviate significantly from $\tilde{\theta}$. With this purpose, we define $R(\theta)$ in (9) as

$$ $R(\theta):=\frac{1}{2}\|\theta-\tilde{\theta}\|_{2}^{2}.$ (11) $$

The next definition is to ensure the parametrized algorithm adopts the convergence properties of gradient descent on the training data.

###### Definition 2 .

We say that $\theta_{t}$ is $BGD$ (better than gradient descent) with regularization parameter ${\lambda_{t}}$ if

$$ $g_{t,{\lambda_{t}}}(\theta_{t})\leq g_{t,0}(\tilde{\theta})=\frac{1}{N}\sum_{k =1}^{N}f_{k}\left(x_{k}^{t}-\tau\nabla f_{k}(x_{k}^{t})\right).$ (12) $$

###### Definition 1 .

###### Definition 2 .

### 4.1 Convergence on training data

###### Theorem 1 .

Convergence on training data. Suppose that ${\lambda_{t}}\geq 0$ and $(\theta_{t})_{t=0}^{\infty}$ is a $BGD$ sequence of parameters. Then with Algorithm (1), we have $\nabla f_{k}(x_{k}^{t})\to 0$ as $t\to\infty$ for all $k\in\{1,\cdots,N\}$.

Convergence rates can also be obtained for training data, see Appendix Section 4.

###### Theorem 1 .

### 4.2 Convergence on unseen data

We now show convergence on unseen data. Firstly, we show that if the regularization parameters $\lambda_{t}$ are eventually non-vanishing, then the learned parameters tend towards $\tilde{\theta}$.

###### Lemma 1 .

If $\lim\inf_{t\to\infty}\lambda_{t}>0$ and $(\theta_{t})_{t=0}^{\infty}$ is $BGD$, then $\theta_{t}\to\tilde{\theta}\text{ as }t\to\infty$.

This result is useful to ensure convergence on unseen data as if $G:\Theta\to\mathcal{L}(\mathcal{X})$ is continuous, then under the same conditions, $G_{\theta_{t}}\to G_{\tilde{\theta}}=\tau I$ as $t\to\infty$, i.e. our learned algorithm gets close to GD for large $t$. The idea is that we start with a method that fits the data very well leading to quick initial convergence, but in the interest of safety, over time we become closer to an algorithm with proven convergence.

###### Theorem 2 .

Convergence on unseen data for regularized parameter learning
Assume that $G:\Theta\rightarrow\mathcal{L}(\mathcal{X})$ is continuous, $\theta_{t}$ is BGD and $\lim\inf_{t\to\infty}\lambda_{t}>0$. Then, there exists a constant $\tilde{L}>0$, and a final training iteration $T$ such that for all $f\in\mathcal{F}_{\tilde{L}}$ and any starting point $x_{0}$, using Algorithm 2 (which depends on $T$), we have $\nabla f\left(x_{t}\right)\rightarrow 0$ as $t\rightarrow\infty$.

In practice, provable convergence can be verified during training. At iteration $t$, the regularization parameter $\lambda_{t}$ may be selected large enough such that $\|G_{\theta_{t}}-\tau I\|<\tau$, which guarantees convergence if training is then stopped with $T=t$. A proof is provided with Proposition 6 in the Appendix. The following theorem presents the convergence rate obtained for test functions.

###### Theorem 3 .

Convergence rates on unseen data for regularized parameter learning
Under the same assumptions as Theorem 2 and if $(x_{t})_{t=1}^{\infty}$ is a bounded sequence then there exists a constant $C>0$, such that

$$ $f(x_{t})-f(x^{*})\leq\frac{C}{t}.$ (13) $$

This result gives the worst-case convergence rate of the learned algorithm. In Section 6 we will see that the empirical performance of the learned algorithms may exceed that of NAG and L-BFGS.

###### Lemma 1 .

###### Theorem 2 .

###### Theorem 3 .

## 5 Linear parametrizations

In this section, we consider ’linear parametrizations’ of $G$, defined below.

###### Definition 3 .

We call $G$ a linear parameterization if $G:\Theta\to\mathcal{L}(\mathcal{X})$ is a linear map.
This means there exists a linear operator $B_{k}^{t}\in\mathcal{L}(\Theta,\mathcal{X})$ such that

$$ $G_{\theta}\nabla f_{k}(x_{k}^{t})=B_{k}^{t}(\theta).$ (14) $$

The motivation is that when $G$ is a linear parametrization, each optimization problem (9) is convex (as it is the composition of a convex function with a linear function (Beck, )). Therefore, learning comprises solving a sequence of convex optimization problems. In this case, there exist fast, provably convergent algorithms to find global solutions. Such problems also require fewer hyperparameters to tune. Due to the speed of training enabled by linear parametrizations and the constant memory of the greedy approach, we are able to learn algorithms up to significantly higher iterations. Four examples of linear parametrizations of $G$ are provided in Table 1. These parametrizations are used for the numerical experiments in Section 6.

**Table 1: Examples of linear parametrizations**
| Label | Description | parametrization | # parameters |
| --- | --- | --- | --- |
| (PS) | Scalar step size | $G_{\alpha_{t}}=\alpha_{t}I,\alpha_{t}\in\mathbb{R}$ | $1$ |
| (PP) | Pointwise operator | $G_{p_{t}}x=p_{t}\odot x,p_{t},x\in\mathcal{X}$ | $\operatorname{dim}(\mathcal{X})$ |
| (PC) | Image convolution | $G_{\kappa_{t}}x=\kappa_{t}\ast x,\kappa_{t}\in\mathbb{R}^{m_{1}\times m_{2}}$ | $m_{1}m_{2}$ |
| (PF) | Full linear operator | $G_{P_{t}}=P_{t}\in\mathcal{L}(\mathcal{X})$ | $\operatorname{dim}(\mathcal{X})^{2}$ |

###### Lemma 2 .

All parametrizations $G_{\theta}$ in Table 1 satisfy the GGD property (10), and are all continuous with respect to their parameters.

###### Corollary 1 .

We obtain the convergence results from Theorems 2 and 3 for linear parametrizations in Table 1, when $\theta_{t}$ is BGD and $\lim\inf_{t\to\infty}\lambda_{t}>0$.

###### Definition 3 .

###### Lemma 2 .

###### Corollary 1 .

### 5.1 Closed-form solutions

If each function $f_{k}\in\mathcal{T}$ can be written as a least-squares function, then the parameters $\theta_{t}$ at iteration $t$ have a closed-form solution. For $k\in\{1,\cdots,N\}$, let $f_{k}$ be given by

$$ $f_{k}(x)=\frac{1}{2}\|A_{k}x-y_{k}\|^{2},$ (15) $$

with corresponding $y_{k}\in\mathcal{Y}$, for a Hilbert space $\mathcal{Y}$, and linear operator $A_{k}\in\mathcal{L}(\mathcal{X},\mathcal{Y})$.

###### Proposition 3 .

For a linear parametrization $G$, let $B_{k}^{t}$ be given as in (14), and $f_{k}$ as in (15). Then $\theta_{t}$ given by

$$ $\theta_{t}=\left({\lambda_{t}}I_{\Theta}+\frac{1}{N}\sum_{k=1}^{N}(A_{k}B_{k}^ {t})^{*}(A_{k}B_{k}^{t})\right)^{\dagger}\left({\lambda_{t}}\tilde{\theta}+ \frac{1}{N}\sum_{k=1}^{N}(B_{k}^{t})^{*}\nabla f_{k}(x_{k}^{t})\right)$ (16) $$

is a solution to (9), where $M^{\dagger}$ represents the Moore–Penrose pseudoinverse of a linear operator $M$.

Note, for this class of functions, we also have a closed-form equation for exact line search for a scalar step size (Nocedal & Wright, ), which we recover with $\lambda_{t}=0,N=1$ for the scalar parametrization (PS) in Table 1. Therefore the optimization problem (9) can be seen as an extension of exact line search to include linear operators. Calculations for the closed-form solutions for the four examples in Table 1 are detailed in Appendix Section 5. In general, we require optimization algorithms to approximate $\theta_{t}$. Due to the optimization problem being convex, we provide gradient calculations and smoothness constants for these parametrizations in Appendix Section C.1.

###### Proposition 3 .

## 6 Numerical Experiments

In this section, we test the four linear parametrizations in Table 1 on two inverse problems in imaging: image deblurring and CT. We consider linear inverse problems, defined by receiving an observation $y\in\mathcal{Y}$, generated from a ground-truth $x_{\text{true}}$ via some linear forward operator $A:\mathcal{X}\to\mathcal{Y}$, such that $y=Ax_{\text{true}}+\varepsilon$, where $\varepsilon\in\mathcal{Y}$ is some random noise, and the goal is to recover $x_{\text{true}}$. For both experiments, $\mathcal{X}=\mathbb{R}^{h_{1}\times h_{2}}$ for $h_{1},h_{2}\in\mathbb{N}$, and $\varepsilon$ is noise sampled from a zero-mean Gaussian distribution. To approximate $x_{\text{true}}$ from $y$, we solve

$$ $\min_{x}\left\{f(x):=\frac{1}{2}\|Ax-y\|_{2}^{2}+\alpha H_{\epsilon}(x)\right\},$ (17) $$

for a given regularization parameter $\alpha$. The regularizer $H_{\epsilon}$ is the Huber Total Variation (Rudin et al., ; Huber, ) defined by

$$ $H_{\epsilon}(x)=\sum_{i,j=1}^{h_{1},h_{2}}h_{\epsilon}\left(\sqrt{(\mathrm{D}x )_{i,j,1}^{2}+(\mathrm{D}x)_{i,j,2}^{2}}\right),\quad h_{\varepsilon}(s)= \begin{cases}\frac{1}{2\epsilon}{s^{2}},\text{ if }|s|\leq\epsilon\\ |s|-\frac{\epsilon}{2},\text{ otherwise},\end{cases}$ (18) $$

where finite difference operator $\mathrm{D}:\mathbb{R}^{h_{1}\times h_{2}}\rightarrow\mathbb{R}^{h_{1}\times h_
{2}\times 2}$ is defined in Chambolle & Pock (). We take $\epsilon=0.01$ and normalize the forward operator in both cases so that $\|A\|=1$. Then, each function $f$ is $L$-smooth, where $L=1+\frac{8\alpha}{\epsilon}$ (Chambolle & Pock, ).

For each parametrization in Table 1, to learn parameters $\theta_{t}$ we apply NAG for solving the optimization problem 9. The initialization is set as $\theta_{t}^{0}=\tilde{\theta}$ for $t=0$, and $\theta_{t}^{0}=\theta_{t-1}$ for $t>0$. NAG is stopped when $\|\nabla g_{t}(\theta_{t}^{\ell})\|_{2}/\|\nabla g_{t}(\theta_{t}^{0})\|_{2}<1
0^{-3}$, or when $\ell=\ell_{\text{stop}}=5000$. For both problems, we use a training set of $100$ functions for parametrizations (PS), (PP), and (PC). For the parametrization (PF), the model is trained using $1000$ functions and is only implemented for the CT problem. Testing is performed on a separate set of $100$ functions for all parametrizations. The learned convolutional kernel (PP) has dimensions $h_{1}\times h_{2}$, matching the size of the images in $\mathcal{X}$.

Given a dataset of functions ${f_{1},\dots,f_{N}}$, the mean value at iteration $t$ is defined as $F(x_{t})=\frac{1}{N}\sum_{k=1}^{N}f_{k}(x_{k}^{t})$ and let $F^{*}:=\min_{x\in\mathcal{X}^{N}}F(x)$. This function is used to visualize the performance of the learned methods. The learned algorithms are compared with NAG, L-BFGS, and GD with constant step size $\tau$. All computations were carried out on an Nvidia RTX 3600 12GB GPU.

### 6.1 Image deblurring

In this example, the forward operator $A$ in (17) is a Gaussian blur with a $5\times 5$ kernel size and a standard deviation of $\sigma=1.5$. We use the STL-10 dataset (Coates et al., ) with greyscale images of size $96\times 96$ as $\mathcal{X}$. The noise $\varepsilon$ is modeled with a standard deviation of $2.5\times 10^{-3}$, and we set $\alpha=10^{-5}$, resulting in $L=1.008$. The initial point $x_{0}$ is taken equal to the observation $y$.

Training with the greedy method was performed up to iteration $T=600$ with $\lambda_{t}=0$ for all $t$ for the parametrizations (PS), (PP), (PC). This resulted in learning $600$ parameters for (PS) and $5529600$ for both (PP) and (PC). For the convolutional parametrization (PC) at iterations $t>50$ we set $\ell_{\text{stop}}=500$. The total training time for (PS) and (PP) was approximately $10$ minutes each, while (PC) required $19$ hours.

Figure 1(a) shows that the learned scalar parameters (PS) eventually fluctuate above and below $2/L$, which is outside of the range of provable convergence of gradient descent with a constant step size. Despite this, the learned algorithm leads to convergence on training data as $t\to\infty$ by Theorem 1. In Figure 1(b), Oscillatory behavior also occurs in the pointwise parametrization (PP), with most fluctuations occurring near the boundary. We also see values outside the interval $(0,2/L)$, including negative values. The learned convolutional kernels (PC) in Figure 1(c) also contain positive and negative values. These kernels are predominantly weighted towards the center, suggesting that information from neighboring pixels is prioritized over more distant ones. As the number of iterations increases, the kernels exhibit increasing similarity, though no formal convergence result for $\theta_{t}$ has been established when $\lambda_{t}=0$.

Figure: (a)
Refer to caption: extracted/5992054/blur_images/learned_scalars.png

Figure 2 shows that the learned parametrizations (PS), (PP), and (PC) generalize well to unseen data due to the closeness of the train and test curves. The learned pointwise parametrization performs very similarly to the learned scalar parameterization for this example. This is despite (PP) having an equal number of parameters as the learned convolutional kernels, which use their parameters to capture global information across the image, rather than focusing on specific pixel-level details. Furthermore, the learned convolutional algorithm significantly outperforms both NAG and L-BFGS on the test data, reaching a tolerance of $10^{-6}$ in just over $100$ iterations, compared with about $400$ and $450$ for L-BFGS and NAG, respectively. More information on the number of iterations to reach specific tolerances can be found in Appendix Section D.3. Additionally, Appendix Section D.1 provides comparisons of the best and worst performance between the learned convolutional parametrization and NAG, while D.2 explores the impact of different kernel sizes on the performance of the (PC) parametrization.

We now consider the effect of regularization on the convergence of the learned Algorithm 2. We use regularization at iteration $T$ for the parametrizations (PS) and (PP) to ensure convergence when applying Algorithm 2 to further iterations. The (PC) parametrization was not considered as it has already reached a suitable tolerance within the training iterations. As discussed in Section 4, we may select $\lambda_{T}$ large enough such that the learned parameters lead to guaranteed convergence. For the (PS) parametrization, we find $\lambda_{T}=2\times 10^{-7}$ guarantees convergence of Algorithm 2, and for the (PP) parametrization, $\lambda_{T}=3.90625\times 10^{-9}$ guarantees convergence. Figure 2 shows that when learned without regularization, the learned algorithm quickly diverges. However, if the regularization parameter $\lambda_{T}$ is chosen large enough, the algorithm converges as proved in Theorem 2.

Figure: Figure 2: Left: Performance comparison of the learned methods on training versus test data for the image deblurring problem within training iteration. Center: Test performance versus benchmark optimization algorithms within training iterations. Right: Performance beyond training iterations for the (PP) and (PS) parameterizations when learned with and without guaranteed convergence. The final training iteration $T$ is indicated by the dotted vertical line.
Refer to caption: extracted/5992054/blur_images/generalisation_performance.png

Figure 3 illustrates the reduced number of iterations needed to obtain a deblurred image reconstruction using the learned convolution parameterization. It requires approximately four times more iterations to achieve a comparable reconstruction quality with NAG compared to the learned convolutional algorithm.

Figure: (a) Observation
Refer to caption: extracted/5992054/blur_images/comparisons/obs_zoom.png

We now compare the time taken for training with the greedy learning approach versus the unrolling approach. In the unrolling approach, we fix the number of iterations to $T=20$ and jointly learn the parameters $\theta_{0},\cdots,\theta_{T-1}$ in the update rule (8) with the (PC) parametrization. The same training dataset as the greedy method is used, using a batch size of $4$, with the objective function defined in equation (4). Parameter updates are performed using Adam (Kingma, ), with the learning rate selected via grid search, while keeping all other parameters at their default values. The unrolling method was trained for $13000$ epochs, taking approximately $41$ hours in total. As shown in Figure 4, the performance of the learned parameters using the unrolled method is similar to that learned with the greedy method over the first $20$ iterations. However, during the $41$-hour training period for the unrolling approach, the greedy approach (9) could learn parameters for at least $1864$ iterations using NAG (with a stopping iteration of $\ell_{\text{stop}}=500$ for $t>50$), showing the benefit of solving a sequence of convex optimization problems.

Figure: Figure 4: Performance of the learned unrolled algorithm versus the greedy learned algorithm on test data.
Refer to caption: extracted/5992054/blur_images/unrol_vs_conv.png

### 6.2 Computed Tomography

In this case, the forward operator $A$ in (17) is the Radon transform in 2D. We simulate CT measurements using ODL (Adler et al., ) with a parallel-beam geometry employing $90$ projection angles evenly distributed over a $180$-degree range. For the dataset, we extract $40\times 40$ pixel crops from the center of each ground-truth image in the SARS-CoV-2 CT-scan dataset (Soares et al., ). The noise $\varepsilon$ is modeled with a standard deviation of $10^{-2}$, and we set $\alpha=10^{-4}$, resulting in $L=1.08$. The initial point $x_{0}$ is taken equal to $0\in\mathcal{X}$.

Greedy training was performed up to iteration $T=200$ with $\lambda_{t}=0$ for all iterations $t$ for the (PS), (PP), and (PC) parametrizations. The total time for training (PS) was about $10$ minutes, for (PP) was about $67$ minutes, and (PC) took approximately $10$ hours. For the (PF) parametrization, training was performed up to iteration $T=11$ with $\lambda_{t}=0$ for al $t$. Furthermore, the (PF) parametrization was trained with regularization such that $\lambda_{t}=10^{-10}$ for $t<T=101$ iterations. At iteration $101$, the learned operator $G_{\theta_{T}}$ satisfied $\|G_{\theta_{T}}-\tau I\|<\tau$, guaranteeing convergence on iterations $t\geq T$. For each iteration $t$, solving the optimization problem (9) with the (PF) parametrization took one hour.

Figure 5 shows that the learned scalar values again eventually fluctuate above and below $2/L$, similar to the behavior observed for the deblurring problem. The learned pointwise operators also exhibit oscillations between consecutive iterations, with many values falling outside the interval $(0,2/L)$. Likewise, the learned convolutional kernels for the CT problem contain both positive and negative values, are predominantly weighted toward the center of the kernel, and become increasingly similar as the iterations progress.

Figure: (a)
Refer to caption: extracted/5992054/CT_images/learned_scalars.png

Figure 6 shows that the learned parametrizations (PS), (PP), and (PC) generalize well to unseen data for the CT problem. In contrast, the full parametrization (PF) shows signs of overfitting, as it does not generalize well. The (PF) parametrization with regularization mitigates this issue, as the generalization performance improves. Similar to the deblurring problem, the learned (PC) parametrization outperforms NAG and L-BFGS on the CT test data, reaching a tolerance of $10^{-10}$ in approximately $30$ iterations, compared with about $80$ and more than $100$ for L-BFGS and NAG, respectively. The (PF) parametrization with regularization initially converges quickly but its speed decreases later due to regularization. This is because, with increasing iterations, the learned update gets closer to gradient descent.

Figure: Figure 6: Left: Train versus test set performance of the learned parametrizations. Right: Test performance versus benchmark optimization algorithms.
Refer to caption: extracted/5992054/CT_images/train_test_all.png

Figure 7 demonstrates that the learned parametrization (PC) achieves high-quality reconstructions with significantly fewer iterations compared to NAG for the CT problem. Also, note that the regularized (PF) parametrization achieves a good visual reconstruction after only two iterations while also providing guaranteed convergence.

Figure: Figure 7: Left: An example reconstruction for the CT problem. Right: The absolute difference between the final reconstruction and the intermediate reconstruction for the full parametrization with regularization at iteration $2$, the convolutional parametrization at iteration $10$, and for NAG at iterations $10$ and $70$.
Refer to caption: extracted/5992054/CT_images/reconstr.png

## 7 Conclusions

In this paper, we presented a novel L2O approach for minimizing unconstrained convex problems with differentiable objective functions. Our method employs a greedy strategy to learn a linear operator at each iteration of an optimization algorithm, meaning that memory requirements are constant with respect to the iterations. Parameter learning in our framework also corresponds to solving convex optimization problems, enabling the use of fast algorithms. Both factors allow training over a large number of training iterations, which would otherwise be prohibitively expensive. Furthermore, we obtain convergence results on the training set, even when the preconditioner is neither symmetric nor positive definite. Convergence is also proven for a class of unseen functions under certain conditions. The numerical results on imaging inverse problems demonstrate that our approach with a novel convolutional parametrization outperforms classical optimization algorithms.

## Appendix A Further Notation

The following notation is required in this section.

A function $f:\mathcal{X}\rightarrow\mathbb{R}$ is strongly convex with parameter $\mu>0$ if $f-\frac{\mu}{2}\|\cdot\|^{2}$ is convex. We say $f\in\mathcal{F}_{L,\mu}$ if $f\in\mathcal{F}_{L}$ and $f$ is $\mu$-strongly convex.

If $\operatorname{dim}(\mathcal{X})=n$ and $\operatorname{dim}(\mathcal{Y})=m$, denote an orthonormal basis of $\mathcal{X}$ by $\{e_{1},\cdots,e_{n}\}$, and an orthonormal basis of $\mathcal{Y}$ by $\{\tilde{e}_{1},\cdots,\tilde{e}_{n}\}$, then $A$ can be uniquely determined by $mn$ scalars $\gamma_{ij}$ for $i\in\{1,\cdots,m\}$, $j\in\{1,\cdots,n\}$: $A(e_{i})=\sum_{j=1}^{n}\gamma_{ij}\tilde{e}_{j}$, and denote $[A]_{ij}=\gamma_{ij}$.

Denote $\mathbf{1}\in\mathcal{X}$ to be such that $[\mathbf{1}]_{j}=1$ for $j\in\{1,\cdots,n\}$. For operators $A,B\in\mathcal{L}(\mathcal{X})$, and elements $x,y,z\in\mathcal{X}$, define the linear operators $A\odot B$ and $x\otimes y$ by

$$ $\displaystyle[A\odot B]_{ij}$ $\displaystyle:=[A]_{ij}[B]_{ij},$ (19) $\displaystyle(x\otimes y)z$ $\displaystyle:=\langle y,z\rangle x,$ (20) $$

with the property that

$$ $[x\otimes y]_{qi}=\langle y,e_{i}\rangle\langle x,e_{q}\rangle=[x]_{q}[y]_{i}.$ (21) $$

For two linear operators $A,B\in\mathcal{L}(\mathcal{X})$, define $A\otimes B$ by

$$ $[A\otimes B]_{ij,kl}=[A]_{ik}[B]_{jl}.$ (22) $$

For a linear operator $A\in\mathcal{L}(\mathcal{X})$ and $x\in\mathcal{X}$,

$$ $\displaystyle[Ax]_{i}=\sum_{j=1}^{n}[A]_{ij}[x]_{j}.$ (23) $$

## Appendix B Proofs for section 4

The following lemma is required to prove the convergence of our learned method.

###### Lemma 3 .

Define $F:\mathcal{X}^{N}\to\mathbb{R}$ by

$$ $F(x)=\frac{1}{N}\sum_{k=1}^{N}f_{k}(x_{k}),\quad x=(x_{1},x_{2},\dots,x_{N}) \in\mathcal{X}^{N}.$ (24) $$

Then

- 1.
Each $f_{k}\in\mathcal{F}_{L_{k}}$ implies $F\in\mathcal{F}_{L_{F}}$, with $L_{F}=\frac{1}{N}L_{\text{train}}$ where $L_{\text{train}}=\max\{L_{1},\cdots,L_{N}\}$.
- 2.
Each $f_{k}\in\mathcal{F}_{L_{k},\mu_{k}}$ implies
$F\in\mathcal{F}_{L,\mu_{F}}$ with $\mu_{F}=\frac{1}{N}\mu_{\text{min}}$ where $\mu_{\text{min}}=\min\{\mu_{1},\cdots,\mu_{N}\}$.

###### Proof.

We have

$$ $\nabla F(x)=\frac{1}{N}\left(\nabla f_{1}(x_{1}),\cdots,\nabla f_{N}(x_{N}) \right),$ (25) $$

and for any $y\in\mathcal{X}^{N}$,

$$ $\displaystyle\|x-y\|=\sqrt{\sum_{k=1}^{N}\|x_{k}-y_{k}\|^{2}}.$ $$

Then

$$ $\displaystyle\|\nabla F(x)-\nabla F(y)\|$ $\displaystyle=\frac{1}{N}\sqrt{\sum_{k=1}^{N}\left\|\nabla f_{k}(x_{k})-\nabla f _{k}(y_{k})\right\|^{2}}$ $\displaystyle\leq\frac{1}{N}\sqrt{\sum_{k=1}^{N}L_{k}^{2}\left\|x_{k}-y_{k} \right\|^{2}}\quad\text{($L_{k}$-smoothness of $f_{k}$.)}$ $\displaystyle\leq\frac{\max\{L_{1},\cdots,L_{N}\}}{N}\|x-y\|,$ $$

which proves $1.$
For strong convexity, it is required to show that $x\mapsto(F(x)-\frac{\min\{\mu_{1},\cdots,\mu_{N}\}}{N}\|x\|^{2})$ is convex.
We have

$$ $F(x)-\frac{\min\{\mu_{1},\cdots,\mu_{N}\}}{N}\|x\|^{2}=\frac{1}{N}\sum_{k=1}^{ N}(f_{k}(x_{k})-{\min\{\mu_{1},\cdots,\mu_{N}\}}\|x_{k}\|^{2}).$ (26) $$

Notice that due to the strong convexity of $f_{k}$ for all $k$, and that $\mu_{k}\geq\min\{\mu_{1},\cdots,\mu_{N}\}$,

$$ $x_{k}\mapsto(f_{k}(x_{k})-\min\{\mu_{1},\cdots,\mu_{N}\}\|x_{k}\|^{2})$ (27) $$

is convex.
Therefore the function $x\mapsto(F(x)-\frac{\min\{\mu_{1},\cdots,\mu_{N}\}}{N}\|x\|^{2})$ is convex as it is the sum of convex functions, as required.
∎

Theorem 1.
Convergence on training data.
Suppose ${\lambda_{t}}\geq 0$. If $\theta_{t}$ is $BGD$ then with Algorithm (1), we have

$$ $\nabla f_{k}(x_{k}^{t})\to 0\text{ as }t\to\infty,$ (28) $$

for all $k\in\{1,\cdots,N\}$.
Bonus: Convergence rates
Furthermore, if we denote

$$ $x_{0}=(x_{1}^{0},\cdots,x_{N}^{0}),\quad x^{*}=(x_{1}^{*},\cdots,x_{N}^{*}),$ (29) $$

then

$$ $F(x_{t})-F(x^{*})\leq\frac{\max\{L_{1},\cdots,L_{N}\}}{2tN}\|x_{0}-x^{*}\|^{2}.$ (30) $$

If, in addition, each $f_{k}$ is $\mu_{k}$-strongly convex, then we have linear convergence given by

$$ $F(x_{t})-F(x^{*})\leq\left(1-\frac{\max\{L_{1},\cdots,L_{N}\}}{\min\{\mu_{1}, \cdots,\mu_{N}\}}\right)^{t}(F(x_{0})-F(x^{*})).$ (31) $$

Note that this result gives a worst-case convergence bound among train functions. However, provable convergence is still acquired. Also, note that this is not an issue for a function class with constant smoothness and strongly convex parameters.

###### Proof.

As $\theta_{t}$ is $BGD$, we have that

$$ $\displaystyle F(x_{t+1})=g_{t,{\lambda_{t}}}(\theta_{t})$ $\displaystyle\leq g_{t,{\lambda_{t}}}(\tilde{\theta})$ $\displaystyle=\frac{1}{N}\sum_{k=1}^{N}f_{k}\left(x_{k}^{t}-\tau\nabla f(x_{k} ^{t})\right)$ $\displaystyle=F\left(x_{t}-\tau\left(\nabla f_{1}(x_{1}^{t}),\cdots,\nabla f_{ N}(x_{N}^{t})\right)\right)$ $\displaystyle=F\left(x_{t}-\tau N\nabla F(x_{t})\right)$ $\displaystyle=F\left(x_{t}-\tau_{F}\nabla F(x_{t})\right),$ $$

where $\tau_{F}=\frac{1}{L_{F}}$.

$F$ is $L_{F}$-smooth as each $f_{k}$ is $L_{k}$-smooth and $\mu$-strongly convex if each $f_{k}$ is $\mu_{k}$-strongly convex, where

$$ $\displaystyle L_{F}$ $\displaystyle=\frac{\max\{L_{1},\cdots,L_{N}\}}{N}$ $\displaystyle\mu_{F}$ $\displaystyle=\frac{\min\{\mu_{1},\cdots,\mu_{N}\}}{N},$ $$

and therefore, using standard convergence rate results of gradient descent (Nesterov et al., ), we have

$$ $F(x_{t})-F(x^{*})\leq\frac{L_{F}}{2t}\|x_{0}-x^{*}\|^{2},$ (32) $$

as $F$ is $L_{F}-$smooth, and if $F$ is also $\mu_{F}$-strongly convex we have

$$ $F(x_{t})-F(x^{*})\leq\left(1-\frac{L_{F}}{\mu_{F}}\right)^{t}(F(x_{0})-F(x^{*} )).$ (33) $$

In both cases, we have that $\|\nabla F(x_{t})\|^{2}=\frac{1}{N^{2}}\sum_{k=1}^{N}\|\nabla f_{k}(x_{k}^{t})
\|^{2}\to 0$ as $t\to\infty$, which implies that $\nabla f_{k}(x_{k}^{t})\to 0$ as $t\to\infty$ for all $k\in\{1,\cdots,N\}$.
∎

We have proved convergence for the mean of our train functions. The following proposition proves the same convergence rate holds for each function in our training set.

###### Proposition 4 .

Suppose we have a convergence rate for $F$ of

$$ $F(x_{t})-F^{*}\leq C\rho(t),$ (34) $$

for some constant $C>0$. Then the convergence rate for all $f_{k}\in\mathcal{T}$ is given by

$$ $f_{k}(x_{k}^{t})-f_{k}^{*}\leq c\rho(t),$ (35) $$

for some constant $c>0$.

###### Proof.

Let $k\in\{1,\cdots,N\}$. Note that by the definition of $F$, we may write

$$ $\displaystyle F(x_{t})$ $\displaystyle=\frac{1}{N}f_{k}(x_{k}^{t})+\frac{1}{N}\sum_{i=1,i\neq k}^{N}f_{ i}(x_{i}^{t}),$ $\displaystyle\implies F^{*}$ $\displaystyle=\frac{1}{N}f_{k}^{*}+\frac{1}{N}\sum_{i=1,i\neq k}^{N}f_{i}^{*}.$ $$

Therefore, using the convergence rate of $F$ gives

$$ $\frac{1}{N}f_{k}(x_{k}^{t})-\frac{1}{N}f_{k}^{*}+\frac{1}{N}\sum_{i=1,i\neq k} ^{N}(f_{i}(x_{i}^{t})-f_{i}^{*})\leq C\rho(t)$ (36) $$

which implies

$$ $\displaystyle f_{k}(x_{k}^{t})-f_{k}^{*}+\sum_{i=1,i\neq k}^{N}(f_{i}(x_{i}^{t })-f_{i}^{*})\leq NC\rho(t),$ $\displaystyle\implies$ $\displaystyle f_{k}(x_{k}^{t})-f_{k}^{*}\leq NC\rho(t),$ $$

as $f_{i}(x_{i}^{t})\geq f_{i}^{*}0$ for all $i$.
∎

###### Proposition 5 .

Assume that $f:\mathcal{X}\to\mathbb{R}$ is convex, continuously differentiable, and has a global minimum. Then for a point $z\in\mathcal{X}$ if there exists some $x^{*}\in\arg\min_{x}f(x)$ such that $[z]_{i}=[x^{*}]_{i}$, then $[\nabla f(x)]_{i}=0$.

###### Proof.

Let $g:\mathbb{R}\to\mathbb{R}$ be defined by $g(t):=f(x+te_{i})$, then $g$ is convex as for $\alpha\in[0,1],t_{1},t_{2}\in\mathbb{R}$, we have $g(\alpha t_{1}+(1-\alpha)t_{2})=f(\alpha(x+t_{1}e_{i})+(1-\alpha)(x+t_{2}e_{i}
))\leq\alpha g(t_{1})+(1-\alpha)g(t_{2})$. Note that $g^{\prime}(0)=[\nabla f(x)]_{i}$. Assume there exists $\delta\neq 0$ such that $g(\delta)<g(0)$, then $g(\delta)<g(0)=f(x)$, which is a contradiction of $x^{*}$ being an optimal point, as one can take $x=x^{*}$. Therefore $g$ achieves a minimum at $t=0$, then $[\nabla f(x)]_{i}=0$.
∎

###### Proof of Proposition 1 . .

Choose the vector $p\in\mathcal{X}$ such that

$$ $[p]_{i}=\begin{cases}\frac{[x_{0}-x_{*}]_{i}}{[\nabla f(x_{0})]_{i}},\text{ if }[\nabla f(x_{0})]_{i}\neq 0,\\ 0,\text{ otherwise},\end{cases}$ (37) $$

and let $I=\{i:[\nabla f(x_{0})]_{i}\neq 0\}$
Then for any $i\in I$, we have

$$ $\displaystyle[x_{0}-p\odot\nabla f(x_{0})]_{i}$ $\displaystyle=[x_{0}]_{i}-[p]_{i}[\nabla f(x_{0})]_{i}$ $\displaystyle=[x_{0}]_{i}-\frac{[x_{0}-x^{*}]_{i}}{[\nabla f_{k}(x_{0})]_{i}}[ \nabla f(x_{0})]_{i}$ $\displaystyle=[x^{*}]_{i}.$ $$

Thus, by proposition 5, $[\nabla f(x_{0}-p\odot\nabla f(x_{0}))]_{i}=0$, for all $i\in I$, and similarly $[\nabla f(x_{0}-p\odot\nabla f(x_{0}))]_{i}=0$, for all $i\notin I$, and therefore $\nabla f(x_{0}-p\odot\nabla f(x_{0}))=0$, meaning that $x_{0}-p\odot\nabla f(x_{0})\in\arg\min_{x}f(x)$ as required.
∎

###### Proof of Proposition 2 . .

We require

$$ $\displaystyle\begin{cases}x_{1}^{*}&=x_{1}^{0}-P\nabla f_{1}(x_{1}^{0}),\\ &\vdots\\ x_{N}^{*}&=x_{N}^{0}-P\nabla f_{N}(x_{N}^{0}).\end{cases}$ $$

Each of these equations gives $n$ linear equations in $n^{2}$ unknowns. There are $N$ such equations and so we have $nN$ linear equations in $n^{2}$ unknowns. Rewritten, these read

$$ $P\begin{bmatrix}\nabla f_{1}(x_{1}^{0})|\cdots|\nabla f_{N}(x_{N}^{0})\end{ bmatrix}=\begin{bmatrix}x_{1}^{0}-x_{1}^{*}|\cdots|x_{N}^{0}-x_{N}^{*}\end{ bmatrix}.$ (38) $$

Such a $P$ exists if $nN\leq n^{2}$, which is equivalent to $N\leq n$, and if the columns of $\begin{bmatrix}\nabla f_{1}(x_{1}^{0})|\cdots|\nabla f_{N}(x_{N}^{0})\end{bmatrix}$ are linearly independent.
∎

###### Lemma 3 .

###### Proof.

###### Proof.

###### Proposition 4 .

###### Proof.

###### Proposition 5 .

###### Proof.

###### Proof of Proposition 1 . .

###### Proof of Proposition 2 . .

### B.1 Unseen data

###### Proof of Lemma 1 . .

Firstly,

$$ $g_{t}(\tilde{\theta},{\lambda_{t}})=\frac{1}{N}\sum_{k=1}^{N}f_{k}\left(x_{k}^ {t}-\tau\nabla f_{k}(x_{k}^{t})\right)\to\frac{1}{N}\sum_{k=1}^{N}f_{k}^{*} \text{ as }t\to\infty.$ (39) $$

Note also that as $(\theta_{t})_{t=0}^{\infty}$ is a $BGD$ sequence of parameters then $\frac{1}{N}\sum_{k=1}^{N}f_{k}^{*}\leq g_{t}({\theta_{t},{\lambda_{t}}})\leq g
_{t}({\tilde{\theta},{\lambda_{t}}})$ and so $g_{t}({\theta_{t},{\lambda_{t}}})\to\frac{1}{N}\sum_{k=1}^{N}f_{k}^{*}$ as $t\to\infty$ as $g_{t}({\theta_{t},{\lambda_{t}}})\geq\frac{1}{N}\sum_{k=1}^{N}f_{k}^{*}$. Furthermore,

$$ $\displaystyle g_{t}(\tilde{\theta},{\lambda_{t}})-g_{t}(\theta_{t},{\lambda_{t }})$ $\displaystyle=-\frac{\lambda_{t}}{2}\|\theta_{t}-\tilde{\theta}\|_{2}^{2}+ \frac{1}{N}\sum_{k=1}^{N}f_{k}\left(x_{k}^{t}-\tau\nabla f_{k}(x_{k}^{t})\right)$ $\displaystyle-\frac{1}{N}\sum_{k=1}^{N}f_{k}(x_{k}^{t}-G_{\theta_{t}}\nabla f_ {k}(x_{k}^{t}))).$ $$

Therefore,

$$ $0=\lim_{t\to\infty}g_{t}(\tilde{\theta},{\lambda_{t}})-g_{t}(\theta_{t},{ \lambda_{t}})=\lim_{t\to\infty}-\frac{\lambda_{t}}{2}\|\theta_{t}-\tilde{ \theta}\|_{2}^{2}.$ (40) $$

Now, $\lim\inf_{t\to\infty}\lambda_{t}>0$ implies that

$$ $\frac{\lambda}{2}\|\theta_{t}-\tilde{\theta}\|_{2}^{2}\to 0\text{ as }t\to\infty.$ (41) $$

In particular, $\theta_{t}\to\tilde{\theta}$ as $t\to\infty$, as required.
∎

###### Lemma 4 .

Suppose $\lim\inf_{t\to\infty}\lambda_{t}>0$, $G:\Theta\to\mathcal{L}(\mathcal{X})$ is continuous and at each training iteration $\theta_{t}$ is $BGT$. Then for any $\nu$ such that $0\leq\nu<\tau$, there exists an iteration $T$ such that

$$ $\|G_{\theta_{T}}-\tau I\|\leq\nu<\tau.$ (42) $$

###### Proof.

By Lemma 1, $\theta_{t}\to\tilde{\theta}$ as $t\to\infty$, therefore as $G_{\theta}$ is continuous in $\theta$ we have $G_{\theta_{t}}\to\tau I$ as $t\to\infty$. Therefore for any $\nu>0$ there exists some iteration $T>0$ such that

$$ $\|G_{\theta_{t}}-\tau I\|\leq\nu$ (43) $$

for all $t\geq T$, in particular for $t=T$.
∎

###### Proof of Theorem 2 . .

By Lemma 4, for any tolerance $\nu<\tau$ there exists an iteration $T$ such that

$$ $\left\|G_{\theta_{T}}-\tau I\right\|\leq\nu.$ (44) $$

Let

$$ $G_{\theta_{T}}=\tau I+M,$ (45) $$

then

$$ $\|M\|\leq\nu.$ (46) $$

Using $L$-smoothness of $f$, we have

$$ $\displaystyle f(x-G_{\theta_{T}}\nabla f(x))$ $\displaystyle\leq f(x)-\left\langle G_{\theta_{T}}\nabla f(x),\nabla f(x) \right\rangle+\frac{L}{2}\|G_{\theta_{T}}\nabla f(x)\|^{2}$ $\displaystyle=f(x)-\left\langle G_{\theta_{T}}\nabla f(x),\nabla f(x)-\frac{L} {2}G_{\theta_{T}}\nabla f(x)\right\rangle$ $\displaystyle=f(x)-\left\langle\left(\tau I+M\right)\nabla f(x),\nabla f(x)- \frac{L}{2}\left(\tau I+M\right)\nabla f(x)\right\rangle$ $\displaystyle=f(x)-\left\langle\tau\nabla f(x)+M\nabla f(x),\nabla f(x)-\frac{ L}{2}\tau\nabla f(x)-\frac{L}{2}M\nabla f(x)\right\rangle$ $\displaystyle=f(x)-\left\langle\tau\nabla f(x)+M\nabla f(x),\left(1-\frac{L \tau}{2}\right)\nabla f(x)-\frac{L}{2}M\nabla f(x)\right\rangle$ $\displaystyle=f(x)-\tau\left(1-\frac{L\tau}{2}\right)\|\nabla f(x)\|^{2}+\frac {L}{2}\left\|M\nabla f(x)\right\|^{2}$ $\displaystyle-\left(1-\tau L\right)\left\langle\nabla f(x),M\nabla f(x)\right\rangle$ $\displaystyle\leq f(x)-\tau\left(1-\frac{L\tau}{2}\right)\|\nabla f(x)\|^{2}+ \frac{L\nu^{2}}{2}\left\|\nabla f(x)\right\|^{2}+\nu\left|1-\tau L\right|\left \|\nabla f(x)\right\|^{2}$ $\displaystyle=f(x)-\left(\tau\left(1-\frac{\tau L}{2}\right)-\frac{L\nu^{2}}{2 }-\nu\left|1-\tau L\right|\right)\left\|\nabla f(x)\right\|^{2}$ $\displaystyle=f(x)-c(\nu,L,\tau)\left\|\nabla f(x)\right\|^{2},$ $$

where

$$ $\displaystyle c(\nu,L,\tau)$ $\displaystyle=\tau\left(1-\frac{\tau L}{2}\right)-\frac{L\nu^{2}}{2}-\nu\left| 1-\tau L\right|$ $\displaystyle=\frac{L}{2}\left(\frac{1}{L}-\nu-\left|\frac{1}{L}-\tau\right| \right)\left(\nu+\left|\frac{1}{L}-\tau\right|+\frac{1}{L}\right).$ $$

Therefore

$$ $f(x_{t+1})\leq f(x_{t})-c(\nu,L,\tau)\left\|\nabla f(x)\right\|^{2}.$ (47) $$

Note that $\nu+\left|\frac{1}{L}-\tau\right|+\frac{1}{L}>0$, so for $c(\nu,L,\tau)$ to be positive, we require

$$ $\frac{1}{L}-\nu-\left|\frac{1}{L}-\tau\right|>0.$ (48) $$

Case 1 - $\frac{1}{L}\geq\tau$
Then we require

$$ $\tau-\nu>0\iff\nu<\tau,$ (49) $$

which is true by assumption.
Case 2 - $\frac{1}{L}<\tau$
Then we require

$$ $\frac{2}{L}-\nu-\tau>0\iff\frac{1}{L}>\frac{\tau+\nu}{2}.$ (50) $$

To conclude both cases, we have $\nu<\tau$ and therefore as $\frac{1}{\tau}<\frac{1}{\tau+\nu}$, we require only case $2$ to be satisfied for $c(\nu,L,\tau)>0$:

$$ $L<\frac{2}{\tau+\nu}$ (51) $$

∎

###### Proposition 6 .

Assume that $G:\Theta\rightarrow\mathcal{L}(\mathcal{X})$ is continuous. Then at any iteration $t$ there exists $\lambda_{t}\geq 0$ and a constant $\tilde{L}>0$ such that for all $f\in\mathcal{F}_{\tilde{L}}$ and any starting point $x_{0}$, using Algorithm 2 gives $\nabla f\left(x_{t}\right)\rightarrow 0$ as $t\rightarrow\infty$.

###### Proof.

Take $\nu\in(0,\tau)$. Define $h(\lambda)=\|G_{\arg\min_{\theta}g_{t,\lambda}(\theta)}-\tau I\|-\nu$ for $g_{t,\lambda}(\theta)$ as in (9). Note that $\lim_{\lambda\to\infty}h(\lambda)=-\nu<0$. If $h(0)<0$ then we are done as for $\lambda_{t}=0$, the corresponding learned parameters $\theta_{t}$ satisfy $\|G_{\theta_{t}}-\tau I\|<\nu$, leading to a provably convergent algorithm for $f\in\mathcal{F}_{\tilde{L}}$ for some $\tilde{L}>0$. Else, suppose that $h(0)>0$. Then as $h$ is continuous in $\lambda$, there exists some $\lambda$ such that $h(\lambda)<0$.
∎

At iteration $t$, to find a $\lambda_{t}$ that is large enough to ensure convergence, we start at an initial point $\lambda=10^{-6}$ and find $\phi\in\arg\min_{\theta}g_{t,\lambda}(\theta)$. If $\|G_{\phi}-\tau I\|<\tau$, then increase $\lambda$ by a multiple and re-evalute. Repeat until this inequality no longer holds, and take $\lambda_{t}$ to be the most recent $\lambda$ such that $\|G_{\phi}-\tau I\|<\tau$. Else if $\lambda=10^{-6}$ and $\phi\in\arg\min_{\theta}g_{t,\lambda}(\theta)$ satisfies $\|G_{\phi}-\tau I\|>\tau$ then reduce $\lambda$ by a multiple and re-evaluate until $\|G_{\phi}-\tau I\|<\tau$, then take $\lambda_{t}=\lambda$. For the (PS) parametrization we take the multiple to be $5$, and for the (PP) parametrization, we take this multiple to be $2$.

###### Proof of Theorem 3 . .

Define $D=\max_{t=0,1,\cdots}\{\|x_{t}-x^{*}\|\}$, which is finite as $(x_{t})$ is bounded. Due to the convexity of $f$ and the Cauchy-Schwarz inequality, we have that

$$ $\displaystyle f(x_{t})-f(x^{*})$ $\displaystyle\leq\langle\nabla f(x_{t}),x_{t}-x^{*}\rangle$ $\displaystyle\leq\|\nabla f(x_{t})\|\|x_{t}-x^{*}\|$ $\displaystyle\leq D\|\nabla f(x_{t})\|.$ $$

Therefore

$$ $\|\nabla f(x_{t})\|^{2}\geq\frac{1}{D^{2}}(f(x_{t})-f(x^{*}))^{2},$ (52) $$

and for $t\geq T$ we have

$$ $\displaystyle f(x_{t+1})$ $\displaystyle\leq f(x_{t})-c(\nu,L,\tau)\|\nabla f(x_{t})\|^{2}$ $\displaystyle\leq f(x_{t})-\frac{c(\nu,L,\tau)}{D^{2}}(f(x_{t})-f(x^{*}))^{2}.$ $$

Denote $\Delta_{t}=f(x_{t+T})-f(x^{*})$, then in the spirit of (Nesterov et al., ), we have for all $t\geq 0$

$$ $\displaystyle\Delta_{t+1}\leq\Delta_{t}-\frac{c}{D^{2}}\Delta_{t}^{2}$ $\displaystyle\implies\frac{1}{\Delta_{t}}\leq\frac{1}{\Delta_{t+1}}-\frac{c}{D ^{2}}\frac{\Delta_{t}}{\Delta_{t+1}}\leq\frac{1}{\Delta_{t+1}}-\frac{c}{D^{2}}$ $\displaystyle\implies\frac{c}{D^{2}}+\frac{1}{\Delta_{t}}\leq\frac{1}{\Delta_{ t+1}}.$ $$

Taking a summation gives

$$ $\displaystyle\sum_{k=0}^{t-1}\frac{c}{D^{2}}\leq\sum_{k=0}^{t-1}\left(\frac{1} {\Delta_{k+1}}-\frac{1}{\Delta_{k}}\right)$ $\displaystyle\implies\frac{c}{D^{2}}t\leq\frac{1}{\Delta_{t}}-\frac{1}{\Delta_ {0}}.$ $$

Therefore

$$ $\displaystyle\Delta_{t}\leq\frac{1}{\frac{1}{\Delta_{0}}+\frac{c}{D^{2}}t}= \frac{D^{2}\Delta_{0}}{D^{2}+c\Delta_{0}t}\leq\frac{D^{2}\Delta_{0}}{c\Delta_{ 0}t}=\frac{D^{2}/c}{t},$ $$

as required.

∎

###### Proof of Lemma 1 . .

###### Lemma 4 .

###### Proof.

###### Proof of Theorem 2 . .

###### Proposition 6 .

###### Proof.

###### Proof of Theorem 3 . .

## Appendix C Proofs for section 5

###### Proof of Lemma 2 . .

- 1.
For scalar step sizes, $G_{\theta}=\theta I$, take $\tilde{\theta}=\tau$.
- 2.
For a pointwise parametrization, $G_{\theta}x=\theta\odot x$, take $\tilde{\theta}=\tau\mathbf{1}$.
- 3.
For full operator parametrization, $G_{\theta}=\theta\in\mathcal{L}(\mathcal{X})$, take $\tilde{\theta}=\tau I$.
- 4.
For the convolutional parametrization, $G_{\theta}x=\theta\ast x$, take
$\theta(i,j)=\begin{cases}\tau,\text{ if }i=j=0,\\
0,\text{ otherwise}.\end{cases}$
(53)

$G_{\theta}$ are clearly continuous in $\theta$ for all listed parametrizations.
∎

###### Proof of Corollary 1 . .

Apply lemma 2.
∎

###### Proof of Proposition 3 . .

Because this problem is convex, if a solution $\theta$ is found by differentiating the objective function and equating equal to zero, this is a global minimizer.
First, note that

$$ $\displaystyle f_{k}(x_{k}^{t}-B_{k}^{t}\theta)$ $\displaystyle=\frac{1}{2}\|A_{k}(x_{k}^{t}-B_{k}^{t}\theta)-y_{k}\|^{2}$ $\displaystyle=\frac{1}{2}\|A_{k}x_{k}^{t}-y_{k}\|^{2}+\frac{1}{2}\|-A_{k}B_{k} ^{t}\theta\|^{2}+\langle-A_{k}B_{k}^{t}\theta,A_{k}x_{k}^{t}-y_{k}\rangle$ $\displaystyle=\frac{1}{2}\|A_{k}x_{k}^{t}-y_{k}\|^{2}+\frac{1}{2}\|A_{k}B_{k}^ {t}\theta\|^{2}-\langle\theta,(B_{k}^{t})^{*}\nabla f_{k}(x_{k}^{t})\rangle.$ $$

Now,

$$ $\displaystyle\nabla_{\theta}\left\{\frac{1}{N}\sum_{k=1}^{N}f_{k}(x_{k}^{t}-B_ {k}^{t}\theta)+\frac{{\lambda_{t}}}{2}\|\theta-\tilde{\theta}\|_{2}^{2}\right\}$ $\displaystyle=\frac{1}{N}\sum_{k=1}^{N}(A_{k}B_{k}^{t})^{*}(A_{k}B_{k}^{t} \theta)-(B_{k}^{t})^{*}\nabla f_{k}(x_{k}^{t})+{\lambda_{t}}(\theta-\tilde{ \theta})$ $$

is equal to zero if and only if

$$ $\displaystyle\left(\frac{1}{N}\sum_{k=1}^{N}(A_{k}B_{k}^{t})^{*}(A_{k}B_{k}^{t })+{\lambda_{t}}I_{\Theta}\right)\theta={\lambda_{t}}\tilde{\theta}+\frac{1}{N }\sum_{k=1}^{N}(B_{k}^{t})^{*}\nabla f_{k}(x_{k}^{t}).$ $$

∎

A bonus proposition regarding the uniqueness of optimal parameters:

###### Proposition 7 .

$g_{t,{\lambda_{t}}}(\theta)$ has a unique global minimizer $\theta_{t}^{*}$ if at least one of the following are satisfied:

- •
${\lambda_{t}}>0$,
- •
$f_{k}$ is twice continuously differentiable for $k\in\{1,\cdots,N\}$, and there exists some $j\in\{1,\cdots,N\}$ for which both $B_{j}^{t}$ is injective and also $f_{j}$ is $\mu_{j}$-strongly convex.

###### Proof.

Case 1 - $\lambda_{t}>0$
$\frac{1}{N}\sum_{k=1}^{N}(A_{k}B_{k}^{t})^{*}(A_{k}B_{k}^{t})$ is self-adjoint and positive semi-definite as it is the sum of self-adjoint operators, $\frac{1}{N}\sum_{k=1}^{N}(A_{k}B_{k}^{t})^{*}(A_{k}B_{k}^{t})+{\lambda_{t}}I$ is a self-adjoint, positive-definite operator and therefore invertible.
Case 2 - $\lambda_{t}=0$
If each $f_{k}$ is twice continuously differentiable; then $g_{t,{\lambda_{t}}}$ is twice continuously differentiable. It is then sufficient to show there exists $m>0$ such that

$$ $\nabla^{2}g_{t,{\lambda_{t}}}(\theta)\succeq mI,$ (54) $$

for all $\theta$, as this implies that $g_{t,{\lambda_{t}}}$ is strongly convex and has a unique global minimizer.
Note that

$$ $\nabla^{2}g_{t,{\lambda_{t}}}(\theta)=\frac{1}{N}\sum_{k=1}^{N}(B_{k}^{t})^{*} \nabla^{2}f_{k}(x_{k}^{t}-B_{k}^{t}\theta)B_{k}^{t}.$ (55) $$

Note that

$$ $\displaystyle\langle v,\nabla^{2}g_{t,{\lambda_{t}}}(\theta)v\rangle$ $\displaystyle=\langle v,\frac{1}{N}\sum_{k=1}^{N}(B_{k}^{t})^{*}\nabla^{2}f_{k }(x_{k}^{t}-B_{k}^{t}\theta)B_{k}^{t}v\rangle$ (56) $\displaystyle=\frac{1}{N}\sum_{k=1}^{N}\langle v,(B_{k}^{t})^{*}\nabla^{2}f_{k }(x_{k}^{t}-B_{k}^{t}\theta)B_{k}^{t}v\rangle$ (57) $\displaystyle=\frac{1}{N}\sum_{k=1}^{N}\langle B_{k}^{t}v,\nabla^{2}f_{k}(x_{k }^{t}-B_{k}^{t}\theta)B_{k}^{t}v\rangle.$ (58) $$

Each $f_{k}$ is convex and so for all $v\in\mathcal{X}$,

$$ $\langle v,\nabla^{2}f_{k}(x_{k}^{t}-B_{k}^{t}\theta)v\rangle\geq 0,$ (60) $$

and $f_{j}$ is $\mu_{j}$-strongly convex, therefore

$$ $\langle v,\nabla^{2}f_{j}(v_{t}^{j}-B_{j}^{t}\theta)v\rangle\geq\mu_{j}\|v\|^{ 2}.$ (61) $$

For $v\in\mathcal{X}$,

$$ $\displaystyle\langle v,\nabla^{2}g_{t,{\lambda_{t}}}(\theta)v\rangle$ $\displaystyle\geq\frac{1}{N}\mu_{j}v^{T}(B_{j}^{t})^{*}B_{j}^{t}v$ $\displaystyle\geq\left(\frac{1}{N}\mu_{j}\rho_{\text{min}}^{j}\right)\|v\|^{2},$ $$

where $\rho_{\text{min}}^{j}$ is the minimum eigenvalue of $M_{j}^{t}=(B_{j}^{t})^{*}B_{j}^{t}$ (a symmetric linear operator). Due to the symmetry of $M_{j}^{t}$, $\rho_{\text{min}}^{j}\geq 0$ and is greater than zero if and only if $B_{j}^{t}$ is injective. As $B_{j}^{t}$ is injective, then $\rho_{\text{min}}^{j}>0$ and therefore $g_{t,{\lambda_{t}}}(\theta)$ is strongly-convex.
∎

Proposition 7 applied to least-square functions:

###### Corollary 2 .

Uniqueness of optimal parameters in the least-squares case
When our $f_{k}$ can be written as least-squares functions (15), then $g_{t,{\lambda_{t}}}(\theta)$ has a unique global minimizer $\theta_{t}^{*}$ if at least one of the following are satisfied:

- •
${\lambda_{t}}>0$,
- •
there exists some $j\in\{1,\cdots,N\}$ for which both $B_{j}^{t}$ and $A_{j}$ are injective.

###### Proof.

If $A_{j}$ is injective then $A_{j}^{*}A_{j}$ is invertible which means that $f_{j}(x)=\frac{1}{2}\|A_{j}x-y^{j}\|^{2}$ is strongly convex.
∎

###### Proposition 8 .

$p_{t}$ given by

$$ $p_{t}=\bigg{(}{\lambda_{t}}I_{\Theta}+\frac{1}{N}\sum_{k=1}^{N}\left(\nabla f_ {k}(x_{k}^{t})\otimes\nabla f_{k}(x_{k}^{t})\right)\odot(A_{k}^{*}A_{k})\bigg{ )}^{\dagger}\bigg{(}{\lambda_{t}}\tilde{\theta}+\frac{1}{N}\sum_{k=1}^{N} \nabla f_{k}(x_{k}^{t})\odot\nabla f_{k}(x_{k}^{t})\bigg{)}$ (62) $$

is a solution to (9) with the pointwise parametrization $G_{p_{t}}x=p_{t}\odot x$ for any $x\in\mathcal{X}$.

###### Proof.

Define for $x\in\mathcal{X}$,

$$ $\displaystyle B_{k}^{t}x$ $\displaystyle=\nabla f_{k}(x_{k}^{t})\odot x$ $$

then, for $x\in\mathcal{X}$, $(B_{k}^{t})^{*}(x)=B_{k}^{t}(x)$. Now,

$$ $\displaystyle(A_{k}B_{k}^{t})^{*}(A_{k}B_{k}^{t})p$ $\displaystyle=(B_{k}^{t})^{*}\left(A_{k}^{*}A_{k}B_{k}^{t}p\right)$ $\displaystyle=\nabla f_{k}(x_{k}^{t})\odot\left(A_{k}^{*}A_{k}(\nabla f_{k}(x_ {k}^{t})\odot p)\right).$ $$

Now,

$$ $\displaystyle[\nabla f_{k}(x_{k}^{t})\odot\left(A_{k}^{*}A_{k}(\nabla f_{k}(x_ {k}^{t})\odot p)\right)]_{j}$ $\displaystyle=[\nabla f_{k}(x_{k}^{t})]_{j}[A_{k}^{*}A_{k}(\nabla f_{k}(x_{k}^ {t})\odot p)]_{j},\text{ by }\eqref{eq::xodoty-identity}$ $\displaystyle=[\nabla f_{k}(x_{k}^{t})]_{j}\sum_{i=1}^{n}[\nabla f_{k}(x_{k}^{ t})\odot p]_{i}[A_{k}^{*}A_{k}]_{ji},\text{ by }\eqref{eq::Ax-identity}$ $\displaystyle=\sum_{i=1}^{n}[\nabla f_{k}(x_{k}^{t})]_{i}[p]_{i}[\nabla f_{k}( x_{k}^{t})]_{j}[A_{k}^{*}A_{k}]_{ji},\text{ by }\eqref{eq::xodoty-identity}.$ $$

Secondly,

$$ $\displaystyle\left[\left(\left(\nabla f_{k}(x_{k}^{t})\otimes\nabla f_{k}(x_{k }^{t})\right)\odot(A_{k}^{*}A_{k})\right)p\right]_{j}$ $\displaystyle=\sum_{i=1}^{n}[p]_{i}[\left(\nabla f_{k}(x_{k}^{t})\otimes\nabla f _{k}(x_{k}^{t})\right)\odot(A_{k}^{*}A_{k})]_{ji},\text{ by }\eqref{eq::Ax-identity}$ $\displaystyle=\sum_{i=1}^{n}[p]_{i}[\nabla f_{k}(x_{k}^{t})\otimes\nabla f_{k} (x_{k}^{t})]_{ji}[A_{k}^{*}A_{k}]_{ji},\text{ by }\eqref{eq::AodotB-identity}$ $\displaystyle=\sum_{i=1}^{n}[p]_{i}[\nabla f_{k}(x_{k}^{t})]_{j}[\nabla f_{k}( x_{k}^{t})]_{i}[A_{k}^{*}A_{k}]_{ji}\text{ by }\eqref{eq::AotimesB-identity}$ $\displaystyle=[\nabla f_{k}(x_{k}^{t})\odot\left(A_{k}^{*}A_{k}(\nabla f_{k}(x _{k}^{t})\odot p)\right)]_{j},\text{ by }\eqref{eq::xotimesy_iq}.$ $$

Finally,

$$ $\displaystyle{\lambda_{t}}\tilde{\theta}+\frac{1}{N}\sum_{k=1}^{N}(B_{k}^{t})^ {*}\nabla f_{k}(x_{k}^{t})$ $\displaystyle={\lambda_{t}}\tilde{\theta}+\frac{1}{N}\sum_{k=1}^{N}\nabla f_{k }(x_{k}^{t})\odot\nabla f_{k}(x_{k}^{t}).$ $$

Then the result follows from proposition 3.
∎

###### Proposition 9 .

Let $B_{k}^{t}:\mathcal{L}(\mathcal{X})\to\mathcal{X}$ be such that for any linear operator $P\in\mathcal{L}(\mathcal{X})$, we have $B_{k}^{t}(P)=P\nabla f_{k}(x_{k}^{t})$. Then its adjoint $(B_{k}^{t})^{*}:\mathcal{X}\to\mathcal{L}(\mathcal{X})$ is given by

$$ $(B_{k}^{t})^{*}(w)=\nabla f_{k}(x_{k}^{t})\otimes w,$ (63) $$

for any element $w\in\mathcal{X}$. Then

$$ $\theta_{t}=\left({\lambda_{t}}I_{\Theta}+\frac{1}{N}\sum_{k=1}^{N}(A_{k}^{*}A_ {k})\otimes(\nabla f_{k}(x_{k}^{t})\otimes\nabla f_{k}(x_{k}^{t}))\right)^{ \dagger}\left({\lambda_{t}}\tilde{\theta}+\frac{1}{N}\sum_{k=1}^{N}\nabla f_{k }(x_{k}^{t})\otimes\nabla f_{k}(x_{k}^{t})\right),$ (64) $$

(16) is a solution to (9) for the full operator parametrization.

###### Proof.

$\theta_{t}\in\mathcal{L}(\mathcal{X})$ and we require $\theta_{t}\nabla f_{k}(x_{k}^{t})=B_{k}^{t}(\theta_{t})$, so can take $B_{k}^{t}(\theta_{t})=\theta_{t}\nabla f_{k}(x_{k}^{t})$. For the adjoint,

$$ $\displaystyle\langle B_{k}^{t}(P),w\rangle$ $\displaystyle=\sum_{i=1}^{n}[P\nabla f_{k}(x_{k}^{t})]_{i}[w]_{i}$ (65) $\displaystyle=\sum_{i=1}^{n}\sum_{j=1}^{n}[P]_{ij}[\nabla f_{k}(x_{k}^{t})]_{j }[w]_{i}$ (66) $\displaystyle=\langle P,(B_{k}^{t})^{*}w\rangle$ (67) $\displaystyle=\sum_{i=1}^{n}\sum_{j=1}^{n}[P]_{ij}[(B_{k}^{t})^{*}w]_{ij},$ (68) $$

and therefore $[(B_{k}^{t})^{*}w]_{ij}=w_{i}[\nabla f_{k}(x_{k}^{t})]_{j}$, which means $(B_{k}^{t})^{*}(w)=w\otimes\nabla f_{k}(x_{k}^{t})$. Now,

$$ $\displaystyle[(A_{k}B_{k}^{t})^{*}(A_{k}B_{k}^{t})\theta]_{ij}$ $\displaystyle=[(B_{k}^{t})^{*}(A_{k}^{*}A_{k}B_{k}^{t}\theta)]_{ij}$ $\displaystyle=[(A_{k}^{*}A_{k}B_{k}^{t}\theta)\otimes\nabla f_{k}(x_{k}^{t})]_ {ij}$ $\displaystyle=[\nabla f_{k}(x_{k}^{t})]_{j}[A_{k}^{*}A_{k}B_{k}^{t}\theta]_{i} ,\text{ by }\eqref{eq::xotimesy_iq}$ $\displaystyle=[\nabla f_{k}(x_{k}^{t})]_{j}\sum_{q=1}^{n}[A_{k}^{*}A_{k}]_{iq} [B_{k}^{t}\theta]_{q}$ $\displaystyle=[\nabla f_{k}(x_{k}^{t})]_{j}\sum_{q=1}^{n}[A_{k}^{*}A_{k}]_{iq} \sum_{\ell=1}^{n}[\theta]_{q\ell}[\nabla f_{k}(x_{k}^{t})]_{\ell},(\text{ definition of }B_{k}^{t}).$ $$

Similarly,

$$ $\displaystyle[((A_{k}^{*}A_{k})\otimes(\nabla f_{k}(x_{k}^{t})\otimes\nabla f_ {k}(x_{k}^{t})))\theta]_{ij}$ $\displaystyle=\sum_{q,\ell=1}^{n}[(A_{k}^{*}A_{k})\otimes(\nabla f_{k}(x_{k}^{ t})\otimes\nabla f_{k}(x_{k}^{t}))]_{ij,q\ell}[\theta]_{q\ell}$ $\displaystyle=\sum_{q,\ell=1}^{n}[A_{k}^{*}A_{k}]_{iq}[\nabla f_{k}(x_{k}^{t}) \otimes\nabla f_{k}(x_{k}^{t})]_{j\ell}[\theta]_{q\ell},\text{ by }\eqref{eq:: AotimesB-identity}$ $\displaystyle=\sum_{q,\ell=1}^{n}[A_{k}^{*}A_{k}]_{iq}[\nabla f_{k}(x_{k}^{t}) ]_{j}[\nabla f_{k}(x_{k}^{t})]_{\ell}[\theta]_{q\ell}$ $\displaystyle=[(A_{k}B_{k}^{t})^{*}(A_{k}B_{k}^{t})\theta]_{ij},$ $$

as required, due to $[A_{k}^{*}A_{k}]_{jq}=[A_{k}^{*}A_{k}]_{qj}$.

∎

###### Proposition 10 .

If each $f_{k}$ can be written as a least-squares function (15), then $\alpha_{t}$ can be given as

$$ $\alpha_{t}=\frac{{\lambda_{t}}\tilde{\theta}+\frac{1}{N}\sum_{k=1}^{N}\|\nabla f _{k}(x_{k}^{t})\|^{2}}{{\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}\|A_{k}\nabla f_ {k}(x_{k}^{t})\|^{2}},$ (69) $$

if ${\lambda_{t}}>0$ or $A_{j}\nabla f_{j}(x_{j}^{t})\neq\underline{0}$ for some $j\in\{1,\cdots,N\}$.

###### Proof.

Take $B_{k}^{t}:\mathbb{R}\to\mathcal{X}$ such that

$$ $\displaystyle B_{k}^{t}(\alpha)$ $\displaystyle=\alpha\nabla f_{k}(x_{k}^{t}).$ $$

Then for $\alpha\in\mathbb{R}$

$$ $\displaystyle\langle B_{k}^{t}(\alpha),w\rangle=\langle\alpha\nabla f_{k}(x_{k }^{t}),w\rangle=\alpha\langle\nabla f_{k}(x_{k}^{t}),w\rangle.$ $$

Therefore

$$ $(B_{k}^{t})^{*}(w)=\langle\nabla f_{k}(x_{k}^{t}),w\rangle.$ (70) $$

then general formula 3 gives the desired result as

$$ $\displaystyle(B_{k}^{t})^{*}(A_{k}^{*}A_{k}B_{k}^{t}(\alpha))$ $\displaystyle=\alpha\langle\nabla f_{k}(x_{k}^{t}),A_{k}^{*}A_{k}\nabla f_{k}( x_{k}^{t})\rangle=\alpha\|A_{k}\nabla f_{k}(x_{k}^{t})\|^{2},$ $\displaystyle(B_{k}^{t})^{*}(\nabla f_{k}(x_{k}^{t}))$ $\displaystyle=\langle\nabla f_{k}(x_{k}^{t}),\nabla f_{k}(x_{k}^{t})\rangle=\| \nabla f_{k}(x_{k}^{t})\|^{2}.$ $$

Then the result follows from proposition 3.
∎

###### Proposition 11 .

For $n_{1},n_{2}\in\mathbb{N}$, let $\mathcal{X}=\mathbb{R}^{n_{1}\times n_{2}}$. Define $B_{k}^{t}:\mathcal{X}\to\mathcal{X}$ be such that for any convolutional kernel $\kappa\in\mathcal{X}$, we have $B_{k}^{t}(\kappa)=\kappa\ast\nabla f_{k}(x_{k}^{t})$. Then its adjoint $(B_{k}^{t})^{*}:\mathcal{X}\to\mathcal{X}$ is given by

$$ $(B_{k}^{t})^{*}(w)=w\ast\overline{\nabla f_{k}(x_{k}^{t})},$ (71) $$

where for $x\in\mathcal{X}$,

$$ $\overline{x}(k,l)=x(-k,-l).$ (72) $$

###### Proof.

For the adjoint of $B_{k}^{t}$, we have

$$ $\displaystyle\langle B_{k}^{t}(\kappa),w\rangle$ $\displaystyle=\langle\kappa\ast\nabla f_{k}(x_{k}^{t}),w\rangle$ (73) $\displaystyle=\sum_{i,j}[\kappa\ast\nabla f_{k}(x_{k}^{t})](i,j)w(i,j)$ (74) $\displaystyle=\sum_{i,j}\sum_{k,l}\kappa(k,l)[\nabla f_{k}(x_{k}^{t})](i-k,j-l )w(i,j)$ (75) $\displaystyle=\sum_{i,j}\sum_{k,l}\kappa(k,l)[\nabla f_{k}(x_{k}^{t})](i,j)w(i +k,j+l)$ (76) $\displaystyle=\sum_{i,j}\sum_{k,l}\kappa(k,l)[\nabla f_{k}(x_{k}^{t})](i,j)w(i +k,j+l)$ (77) $\displaystyle=\sum_{i,j}[\nabla f_{k}(x_{k}^{t})](i,j)\left(\sum_{k,l}\kappa(k ,l)w(i+k,j+l)\right)$ (78) $\displaystyle=\sum_{i,j}[\nabla f_{k}(x_{k}^{t})](i,j)\left(\sum_{k,l}\kappa(- k,-l)w(i-k,j-l)\right)$ (79) $\displaystyle=\langle\nabla f_{k}(x_{k}^{t}),\overline{\kappa}\ast w\rangle,$ (80) $$

where $\overline{\kappa}(k,l)=\kappa(-k,-l)$.
∎

###### Proof of Lemma 2 . .

###### Proof of Corollary 1 . .

###### Proof of Proposition 3 . .

###### Proposition 7 .

###### Proof.

###### Corollary 2 .

###### Proof.

###### Proposition 8 .

###### Proof.

###### Proposition 9 .

###### Proof.

###### Proposition 10 .

###### Proof.

###### Proposition 11 .

###### Proof.

### C.1 Approximating optimal linear parameters

For general functions $f_{k}$, a closed-form solution does not exist for calculating linear parameters. Instead, we require an optimization algorithm to approximate these quantities. With information of $\nabla g_{t,{\lambda_{t}}}(\theta)$, and $L_{g_{t,{\lambda_{t}}}}$, the Lipschitz constant of $\nabla g_{t,{\lambda_{t}}}(\theta)$, one can use any first-order convex optimization algorithm, such as gradient descent, Nesterov accelerated gradient (Nesterov et al., ), or stochastic optimization methods such as SGD, and SVRG (Gower et al., ) (especially for large $N$, due to both speed and memory considerations) to approximate $\theta_{t}^{*}$. For example, one can start at an initial point $\theta_{t}^{0}$ at iteration $t$ and update via gradient descent

$$ $\theta_{t}^{w+1}=\theta_{t}^{w}-\frac{1}{L_{g_{t,{\lambda_{t}}}}}\nabla g_{t,{ \lambda_{t}}}(\theta_{t}^{w}).$ (81) $$

The following result illustrates how $\nabla g_{t,{\lambda_{t}}}(\theta)$ and $L_{g_{t,{\lambda_{t}}}}$ can be calculated.

###### Proposition 12 .

For a general linear parametrization $G$, the gradient of $g_{t,{\lambda_{t}}}$ with respect to $\theta$ and its associated Lipschitz constant can be calculated as

$$ $\displaystyle\nabla g_{t,{\lambda_{t}}}(\theta)$ $\displaystyle={\lambda_{t}}(\theta-\tilde{\theta})-\frac{1}{N}\sum_{k=1}^{N}(B _{k}^{t})^{*}\nabla f_{k}(x_{k}^{t}-G_{\theta}\nabla f_{k}(x_{k}^{t})),$ (82) $\displaystyle L_{g_{t,{\lambda_{t}}}}$ $\displaystyle={\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}\|B_{k}^{t}\|^{2}.$ (83) $$

###### Proof.

As

$$ $\displaystyle g_{t,{\lambda_{t}}}(\theta)$ $\displaystyle=\frac{1}{N}\sum_{k=1}^{N}f_{k}(x_{k}^{t}-G_{\theta}\nabla f_{k}( x_{k}^{t}))+\frac{\lambda_{t}}{2}\|\theta-\tilde{\theta}\|^{2}$ $\displaystyle=\frac{1}{N}\sum_{k=1}^{N}f_{k}(x_{k}^{t}-B_{k}^{t}\theta)+\frac{ \lambda_{t}}{2}\|\theta-\tilde{\theta}\|^{2},$ $$

then by the chain rule

$$ $\nabla g_{t,{\lambda_{t}}}(\theta)=-\frac{1}{N}\sum_{k=1}^{N}(B_{k}^{t})^{*} \nabla f_{k}(x_{k}^{t}-B_{k}^{t}\theta)+{\lambda_{t}}(\theta-\tilde{\theta}),$ (84) $$

as required. To calculate the smoothness constant, we have

$$ $\displaystyle\|\nabla g_{t,{\lambda_{t}}}(\theta_{1})-\nabla g_{t,{\lambda_{t} }}(\theta_{2})\|$ $\displaystyle=\left\|{\lambda_{t}}(\theta_{1}-\theta_{2})+\frac{1}{N}\sum_{k=1 }^{N}(B_{k}^{t})^{*}(\nabla f_{k}(x_{k}^{t}-B_{k}^{t}\theta_{2})-\nabla f_{k}( x_{k}^{t}-B_{k}^{t}\theta_{2}))\right\|$ $\displaystyle\leq{\lambda_{t}}\|\theta_{1}-\theta_{2}\|+\frac{1}{N}\sum_{k=1}^ {N}\left\|(B_{k}^{t})^{*}(\nabla f_{k}(x_{k}^{t}-B_{k}^{t}\theta_{2})-\nabla f _{k}(x_{k}^{t}-B_{k}^{t}\theta_{1}))\right\|$ $\displaystyle\leq{\lambda_{t}}\|\theta_{1}-\theta_{2}\|+\frac{1}{N}\sum_{k=1}^ {N}\|B_{k}^{t}\|\left\|\nabla f_{k}(x_{k}^{t}-B_{k}^{t}\theta_{2})-\nabla f_{k }(x_{k}^{t}-B_{k}^{t}\theta_{1})\right\|$ $\displaystyle\leq{\lambda_{t}}\|\theta_{1}-\theta_{2}\|+\frac{1}{N}\sum_{k=1}^ {N}L_{k}\|B_{k}^{t}\|\left\|B_{k}^{t}(\theta_{1}-\theta_{2})\right\|$ $\displaystyle\leq\left({\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}\|B_{k}^{t} \|^{2}\right)\left\|\theta_{1}-\theta_{2}\right\|$ $$

Due to the properties of the triangle inequality, the Cauchy-Schwarz inequality, and the operator norm, this bound is tight. Therefore the Lipschitz constant of $\nabla g_{t,{\lambda_{t}}}(\theta)$ is given by

$$ ${\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}\|B_{k}^{t}\|^{2}$ (85) $$

as required.
∎

Using this general result, we can calculate these values for specific parametrizations of $G$.

###### Corollary 3 .

Suppose each $f_{k}\in\mathcal{F}_{L_{k}}$.
Pointwise parametrization
For the pointwise parametrization, $\theta\in\mathcal{X}$, and

$$ $\nabla g_{t,{\lambda_{t}}}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})-\frac{ 1}{N}\sum_{k=1}^{N}\nabla f_{k}(x_{k}^{t}-\theta\odot\nabla f_{k}(x_{k}^{t})) \odot\nabla f_{k}(x_{k}^{t}),$ (86) $$

and an upper bound of the Lipschitz constant of $\nabla_{\theta}g$ is given by

$$ $L_{\nabla_{\theta}g}={\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}(\max\{|[ \nabla f_{k}(x_{k}^{t})]_{1}|,\cdots,|[\nabla f_{k}(x_{k}^{t})]_{n}|\})^{2}.$ (87) $$

Full operator parametrization
In this case we have $\theta\in\mathcal{L}(\mathcal{X})$. The gradient of $g_{t,{\lambda_{t}}}(\theta)$ is given by

$$ $\nabla g_{t,{\lambda_{t}}}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})-\frac{ 1}{N}\sum_{k=1}^{N}\nabla f_{k}(x_{k}^{t}-\theta\nabla f_{k}(x_{k}^{t})) \otimes\nabla f_{k}(x_{k}^{t}),$ (88) $$

and an upper bound of the Lipschitz constant of $\nabla g_{t,{\lambda_{t}}}(\theta)$ is given by

$$ ${\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}\left\|\nabla f_{k}(x_{k}^{t}) \right\|^{2}.$ (89) $$

Scalar step size
We now take $\theta\in\mathbb{R}$. The derivative of $g_{t,{\lambda_{t}}}$ with respect to $\theta$ is given by

$$ $g_{t,{\lambda_{t}}}^{\prime}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})- \frac{1}{N}\sum_{k=1}^{N}\langle\nabla f_{k}(x_{k}^{t}-\theta\nabla f_{k}(x_{k }^{t})),\nabla f_{k}(x_{k}^{t})\rangle,$ (90) $$

and the Lipschitz constant of $g^{\prime}(\theta)$ is given by

$$ ${\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}\|\nabla f_{k}(x_{k}^{t})\|^{2}.$ (91) $$

Convolution
In this case we have $\theta\in\mathbb{R}^{n_{1}\times n_{2}}$. The gradient of $g_{t,{\lambda_{t}}}(\theta)$ is given by

$$ $\nabla g_{t,{\lambda_{t}}}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})-\frac{ 1}{N}\sum_{k=1}^{N}\nabla f_{k}(x_{k}^{t}-G_{\theta}\nabla f_{k}(x_{k}^{t})) \ast\overline{\nabla f_{k}(x_{k}^{t})}.$ (92) $$

###### Proof.

Pointwise parametrization

In this case, we have $\theta\in\mathcal{X}$ and $B_{k}^{t}(x)=\nabla f_{k}(x_{k}^{t})\odot x$ and $(B_{k}^{t})^{*}(x)=B_{k}^{t}x$ for $x\in\mathcal{X}$. Furthermore,

$$ $\displaystyle\|B_{k}^{t}\|$ $\displaystyle=\max_{x\neq 0}\frac{\|x\odot\nabla f_{k}(x_{k}^{t}))\|}{\|x\|}= \max_{x\neq 0}\sqrt{\frac{\sum_{i=1}^{n}[x]_{i}^{2}[\nabla f_{k}(x_{k}^{t})]_{ i}^{2}}{\sum_{i=1}^{n}[x]_{i}^{2}}}$ $\displaystyle\leq\max_{q}|[\nabla f_{k}(x_{k}^{t})]_{q}|\max_{x\neq 0}\sqrt{ \frac{\sum_{i=1}^{n}[x]_{i}^{2}}{\sum_{i=1}^{n}[x]_{i}^{2}}}=\max\{|[\nabla f_ {k}(x_{k}^{t})]_{1}|,\cdots,|[\nabla f_{k}(x_{k}^{t})]_{n}|\}.$ $$

Full operator parametrization

In the case of the full operator parametrization, we have $(B_{k}^{t})^{*}(w)=w\otimes\nabla f_{k}(x_{k}^{t})$. Therefore, using Proposition 12 gives (88). For the Lipschitz constant, note that

$$ $\displaystyle\|B_{k}^{t}(P)\|$ $\displaystyle=\|P\nabla f_{k}(x_{k}^{t})\|\leq\|P\|\|\nabla f_{k}(x_{k}^{t})\|,$ $$

and therefore

$$ $\displaystyle\|B_{k}^{t}\|$ $\displaystyle=\max_{P\neq 0}\frac{\|B_{k}^{t}(P)\|}{\|P\|}\leq\|\nabla f_{k}(x _{k}^{t})\|.$ $$

Scalar step size
Let $B_{k}^{t}$ be defined for any $\alpha\in\mathbb{R}$ by $B_{k}^{t}(\alpha)=\alpha\nabla f_{k}(x_{k}^{t})$, then for an element $w\in\mathcal{X}$, $(B_{k}^{t})^{*}(w)=\langle\nabla f_{k}(x_{k}^{t}),w\rangle$. Furthermore,

$$ $\displaystyle\|B_{k}^{t}(\alpha)\|$ $\displaystyle=\|\alpha\nabla f_{k}(x_{k}^{t})\|$ $\displaystyle=|\alpha|\|\nabla f_{k}(x_{k}^{t})\|,$ $$

and so

$$ $\displaystyle\|B_{k}^{t}\|$ $\displaystyle=\max_{\alpha\neq 0}\frac{\|B_{k}^{t}(\alpha)\|}{|\alpha|}=\| \nabla f_{k}(x_{k}^{t})\|.$ $$

Convolution
For the gradient,

$$ $\displaystyle(B_{k}^{t})^{*}(\nabla f_{k}(x_{k}^{t}-G_{\theta}\nabla f_{k}(x_{ k}^{t})))=\nabla f_{k}(x_{k}^{t}-G_{\theta}\nabla f_{k}(x_{k}^{t}))\ast \overline{\nabla f_{k}(x_{k}^{t})}.$ $$

∎

For any chosen linear parametrization, one can approximate the operator norm of $B_{k}^{t}$ using the power method (Golub & Van Loan, ). The following table summarises the previous propositions:

**Table 2: Example parametrization properties**
| Parametrization | Equations |
| --- | --- |
| Pointwise | •<br>$\tilde{\theta}=\tau\mathbf{1}\in\mathcal{X}$<br>•<br>$B_{k}^{t}(x)=\nabla f_{k}(x_{k}^{t})\odot x$, $(B_{k}^{t})^{*}(x)=B_{k}^{t}(x)$<br>•<br>$\nabla g_{t,{\lambda_{t}}}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})-\frac{<br>1}{N}\sum_{k=1}^{N}\nabla f_{k}(x_{k}^{t}-\theta\odot\nabla f_{k}(x_{k}^{t}))<br>\odot\nabla f_{k}(x_{k}^{t})$<br>•<br>$L_{g_{t,{\lambda_{t}}}}\leq{\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}(\max\{<br>|\nabla f_{k}(x_{k})_{1}|,\ldots,|\nabla f_{k}(x_{k})_{n}|\})^{2}$ |
| Full operator | •<br>$\tilde{\theta}=\tau I\in\mathcal{L}(\mathcal{X})$<br>•<br>$B_{k}^{t}(P)=P\nabla f_{k}(x_{k}^{t})$, $(B_{k}^{t})^{*}(w)=w\otimes\nabla f_{k}(x_{k}^{t})$<br>•<br>$\nabla g_{t,{\lambda_{t}}}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})-\frac{<br>1}{N}\sum_{k=1}^{N}\nabla f_{k}(x_{k}^{t}-\theta\nabla f_{k}(x_{k}^{t}))<br>\otimes\nabla f_{k}(x_{k}^{t})$<br>•<br>$L_{g_{t,{\lambda_{t}}}}\leq{\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}\|<br>\nabla f_{k}(x_{k}^{t})\|^{2}$ |
| Scalar | •<br>$\tilde{\theta}=\tau\in\mathbb{R}$<br>•<br>$B_{k}^{t}(\alpha)=\alpha\nabla f_{k}(x_{k}^{t})$, $(B_{k}^{t})^{*}(w)=\langle w,\nabla f_{k}(x_{k}^{t})\rangle$<br>•<br>$g_{t,{\lambda_{t}}}^{\prime}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})-<br>\frac{1}{N}\sum_{k=1}^{N}\langle\nabla f_{k}(x_{k}^{t}-\theta\nabla f_{k}(x_{k<br>}^{t})),\nabla f_{k}(x_{k}^{t})\rangle$<br>•<br>$L_{g_{t,{\lambda_{t}}}}={\lambda_{t}}+\frac{1}{N}\sum_{k=1}^{N}L_{k}\|\nabla f<br>_{k}(x_{k}^{t})\|^{2}$ |
| Convolution | •<br>$\tilde{\theta}(i,j)=\begin{cases}\tau,\text{ if }i=j=0,\\<br>0,\text{ otherwise}.\end{cases}$<br>•<br>$B_{k}^{t}(\kappa)=\kappa\ast\nabla f_{k}(x_{k}^{t})$, $(B_{k}^{t})^{*}(\kappa)=\kappa\ast\overline{\nabla f_{k}(x_{k}^{t})}$<br>•<br>$g_{t,{\lambda_{t}}}^{\prime}(\theta)={\lambda_{t}}(\theta-\tilde{\theta})-<br>\frac{1}{N}\sum_{k=1}^{N}\nabla f_{k}(x_{k}^{t}-\theta\ast\nabla f_{k}(x_{k}^{<br>t}))\ast\overline{\nabla f_{k}(x_{k}^{t})}$ |

###### Proposition 12 .

###### Proof.

###### Corollary 3 .

###### Proof.

## Appendix D Additional Numerical Results

### D.1 Best and worst performance of learned convolution

In this section, we compare the performance of the learned algorithm on individual functions in the test dataset (of $100$ functions). Figure 8 shows that for the deblurring problem best performance of the learned convolutional algorithm take less than one quarter of the iterations as NAG to reach a tolerance of $10^{-7}$, and the worst performance takes approximately one third the iterations. Therefore, across all $100$ functions in the test set the learned convolutional algorithm performs significantly better and quite similarly for different functions.

Figure: Figure 8: Best and worst case convergence of the learned convolutional algorithm compared to NAG for the deblurring problem, considering the ratio of number of iterations to reach a tolerance of $10^{-7}$.
Refer to caption: extracted/5992054/blur_images/best_case_blur.png

Figure 9 shows that for the CT problem, the worst performance of the learned convolutional algorithm takes less than one third of the iterations as NAG to reach a tolerance of $10^{-9}$. The best case reaches a tolerance of $10^{-9}$ in about one fifth of the iterations. However, a tolerance of $10^{-5}$ is a sufficient tolerance for reconstruction quality and is achieved in only $4$ iterations with the learned method, whereas NAG requires $28$ iterations. Thus, for both of the problems, we achieve superior performance with the learned convolutional algorithm for every single function in both of the test sets. We also see similar speedups for both problems for all functions when using the learned algorithm with a convolutional parametrization.

Figure: Figure 9: Best and worst case convergence of the learned convolutional algorithm compared to NAG for the CT problem, considering the ratio of number of iterations to reach a tolerance of $10^{-9}$.
Refer to caption: extracted/5992054/CT_images/best_case_ct.png

### D.2 Ablation study: size of learned kernels

Figure 10 shows the many of the learned convolutional algorithms outperform NAG for the deblurring problem. The $3\times 3$ learned kernels outperform NAG initially but are eventually overtaken before a tolerance of $10^{-7}$. We see that the $5\times 5$ kernels significantly outperform the NAG kernels, with then incremental improvements with larger kernel sizes as the $96\times 96$ kernels achieve a tolerance of $10^{-8}$ is half the iterations as the $5\times 5$ kernels, while the $3\times 3$ kernels achieve of tolerance of $10^{-8}$ in approximately $900$ iterations, while the $5\times 5$ kernels achieve this tolerance in about $500$ iterations, compared to just over $200$ iterations for the $96\times 96$ kernels.

Figure: Figure 10: Test performance of different kernel sizes in the convolutional parametrisation, averaged over the test dataset in the deblurring problem. Tested kernel sizes are $3\times 3$, $5\times 5$, $7\times 7$, $11\times 11$, $3\times 3$, $96\times 96$.
Refer to caption: extracted/5992054/blur_images/kernel_sizes_test_semilogy.png

### D.3 Tolerance Tables

**Table 3: The first row shows error thresholds for the deblurring problem. The entries in the table show the number of required iterations to fall below the respective error threshold. "na" means that the threshold was not reached within 2000 iterations.**
|  | $10^{-1}$ | $10^{-2}$ | $10^{-3}$ | $10^{-4}$ | $10^{-5}$ | $10^{-6}$ | $10^{-7}$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Convolution | 1 | 2 | 14 | 40 | 72 | 115 | 176 |
| L-BFGS | 3 | 10 | 39 | 104 | 220 | 384 | 565 |
| NAG | 4 | 15 | 56 | 137 | 266 | 432 | 647 |
| Backtracking Line Search | 2 | 12 | 105 | 553 | 1970 | na | na |
| Scalar | 3 | 17 | 173 | 930 | na | na | na |
| Pointwise | 2 | 14 | 190 | 1368 | na | na | na |
| GD | 5 | 32 | 346 | 1758 | na | na | na |

**Table 4: The first row shows error thresholds for the deblurring problem. The entries in the table show the number of required iterations to fall below the respective error threshold. "na" means that the threshold was not reached within 1000 iterations. The convolutional method clearly performs best**
|  | $10^{-3}$ | $10^{-4}$ | $10^{-5}$ | $10^{-6}$ | $10^{-7}$ | $10^{-8}$ | $10^{-9}$ | $10^{-10}$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Convolution | 1 | 2 | 4 | 5 | 8 | 14 | 22 | 29 |
| L-BFGS | 4 | 8 | 15 | 26 | 39 | 55 | 71 | 83 |
| NAG | 5 | 11 | 19 | 33 | 49 | 71 | 110 | 141 |
| Backtracking Line Search | 5 | 9 | 21 | 47 | 87 | 141 | 201 | na |
| Pointwise | 2 | 8 | 22 | 56 | 113 | 189 | na | na |
| Scalar | 3 | 8 | 23 | 56 | 111 | 184 | na | na |
| GD | 6 | 19 | 52 | 129 | na | na | na | na |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: * Equation 2 seems to only vary the y over the whole space. I believe this should be rewritten to emphasise what kind of functions you expect to be varying over. I.e. y should be from some underlying distribution? Do regularisers get varied? Do operators get varied?
* Line 108, you seem to choose X to be a finite dim Hilbert space - what is the value of this? If X is finite dim, then practically there is not point in distinguishing it from Euclidean space or am I missing something?
* Proposition 2 - linear independence seems like a rather arbitrary assumption. Can anything be said when it does not hold?
* Equation 9 - this is a different regularizer from the one in equation 2, so worth using different notation for the two.
* Theorem 3 seems to assume that $x_t$ has to now be a bounded sequence, which seems to be a relatively strange assumption to appear in such a context. Can you explain why this is necessary?
* Random question to authors: do you believe that L2O can overcome the problems discussed in https://arxiv.org/pdf/2301.06148? 
* Section 6 - in this section what is Y? 
* Figure 3a - this is not reconstruction, but initialization presumably? 
* Figure 5 - where is the ground truth? Why is it only sinogram and reconstruction shown?
soundness: 4
strengths: The theoretical part of the paper is well-written and provides a very clear story. The theoretical framework also provides a good starting point for further generalising analysis of L2O schemes, especially when moving to non-convex settings.

The proposed preconditioner learning is memory-efficient, fast, with convergence guarantees and empirical evidence on small problems.
summary: The paper proposes to use greedy learning to help scale L2O methods by avoiding memory constraints of unrolling, allowing to train for more iterations. The paper focuses on learning a linear operator preconditioning operator, minimizing the function value at subsequent iteration. The paper shows that such parametrisation, by virtue of generalising gradient descent, the iterations admit provable convergence guarantees, even on unseen data. Such preconditioner parametrization is shown to outperform classical optimization algorithms, such as NAG and L-BFGS, in experiments on two image inverse problems: image deblurring and Computed Tomography.
weaknesses: Major: 

Main weaknesses appear in the numerical section of this paper. Overall, the numerical section is very difficult to read, as it seems to mix the exact parameter choices with the rest of the explanations, further obfuscating everything. The various details seemed to have been mixed into a single soup of information - this should be summarised better. 

The numerical comparison seems to be missing the main evaluation - it is unclear whether the method actually generalises. Only a small example is portrayed, and no evaluation over the whole dataset is provided. It is not clear to me whether Figure 2 converges and Figure 6 diverges simply due to a different example being provided. Preferably there would also be some analysis of the initialisation question for the problems of consideration. 

There is also a summary/interpretation missing from the numerical section - it is not clear what the numerical results seem to be showing beyond efficiency over classical methods? Why are fully learned preconditioners bad? Is this observed for all problems? 

There also seems to be very little mentioned about the limitations of this approach - can this be expanded upon? Currently (at least to me), the question of computational costs behind hyperparameter tuning is not very clear. Also, memory footprint is unclear to me - the greedy approach needs to store a matrix for each step - thus linearly increasing memory requriement with T?

Minor: 

The method has only been illustrated on rather small scale problems (limited to 40x40 and 96x96) - it would be interesting to see whether same behaviour is observed on more realistic image sizes like 256x256 or 1024x1024. 

The paper is limited to convex problems, and while this is a good starting point I do believe that this is a significant weakness, especially given the interest in using L2O for optimization of non-convex problems. In the same vain, this approach (or at least the analysis) seems to be limited to differentiable functions.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: ## Smaller concerns

- **BGD assumption** Please correct me otherwise, but the "Better than Gradient Descent" (BGD) is assumed rather rather than proved, and theorem 2 both uses that $t \to \infty$ while requiring a final training iteration $T$. These two statements seem contradictory? 
- The BGD assumption seem unecessary for the unseen problem proof? A similar argument could be made without the $t \to \infty$ or BGD assumption if the preconditioner is PD. For a given training budget, the algorithm repeats the last learned preconditioner $G_{\theta_{T-1}}$, so it will eventually converge on any unseen function that is smooth relative to $G_{\theta_{T-1}}$ in the sense that $\nabla^2 f(x) \preceq [G_{\theta_{T-1}}]^{-1}$. This doesn't seem much weaker than the current proposition, which only guarantees that there exists a smoothness constant $\tilde L$ such that the algorithm will convergence on $\tilde L$-smooth functions.

- **Related work in optimization**: The discussion of related work in optimization only touches on L-BFGS and ignores relevant work that attempt to achieve a similar goal, "find a better step-size/preconditioner for the problem", but by running additional computation before/while solving the problem rather than by taking the learning to optimize approach. While the approaches are different, this literature should at least be acknowledged in a paragraph in the introduction as it shares a similar goal. Examples include a simple [Armijo line-search](https://projecteuclid.org/journals/pacific-journal-of-mathematics/volume-16/issue-1/Minimization-of-functions-having-Lipschitz-continuous-first-partial-derivatives/pjm/1102995080.full), 
[optimal diagonal preconditioners](https://arxiv.org/abs/2209.00809) for quadratic, 
[AdaGrad](https://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) and [adaptive bound optimization](https://arxiv.org/abs/1002.4908) for convex non-smooth functions, [multidimensional Armijo](https://arxiv.org/abs/2306.02527) for smooth strongly convex functions, or [parameter-free methods in online learning](Coin Betting and Parameter-Free Online Learning).


## Questions

- Please clarify what is meant by "maintaining constant memory usage" and "memory is constant with increasing training iteration". My understanding of the algorithm is that the methods learns a different preconditioner for each iteration. This scales at least linearly with the number of training iterations. A more detailled explaination of where the memory is used and how it differs from the unrolling strategy would help.

- I don't understand Eq. 16. What is $B_k$? Equation (14) treats both $G_\theta$ and $B_k^t(\tehta)$ as equations of $\theta$, which makes 

## Minor

- Proposition 1 & 2 are missing the assumption that no entries of $\nabla f(x)$ is $0$. It is possible for the $j$th entry of $\nabla f(x)$ to be 0 while having $x[j] \neq x^*[j]$ where $x^* \in \arg\min_x f(x)$.
soundness: 2
strengths: The learning to optimize literature is still in development, and does not yet have well defined formalism or benchmark problems. As such, the goal of the paper as bringin formal guarantees to this setting is a new and relevant contribution to the community. The proposed algorithms show promising results on the experimental setup.
summary: The submission introduces a "learning to optimize" algorithm that learns a sequence of fixed preconditioners to apply to gradient descent by fitting them greedily to maximize the one-step progress on a set of training problems. The paper presents convergence guarantees that the algorithm can fit and is guaranteed to converge when run for longer than trained on some problems that are not in the training set. The submission introduces multiple types of preconditioners and shows experimental results on deblurring and tomography problems.
weaknesses: **Safeguards.** The proposed algorithm still relies on hand-crafted a-priori knowledge of the optimization problem in the form of the maximum step-size that would work for all training problem, $\tau$. As such, it does not significantly differ from alternative approaches that require safeguarding or search within a predefined set that guarantees convergence.

**Experimental validation.** The definition of train and test problems in the numerical experiments is not sufficiently transparent. If I understood it correctly, both problems are of the form $\|Ax-b\|^2$ but the training and test "samples" only differ in $x$ and $b$; the linear operator $A$ is fixed. This seems like an ``easy'' problem, as the goal of learning algorithm is to learn a preconditioner that approximates the inverse of $A^TA$, which is fixed across both training and test. The paper should make the distinction clear and highlight possible limitations. A more thorough experimental evaluation that tests the trained algorithm against other blurring operators, at least Gaussian blur with different parameters, would help make the claim that the learning algorithm can indeed generalize.

**Generalization guarantees.** The submission claims that the given algorithm "ensures convergence on unsees data" but the guarantees seem weak. Theorem 2 guarantees that there exists functions for which the given algorithm will work, but this is weaker than typical generalization guarantee. Translated to the classification setting, I would take to mean "there exists samples that were not in the training set that the algorithm classifies correctly", which is not a strong statement about the performance of the algorihm. Would it be possible to guarantee that the algorithm would work on all $1/\tau$ smooth problems, or to characterize that the preconditioner learned by the algorithm will lead to better performance on other problems if the train and test problems use the same operator $A$?

**Convolutional preconditioners** That the experimental results show a significant improvement in performance when used the learned convolutional preconditioner makes it unclear whether the benefit arises from the convolutional parameterization or the "learning to optimize" approach, and one could envision a BFGS-like algorithm using the convolutional structure. Although this is not my area of expertise, my understanding is that specialized approaches for deblurring using convolutional preconditioners exist, see for example [the work of Eboli et at.](https://arxiv.org/pdf/2007.01769). A discussion of, and a comparison with, specialized algorithms would be a welcome addition.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: As mentioned above, please clarify the doubts regarding the theoretical analysis and include comparision with existing provable L2O schemes.

Meanwhile, the proposed scheme is tailored for inverse problems in imaging, particularly the convolutional parameterization of the preconditioner -- should this be make clearer in title/abstract/introduction?
soundness: 3
strengths: The greedy training approach is a novel and interesting scheme for L2O. Indeed, current L2O methods are mostly depending on unrolling several iterations, the proposed scheme is very timely for L2O area for scalability of training the optimizer.

The numerical performance of the proposed scheme on inverse problems is very impressive.
summary: This paper proposed a novel L2O approach for inverse problems. To mitigate the computational challenges in current approaches in L2O, which are mostly based on unrolling, this paper investigate a greedy training approach which decouples the iterates leading to more efficient training. The proposed scheme provides an effective approach for training a preconditioner for gradient descent. Theoretical analysis demonstrate that the proposed preconditioned gradient descent converges under the BGD condition. A specialized parameterization using convolution has been proposed in the paper which is tailored for imaging applications. Numerical experiments on inverse problems have demonstrated superior performance over classical methods such as Nesterov's accelerated gradient and L-BFGS.
weaknesses: The theoretical part of the paper seems to be weak. The convergence analysis is relying on an unrealistic assumption named BGD (better than gradient descent) assumption across each iteration -- you can't just simply assume what you wish to proof. Corollary 1 seems to give a very strong claim but no explicit proof is given (it is unclear how simply applying Lemma 2 can lead to such claim). The reviewer believes that whether or not the learned linear preconditioner is BGD should be non-trivial to show -- it should certainly depend on the actual values of the trained parameters. The "with Convergence Guarantees" part of the claimed contribution is unfortunately not valid.

In terms of further development, a limitation of the proposed scheme is inability to learn those iterative schemes which utilize memory of past iterates -- that is, in fact an advantage of unrolling which deserves acknowledgement. The reviewer believes that, ultimately, the proposed greedy scheme should be jointly applied with unrolling for the best performance.

In terms of experiments, although classical hand-crafted optimizers are included as baselines, there is no comparison with other existing L2O methods. For example, the author(s) could consider the SOTA (truly) provably convergent method by: Banert S, Rudzusika J, Öktem O, Adler J. Accelerated forward-backward optimization using deep learning. SIAM Journal on Optimization. 2024 Jun 30;34(2):1236-63.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: * (l.327) Should this be (PC) instead of (PP)?
* How are the ground truths $x^*$ computed?
* While the linear parameterization should give minimal overhead, a wall-clock time comparison of test-time might be useful for further comparison.
* For quadratic problems, Tan et al. (2023b) consider the preconditioning $G = (A^* A)^{-1}$, which can be used to motivate the convolutional structure. How does the proposed method compare to perhaps a regularized version, say $(I + \delta^{-1} A^* A)^{-1} \approxeq I - \delta A^\dagger A^{*\dagger}$. I am not sure if this is available in standard libraries.
* The convolution-based preconditioner seems to generalize something called "Laplacian smoothing gradient descent" (Osher et al., 2022). Have the authors considered other possible instances of such linear preconditionings that may have better empirical performance?
* Related: perhaps a short reference to App. D.2 would be helpful in the main text, as the choice of full-image convolution is not motivated.
* (l.737) RHS of first inequality should be $g_{t,0}(\tilde{\theta})$, and $\nabla f_k$ on the second line. Proof of Thm 1 perhaps needs a short clarification on telescoping arguments (relating objective of $x_{t+1}$ in terms of $x_t$) since $x_t$ is not generated using GD, but the result should be the same.
* Prop 4 can be more proved more succinctly by noting the residual of $F$ is the average of residuals of $f_k$, which are non-negative, so $f_k^t - f_k^* \le N F_t - F^*$.
* Prop 5 can directly use coordinate projection and remove the elementary calculation: consider the convex function $\pi_i f$ and first-order optimality, clearly minimized at $x^*$ and with derivative equal to $[\nabla f]_i$.
* Proof of Lem 1: notation from $g_{t,\lambda_t}$ to $g_t(\cdot, \lambda_t)$. Inconsistency in the first inequality with definition of BGD: should be $g_t(\theta_t, \lambda_t) \le g_t(\tilde{\theta}, 0)$.
* Line 40: perhaps reference Theorem 1 here when claiming that the sums of the $f_k$ converge to the optimal values.

[1] Osher, S., Wang, B., Yin, P., Luo, X., Barekat, F., Pham, M., & Lin, A. (2022). Laplacian smoothing gradient descent. Research in the Mathematical Sciences, 9(3), 55.
soundness: 4
strengths: * The paper reads well and the theorems are very intuitive. Overall, a good submission to ICLR.
* The experiments are convincing and the details are very comprehensive. The speed of training is especially good due to the simple parameterization, and comparisons are made with the greedy training vs standard unrolled training.
* The clear structure of the learned kernels (e.g. Figure 1c, 5c) is good cause for further investigation towards optimal preconditioning for certain imaging problems.
summary: This work proposes a provably convergent learning-to-optimize method based on preconditioned gradient descent. By considering gradient descent and regularizing the proposed algorithm such that it is majorized by and eventually becomes GD, they demonstrate significant speedups on various convex optimization problems, while maintaining provable convergence guarantees from GD. Moreover, the linear parameterization allows for convex solvers at each timestep which gives speedup compared to other L2O methods.
weaknesses: * The theory seems to require the knowledge of the maximum Lipschitz constant over all training examples in its regularization. How well does the method work when applied to problems with large or unknown constant?
* Related: finding the regularization parameter $lambda_T$ as in p.18 seems arduous. Is this done at every iteration $t$, and how many problems need to be solved to find $\lambda_t$? It seems also that the condition for choosing $\lambda_t$ gives that $G_\phi$ is positive definite, which is confusing with the initial claim that the method converges for non p.d. conditioners.
* While stepsize $\tau = 1/L$ does indeed give convergence for NAG, in practice, the smoothness is unknown and larger step-sizes can be taken while maintaining empirical convergence. Same for L-BFGS. Since the convergence profiles are quite similar, a proper comparison with differing parameters for NAG and L-BFGS would greatly help in ascertaining the role of interacting pixels in helping optimization.

Meta-review: The paper studies a greedy training method for learning to optimize. In the proposed approach, parameters are determined sequentially. Given the parameters for the first k iterations, parameters for iteration k+1 are chosen to minimize the average objective value over the training set. Greedy training has efficiency advantages compared (e.g.) to loop unrolling, since one doesn’t have to backpropagate through multiple iterations. For linear parameterizations (i.e., a step size, element wise scale, or preconditioning matrix) this parameter selection operation is a convex program (indeed, simply least squares if the objective is the squared error). 

The paper analyzes the greedy scheme theoretically, arguing that (1) if the class of admissible step rules includes gradient rules, then greedy learning performs at least as well as gradient decent on training data, and (2) if the regularization parameters are chosen such that for large iterations, the chosen step rule tends to the gradient rule, then the performance on unseen (test) data inherits the convergence rate of gradient descent. The proposed approach is applied to inverse problems in image deblurring and computed tomography, where it used to learn convolutional preconditioners. 

The main strength of the paper is its relatively simple, practical proposal for greedy learning to optimize. Compared to unrolling approaches, this method is scalable to large numbers of iterations without needing to backpropagate, and with constant memory. As described below, the discussion clarified a number of issues around the paper’s theory — in particular the meaning of the “better than gradient descent” condition. At the same time, reviewers retained concerns about the paper’s generalization theory and its experiments, which would be stronger with comparisons across problem types (e.g., different blur operators) and with comparisons to existing approaches to L2O (not just classical optimization baselines).

Decision: Reject