Gpytorch: Predictions flipping sign when calling backward with model in eval mode

[darbour]

I am working through a toy problem that attempts to maximize the posterior mean w.r.t. the input parameters. However it seems that after putting the model into evaluation mode, calling backward changes the sign of the gradient w.r.t. the input (it will also change the sign of the predictions). I've attached a NB to reproduce.
sign_bug.pdf

[Balandat]

@gpleiss, @jacobrgardner this seems very problematic for any kind of BayesOpt functionality - any idea what could be going on here?

[Balandat]

@bkarrer

[jacobrgardner]

I'll take a look. The first backward looks like it worked -- are things fixed by zeroing out the gradients between backward calls?

[Balandat]

No, running model.zero_grad() between the backward calls doesn't help. The interesting thing is that the sign of the gradient alternates between calls.

[jacobrgardner]

Took a look. 0eb14c206d5c70503229571d3e5fe302373080d0 should fix this. We were doing inplace operations in the backward pass of inv_matmul on some things we cache, which meant that calling backward actually ruined the entire model, not just derivatives. Yikes!

I don't think this would have shown up before recent PyTorch changes, because there's no reason a priori to expect we would ever take derivatives through the caches (e.g., you'd need to be doing something strange like taking derivatives with respect to the data the model was trained on). To be honest, I'm still not entirely sure why we'd be entering that backward pass -- in the example notebook, none of the training data is set to requires_grad=True.

[jacobrgardner]

Closing this for now since running the supplied code gets identical gradients every time for me now.

[Balandat]

@jacobrgardner Hmm so you think you shouldn't enter that backward pass even when taking the gradient of some function of the output w.r.t. the model input? That's what the notebook does and that's what you'd be doing in any gradient-based optimization of an acquisition function.

[jacobrgardner]

@Balandat I agree that taking the derivative with respect to the __test__ input is expected and necessary for our applications.

The predictive mean looks something like k'K^{-1}y. We compute and store K^{-1}y as self.mean_cache. What is confusing to me is why we'd enter the backward pass of InvMatmul (the function we'd call to get K^{-1}y) at all to get the derivative w.r.t. x, since nothing about self.mean_cache depends on the test data we want the derivative for. Is it because the noise term (the sigma in K+sigma*I) is still a model parameter, so we always get derivatives for it regardless of whether those are the ones we want? If so, we should be able to make derivatives for acquisition functions much more efficient by not computing those derivatives.

Either way, doing in place operations there is clearly wrong so the commit I made 0eb14c2 fixing the actual issue raised here is a necessary change.

[Balandat]

Hmm that's am interesting point, we'd definitely want to avoid computing unnecessary stuff. Let me see if I can figure out what's going on in this toy example.

Would you mind creating a follow-up issue so we can track this?