← Home

Minimizing matrix-aware regret

Here we look at minimizing regret of gradient descent on a matrix. This interestingly points to a few 'modern' optimizers.

Note: all we are really doing below is taking a matrix view of Adagrad.

Consider gradient descent with preconditioning on convex $f$, $$ X_{k+1}=X_k-\rho\cdot w_kH^{-1}G_k, $$ where $G_k = \partial f / \partial X_k$ and $X_k \in \mathbb R^{m\times n}$. Weights $w_k > 0$ and global scale $\rho > 0$ encode the learning rate. Wlog assume $m \leq n$ (if not, take the transpose). $H \in \mathbb R^{m\times m}$ is a fixed symmetric positive definite preconditioner. Throughout, $\|\cdot\|_2$ denotes the spectral norm. Let $X^\star$ be the optimum.

Goal: we want to find a suitable preconditioner $H$.

Similarly to this analysis, define (weighted) regret after $T$ steps as $$ R=\sum_{k<T}w_k\bigl(f(X_k)-f(X^\star)\bigr). $$ Regret punishes us for being far away from the solution. Since we assume no knowledge of $f$ other than convexity, the best we can do is use the first-order condition for convexity, $$ R\leq\sum_{k<T}w_k \langle G_k,X_k - X^\star\rangle. $$ Perhaps an intuitive way to think about $R$ is by treating $G_k$ as the signal, and $E_k = X_k - X^\star$ as the error to the true solution. If there is large signal and large error, we are doing a bad job!

We have $E_{k+1} = E_k - \rho w_k H^{-1} G_k$, so taking the squared norm of both sides gives $$ \begin{aligned} \|E_{k+1}\|_H^2 =\|E_k\|_H^2-2\rho w_k\langle G_k,E_k\rangle +\rho^2w_k^2\operatorname{tr}(G_k^\top H^{-1}G_k). \end{aligned} $$ Note: use the Mahalanobis norm $\|Z\|_H^2 = \mathrm{tr}(Z^\top HZ)$, so that $-2\langle E_k,\,\rho w_k H^{-1} G_k\rangle_H = -2\rho w_k \langle E_k ,\,G_k\rangle$.

Skipping the details, but rearranging, summing, and telescoping gives $$ R\leq \frac{\|E_0\|_H^2-\|E_T\|_H^2}{2\rho} +\sum_{k < T}\frac{\rho w_k^2}{2}\operatorname{tr}(H^{-1}G_k G_k^\top ). $$ We massage a bit more to get a clean expression in terms of $H$. First, we have $$ \|E_0\|_H^2 = \operatorname{tr}(HE_0 E_0^\top) \leq \|E_0\|_2^2 \operatorname{tr}(H). $$ This inequality comes from the fact that the trace is the sum of the eigenvalues, and $H\mapsto H E_0E_0^\top$ increases each eigenvalue by at most a factor of $\|E_0\|_2^2$. Second, use the fact that trace is linear, defining $C_T = \sum_{k<T} w_k^2 G_k G_k^\top$, giving, $$ R\leq \frac{\|E_0\|_2^2}{2\rho}\operatorname{tr}(H) +\frac\rho2\operatorname{tr}(H^{-1}C_T). $$ Finally, differentiate w.r.t. $H$ and set to zero (use $\mathrm{d}H^{-1} = -H^{-1}\, \mathrm{d}H \, H^{-1}$), $$ \frac{\|E_0\|^2_2}{2\rho} I - \frac{\rho}{2} H^{-1} C_T H^{-1} = 0. $$ This gives, $H = \rho C_n^{1/2} / \|E_0\|_2$. Thus the implied adaptive descent step is, $$ \Delta X_{k+1}=-D w_k\,\bigg(\sum_{i=0}^{k} w_i^2 G_i G_i^\top\bigg)^{-1/2}\,G_k. $$ Inspired by Prodigy/D-adaptation, we write $D = \|X_0 - X^\star\|_2$. This result is interesting because we can compare it to other optimizers: