Why aren't Krylov subspace methods popular in the Machine Learning community compared to Gradient Descent?

Question

Historically, iterative methods for solving relatively simple-structured systems $Ax=b$ with $A$ being a $4\times 4$ matrix or to find the eigenvalues of that matrix assuming in both problems that $A$ is non-singular were not quite popular and were not needed in pre-1950s as stationary methods or perhaps classical iterative algorithms such as Jacobi, Gauss-Seidel would handle the task easily.

Then in the $1950$s the rapid development of computational tools made it possible to deal with large and sparse matrix system (possibly a $100\times 100$ at that time in engineering applications). This is all thanks to the development of Krylov Subspace tools which at that time was considered to be "high-tech". This is seen from various celebrated algorithms such as Arnoldi Iteration (Arnoldi-Lanczos 1950) and Conjugate Gradient (Stiefel,1952). During the 1970s and 1980s when we had tons and tons of transistors inside computers this made it possible for us to witness the birth of GMRES and BiCGSTAB which are known to handle relatively large and sparse systems.

In fact, it was during that period that BLAS and LAPACK were developed and ever since have continued to upgrade various Krylov algorithms and achieve a high level of robustness.

In machine learning, many problems such as regression, neural networks, and even recommender systems can one way or another be formulated as an optimization problem which in turn can be expressed as a linear system $Fx=z$. The apparent default choice for solving such problems is the "Gradient Descent Algorithm" (and its derivatives) a popular optimizer algorithm that seems to be used literally everywhere in machine learning problems (movie recommender system, linear regression,...) even a famous variant is the ADAM Optimizer which in many deep learning libraries is set as the default optimizer.

Given that numerical linear algebra libraries have been for the past 40 years enhancing Krylov algorithms, why haven't they made an impact in today's world of machine learning as they did in engineering communities?

Let's say linear regression is the underlying machine learning problem, why is the gradient descent algorithm preferred over using BiCGSTAB with some restart when rewriting the regression problem as $Fx=z$ where $F$ could be sparse and $10000\times 10000$?

I believe one answer (though it may be biased) is that the majority of the machine learning community is not paying attention to the computational mechanism of the algorithm. For instance, you can go ahead and apply ADAM optimizer to solve a simple regression problem (overkill) or a complicated large regression problem and may end up witnessing your algorithm diverging without properly understanding why. I found a meme from medium.com that jokingly describes this:

Source: https://medium.com/analytics-vidhya/mechanism-of-gradient-descent-optimization-algorithms-7baed9c9c35e

Update 1: I would like to thank davidhigh for the detailed answer. I would like to add that my experience with Krylov methods were mainly focused on modelling dynamical system. Tim Davis has a very interesting collection of sparse matrices with detailed information on where the problem arises from (in particular engineering applications) Krylov methods are highly applicable there.

Update 2: Neil Lindquist made an important correction regarding BLAS and LAPACK not having Krylov solvers.

Final Remark: I welcome any additional answer especially if some users have experience with using Krylov methods in ML applications.

Answer 1

On a basic level, I don't buy the argument that you have to "solve a linear system for many machine learning algorithms". Much more, you usually have to optimize a non-linear equation which is obtained from a given functional in order to derive the parameters of the applied model.

Thus, the basic problem is, that you have a functional depending on $N$ parameters as well as on the training data,

$$ \min_\alpha L(\alpha_1,\ldots, \alpha_N; \text{data}) $$

How do you solve these kind of problems? In a not all too satisfying way: you pick a start configuration for your parameters, and then go "downhill". That's basically gradient descent, which ADAM is a sophisticated variant of.

As you say, you can formulate this problem as a linear equation. But the whole benefit of this linear system is to derive an "optimal way downhill". And, remember, as the functional is non-linear, this optimal direction holds only for the current parameter configuration (you only go along the current tangent space of the manifold). So you'd have to solve a linear system in each step.

Now, it turned out that the benefit in using the optimal direction at each step is not worth the overhead of solving a linear system. So people still simply go downhill. Stochastic gradient descent and ADAM are more sophisticated approaches based on the same idea.

If you have particular applications in mind where you would expect a linear system to be advantageous, please feel free to add it to the question pr in the comments.

Answer 2

They aren't popular because they don't work.

Nicol N. Schraudolph spent a few years on Krylov-like methods for Machine Learning. I first learned of his work at 2004 Machine Learning Summer School in Canberra.

Looking at his later work (A Stochastic Quasi-Newton Method for Online Convex Optimization) you can find the following explanation:

More generally, thousands of people have tried to improve on SGD+momentum in Machine Learning setting, without success. Adam is slightly better on many tasks, but slightly worse on others, like Imagenet training.

Personally I've spent significant time working on ways of improving SGD for ML (contributing to tools like this), and came to believe that for optimized architectures like Resnet-50, optimally preconditioned direction would only give a modest improvement, perhaps $\approx$ 5x reduction in terms of the number of steps needed. This means if your advanced method does more than "5 SGD steps" worth of work per step, it'll lose to SGD. Additionally there's a storage constraint which means you can't store more than a couple gradients worth of statistics throughout training. Models are typically scaled up until gradient computation exceeds GPU memory constraint.

General reasons why it's hard to improve on SGD in ML. Typical problems are:

• overparameterized and very high-dimensional. SGD in this setting is already pretty good, converging at an exponential rate -- https://arxiv.org/abs/1811.02564

• non-stationary. Local curvature statistics get stale after a few of steps.

• not optimized to convergence. Many "advanced methods" are theoretically superior to SGD only in the "tail regime", but we don't get there. For instance, in BigScience model training, the loss went down from $\approx 10$ to 0.13 in the last 12 months of training. So we care about the time it takes to decrease loss by a factor of 100. Meanwhile second-order methods like Newton method starts to pull away in the tail regime, going from $10^{-5}$ loss to $10^{-10}$.

• tailor made for SGD. If SGD fails to achieve good results, corresponding network architecture won't see the light of day. For instance, autoencoder with a bottleneck layer is near impossible to train with SGD, but easy with a factored Gauss-Newton method. Since advanced optimizers aren't accessible in ML frameworks, such architecture wouldn't get a shot at evaluation.

An example of method which satisfies these constraints is PSGD which:

1. stores preconditioner $P$ in a factored form, in extreme case, parameterizing it by a single number recovers SGD with "automatic step size"
2. updates $P$ in an online fashion by taking cheap SGD steps on $P$ to keep $P^{-1}$ close to Hessian $H$, using Hessian-vector products.

(I had no hand in PSGD, other than having evaluated it and found it to be the most promising out of all second-order methods I've used. I think it didn't get attention like KFAC or Shampoo because it didn't come from a famous ML lab)

Answer 3

Posts 3 in the following series on "regret minimization" talks about gradient descent, and the later posts come back to that topic from time to time.

• Online Optimization Post 1: Multiplicative Weights
• Online Optimization Post 2: Constructing Pseudorandom Sets
• Online Optimization Post 3: Follow the Regularized Leader
• Online Optimization Post 4: Regularity Lemmas
• Online Optimization Post 5: Bregman Projections and Mirror Descent
• Online Optimization Post 6: The Impagliazzo Hard-Core Set Lemma
• Online Optimization Post 7: Matrix Multiplicative Weights Update

One conclusion of those posts is that gradient descent can be seen as a regret minimization algorithm in the “Follow the Regularized Leader” framework.