Gpytorch: Robust model fitting (RuntimeError when learning rate is too high)

[bletham]

Thanks for the continued development on GPyTorch. It's fantastic for constructing models but one area that we've recently been struggling with is robustly fitting models, due to numerical issues during optimization. Even a simple ExactGP model errors out during fitting with a RuntimeError if the learning rate is too large.

This makes model fitting a real challenge: too small of a learning rate and it takes forever, too large and it errors. To get around learning rate tuning, we've been using L-BFGS (both from https://github.com/hjmshi/PyTorch-LBFGS and from Scipy), but even this can fail when the L-BFGS takes too large of a step. We want to run this for a large, heterogeneous collection of problems and so really need robust model fitting that doesn't require hyperparameter optimization.

Attached is a notebook with a simple GP example that fails to fit with default L-BFGS setup, or when the Adam lengthscale is too long. It also includes a plot showing the unfortunate hyperparameter sensitivity of Adam.

I've been in touch with @hjmshi about this. Ideally, the L-BFGS would be given a chance to backtrack from the overly large step by getting a NaN for the function return and leaving the model in a useable state. Is there anything that can be done along these lines to make fitting more robust?

gpytorch_optim_test5.zip

[jacobrgardner]

I actually think it should be fairly straightforward to get NaNs out instead of runtime errors so that you could backtrack. If you think this would be a good idea, we can get this done

[Balandat]

@bletham, @hjmshi, would getting a NaN function value be enough for the algorithms to backtrack? If that's the case and the occurrence of the failure due to overly large steps are relatively consistent, what we could potentially do is add a debug option to the settings, which when enabled allows us to catch the Errors and return a NaN

[jacobrgardner]

@Balandat totally. It may actually be easier than catching the errors: most of the errors that pop up happen downstream from getting a NaN in a result from CG or the preconditioner calculation. We could probably simply detect a NaN in one of those places, and return NaN for the overall loss if we do.

[Balandat]

You guys probably know best where to do these checks and how to do them efficiently - do you mind putting up a PR for this? This would be super helpful.

[hjmshi]

@Balandat Yes, this should be sufficient for L-BFGS to know to backtrack.

[jacobrgardner]

I tested things out using the new param_transform features that #371 will add with a softplus nonlinearity, and I did binary search to find the maximum learning rate I could run Adam with before crashing.

I was able to get Adam working with learning rates up to 267 without crashing using the softplus nonlinearity. I will still look in to returning NaNs in certain cases, but I would say that model training is much more robust with softplus.

Here's the new notebook attached:
gpytorch_optim_test6.zip

[bletham]

@jacobrgardner this is great, thanks a bunch! I'll let you know if I can break it again :-)