Online Learning to Precondition III. Concise analysis of a quasi-Newton type algorithm
This is a short post that continues the previous discussion on “online learning to precondition”, and I’ll present a very concise quasi-Newton type method with a two-line proof. The algorithm follows exactly the OL2P intuition.
AI-disclosure. All my blogs marked by 
emphasize insights behind technical results, and I write the whole blog draft by myself. Generative AI tools are only used for result verification/proofreading.
Introduction and background
Again and again, to simplify the discussion, I’ll resort to the quadratic function convention
\[ \min_{x \in \mathbb{R}^n} \quad \tfrac{1}{2} \langle x, A x \rangle - \langle b, x \rangle =: f(x) \]
where \(\mu I \preceq A \preceq L I\), and it corresponds to solving a linear system \(A x = b\).
Sparse approximate inverse preconditioner. There are many ways to construct preconditioners for \(A\), and one of them is the sparse approximate inverse preconditioner (SPAI) (Benzi and Tuma 1999) that aims to find \(P\) such that \(\|PA-I\|_F\) (or \(\|AP-I\|_F\)) is reduced. We will abuse the notation of distance and define
\[ \operatorname{dist}(P,A) := \|AP-I\|_F^2. \]
A SPAI preconditioner aims to solve \(\min_P \|AP-I\|_F^2\). Now, suppose we are somewhat “lazy”, and we don’t want to solve the minimization problem. Say we just randomly pick some \(P\) and obtain
\[ x^+ = x - P\nabla f(x). \]
Again, there are basically two cases, categorized by the contraction \(r = \tfrac{\|\nabla f(x^+)\|}{\|\nabla f(x)\|} = \left\|(I-AP){\color{red}\tfrac{\nabla f(x)}{\|\nabla f(x)\|}}\right\|\).
Case 1. Either the contraction is good. i.e., \(r = \tfrac{\|\nabla f(x^+)\|}{\|\nabla f(x)\|}\) is small. It implies our random choice of \(P\) is lucky .
Case 2. Or the contraction is bad. i.e., \(r\) is large. In this case, we know that \(P\) is not a good preconditioner, and we need to fix it. Suppose we want to update \(P\) with a low-rank modification:
\[ P_+ = P + \eta u v^{\top}, \qquad \|u\|=\|v\|=1, \]
where \(u,v\) will be specified later. Given \(u,v\), the parameter \(\eta\) is chosen to minimize \(\|AP-I\|_F^2\):
\[ \eta^{\star} = \arg\min_{\eta}\|A(P+\eta u v^{\top})-I\|_F^2 = \tfrac{\langle Au,(I-AP)v\rangle}{\|Au\|^2} \]
and by taking \(u = \tfrac{(I-AP)v}{\|(I-AP)v\|}\), we have \(\|AP_+-I\|_F^2 \leq \|AP-I\|_F^2 - \tfrac{1}{\kappa}\|(I-AP){\color{red}v}\|^2\).
The question is how to choose \(\|v\|=1\) with large \(\|(I-AP){\color{red}v}\|\). The answer is obvious given
\[ r = \left\|(I-AP){\color{red}\tfrac{\nabla f(x)}{\|\nabla f(x)\|}}\right\|. \]
Taking \(v = \tfrac{\nabla f(x)}{\|\nabla f(x)\|}\), we have \(\|AP_+-I\|_F^2 \leq \|AP-I\|_F^2 - \tfrac{1}{\kappa}r^2\). The algorithm is very simple.
\begin{algorithm}
\caption{Two-line quasi-Newton}
\begin{algorithmic}
\Require $x^1,P_1$
\For{$k=1,\ldots,K$}
\State \textbf{compute} $v=\tfrac{\nabla f(x^k)}{\|\nabla f(x^k)\|}$, $u=\tfrac{(I-AP)v}{\|(I-AP)v\|}$, $\eta=\tfrac{\langle Au,(I-AP)v\rangle}{\|Au\|^2}$
\State $x^{k+1}=x^k-P\nabla f(x^k)$
\State $P_{k+1}=P_k+\eta u v^{\top}$
\EndFor
\end{algorithmic}
\end{algorithm}
Again, the proof follows from the “online learning to precondition” relation. Define the learning potential
\[ \phi(x,P) := \log\|\nabla f(x)\| + \kappa\|AP-I\|_F^2. \]
We have the following reduction in the potential.
Lemma 1 We have \(\phi(x^{k+1},P_{k+1})\leq\phi(x^k,P_k)-\tfrac12\log 2\mathrm{e}\) for all \(k\geq1\).
Proof. Since \(\|AP_{k+1}-I\|_F^2\leq\|AP_k-I\|_F^2-\tfrac{1}{\kappa}r_k^2\), we have
\[ \phi(x^{k+1},P_{k+1})-\phi(x^k,P_k)=\log r_k-r_k^2\leq\max_r\{\log r-r^2\}=-\tfrac12\log 2\mathrm{e}. \]
Complexity. The complexity is straightforward
\[ \mathcal{O}\left(\underbrace{\kappa\|AP-I\|_F^2}_{\text{Learning}}+\log\left(\tfrac1\varepsilon\right)\right) \]
And superlinear convergence follows from \(\sum_{k=1}^K r_k^2\leq\kappa\|AP_1-I\|_F^2\) and the quadratic-mean inequality
\[ \tfrac{\|\nabla f(x^{K+1})\|}{\|\nabla f(x^1)\|}=\prod_{k=1}^K r_k\leq\left(\sqrt{\tfrac1K\sum_{k=1}^K r_k^2}\right)^K\leq\left(\tfrac{\sqrt\kappa\|AP_1-I\|_F}{\sqrt K}\right)^K. \]