What makes a good preconditioner when only a few approximate iterations are needed?

Question

For deterministic solver of $Xw=y$, one recommendation is to pick $P$ such that $P^{-1}X$ has a low condition number. However, this condition only really matters when you want to reduce initial error a lot, ie, by a factor of $10^{-15}$. There are scenarios where you are only need to reduce the error by a factor of $10$. What makes a good preconditioner in such case?

Examples:

• huge linear systems, a randomized iterative solver only touches a fraction of rows of $X$
• non-linear system, linear approximation $X$ becomes irrelevant after couple of steps
• huge non-linear system -- use only a fraction of rows of $X$, and $X$ is changing

The last case is common in machine learning.

"Condition-number" preconditioners focus on performance of the last step of the solver, whereas I care about better performance of the first step, what's a better criterion to use for this case?

It seems that performance of the first step of randomly initialized gradient descent is determined by the following ratio of first and second moments of Hessian eigenvalues $\lambda$: $$\frac{\left(E\lambda\right)^2}{E\lambda^2}$$

A good preconditioner may try to optimize this value directly. Does if this criterion occur in the literature?

Setup for showing result above:

• translation invariance, WLOG $w^*=0$ so $\|\text{error}\|=\|w-w^*\|=\|w\|$
• isotropic initialization, initialize $w_i$ using standard normal
• $w_1=(I - \alpha X^T X)w_0=Tw_0$, estimate after one step of gradient descent with Hessian $X^T X$
• Find $\alpha$ which minimizes norm of expected error after 1 step
• Find expected improvement in error squared using this $\alpha$

and weights initialized with standard normal, we can show that expected improvement in error squared achievable by single step of gradient descent is determined by the ratio of first and second moments of eigenvalues

Using result on expectation of quadratic form on a unit sphere in $d$ dimensions we get expression of error norm after one step

$$E[\|\text{error}_1\|^2]=\frac{1}{d}\text{Tr}(T^TT)$$

Minimize this to get the formulation above

Answer 1

Edit there seem to be a few quantities that predict the difficulty based on how "flat" the spectrum is. Analysis below corresponds to the purity measure of spectrum flatness, but one could also use regular von Neumann entropy.

Sampling random quadratic problems and plotting number of steps to achieve relative target against 1/entropy or 1/purity we see a general correspondence between "flatness" and ease of optimization. (note, decreasing entropy corresponds to increasing purity)

One proxy for a preconditioner may be to look for the quantities which allow for a gradient step with a large step size. The idea is that if the problem is normalized to have $\text{Tr}(H)=1$, ability to take a large step along gradient direction implies that gradient is pointing in the direction of the optimum.

Step step size depends on the direction of true optimum in relation to current position. We can consider the following cases

1. Average case -- every direction is equally likely
2. Worst case -- $w-w^*$ is along the direction of highest curvature of $H$

We can also consider the of stochastic gradient, where direction at current point is based on random subset of all examples.

1. Stochastic -- computing step based on 1 draw from the distribution
2. Deterministic -- compute gradient based on all examples

Taking linear least squares with examples normally distributed, the largest learning rate $\alpha$ which guarantees expected reduction in objective:

1. Stochastic, average case $$\alpha<\frac{1}{\text{Tr}H^2+0.5}$$

2. Deterministic, average case $$\alpha<\frac{2}{\text{Tr}H^2}$$

3. Stochastic, worst case $$\alpha<\frac{1}{\|H\|+0.5}$$

4. Deterministic, worst case $$\alpha<\frac{2}{\|H\|}$$

To summarize:

1. "worst-case" scenario is determined by ratio of largest and average eigenvalue (some kind of condition number?)

2. "average-case" scenario is determined by linear entropy of the eigenvalues, since $\text{Tr}H^2=1-H_L$

Answer 2

Here's a quantity which seems analogous to "condition number" but for $k=1$ regime (instead of $k=\infty$ for condition number)

$$\Delta=\frac{\operatorname{Tr}(H^2)^2}{\operatorname{Tr}(H^3)\operatorname{Tr}H}$$

The way to interpret the quantity: suppose

1. Someone gives us quadratic $H$
2. We pick step size optimally for this task
3. Initialize parameter randomly with starting expected loss=1
4. Do single step of gradient descent.

Expected loss is now $(1-\Delta)$.

This quantity seems to be minimized by a quadratic which puts bulk of the mass on largest dimension and splits the rest equally.