2
$\begingroup$

I train a neural network - one of the Resnet variations ($\approx 10^7$ parameters) on the CIFAR-10 dataset - and after each epoch, I would like to find the smallest/largest eigenvalues of its Hessian. Formally, we can view our NN as a function: let our dataset (training set) be $D$, let $x$ be NN parameters, and we want to minimize

$$f(x) = \frac 1 {|D|} \sum_{d \in D} \text{loss}(x, d)$$

Now, we consider Hessian of this function $$\nabla^2 f(x) = \frac 1 {|D|} \sum_{d \in D} \nabla^2 \text{loss}_{xx}(x, d)$$

For a given NN parameters $x$, I want to find the smallest and largest eigenvalues $\nabla^2 f(x)$. For that, I can use hessian-vector products, i.e. for any vector $v$ I can compute $\nabla^2 f(x) \cdot v$ (PyTorch has a built-in mechanism to compute $\nabla^2 \text{loss}_{xx}(x, d) \cdot v$ for any batch $d$), so, for example, I can use the power method. However, the pure power method will be very slow: one needs a pass over the entire data $D$ to compute the true Hessian-vector product.

Question: do you know of an efficient algorithm for this task? To be precise, I mean that both eigenvalues can be computed with a reasonable multiplicative error within at most 1 minute each. (The run aside from that takes 5 hours and I need to compute both eigenvalues about $40$ times per run, so it's the time I'm ready to spare. There are a lot of runs, so this is my limit.)

Note that I'm asking about an algorithm that you know to be efficient for this (or similar) problem. I tried the power method, accelerated power method, stochastic power method, Oja's algorithm, gradient-based algorithm, its accelerated version, algorithms from https://arxiv.org/abs/1707.02670. All these experiments take a lot of time and so far I didn't have any success, no matter how much engineering I used. When eigenvalues are close to $0$ (e.g. of order $-\frac 12$, when the largest eigenvalue is of order $100$), either convergence takes a lot of time or the results are unstable/unreliable. Just in case, I'm aware of PyHessian (and the first version of my code is based on theirs).

My current best solution is the one that performs the power method with increasing batch size. In a good scenario, it looks like this: enter image description here enter image description here enter image description here The left plot shows eigenvalues, the right plot shows what I call "error": $\|v - \frac {H v} {v^\top H v}\|$, where $Hv$ is a hessian-vector product computed at the last iteration, and $v$ is the current eigenvector candidate of norm $1$. Red lines separate cases when I dramatically increase batch size: I first use a lot of iterations with a small batch size (64 images), then 1000 images per batch, and in the end, I increase the batch size until it reaches the whole dataset. It can be seen that the power method with a smaller batch size always has a limit beyond which it can't grow.

The problem is that this approach takes $4$ minutes for one eigenvalue - $4$ times more than what I'm ready to spare. Decreasing the number of batches at every step is detrimental to quality when eigenvalues are small. If the eigenvalue is closer to $0$, it already doesn't finish converging to it with this number of iterations: enter image description here

$\endgroup$
2
  • $\begingroup$ Thank you! I understand you're only looking for things we know will work, not ideas for something to try, so this isn't what you're looking for, but: have you tried a SGD version of power iteration? (i.e., GD : SGD as power iteration : it) In other words, in each iteration you pick a batch $B$ that is a randomly chosen subset of the training set, compute $1/|B| \sum_{d \in B} \nabla^2 \text{loss}_xx(x,d) \cdot v$, and use that as an estimate of $\nabla^2 f(x) \cdot v$ in your power iteration. Does that do anything useful? I apologize if this is not helpful. $\endgroup$ – D.W. Jan 21 at 6:56
  • $\begingroup$ @D.W., I've edited my post to show my best solution. It includes what you suggest. I've also noticed that I specified the wrong computation time: it should be 2 minutes for both eigenvalues instead of 10. $\endgroup$ – user114966 Jan 21 at 8:58

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy