Nesting GPyTorch into another NN

[isaactian]

I want to attach GPyTorch to the end of a feature extractor like specified in deep kernel regression:

https://docs.gpytorch.ai/en/stable/examples/06_PyTorch_NN_Integration_DKL/KISSGP_Deep_Kernel_Regression_CUDA.html

In the example a feature regression is nested into a ExactGP object. Is the reverse possible? I.e. nest the GPyTorch model into a feature extractor network to get the same results. Thanks.

[Balandat]

I.e. nest the GPyTorch model into a feature extractor network to get the same results

What do you mean by "get the same results"? What are you trying to achieve? Use a GP model as a differentiable component inside some feature extractor / more genera NN architecture?

[isaactian]

Exactly, my intention is to have this architecture:

NN_feature extractor -> GPyTorch

The linked tutorial appears to do exactly that and it has this structure, simplified:

ExactGP model {
feature_extractor
Mean_module
Covar_module
...
Return multivariatenormal(mean(feature_extractor(x)), covar(feature_extractor(x))
}

Can you flip the nesting and train an equivalent network? Such that the GPyTorch model is a component of some wrapping feature extractor and not the other way around.

SomeGenericNN {
Layers...
GPRModel
Return(GPRModel(layers(x))
}

[Balandat]

Can you flip the nesting and train an equivalent network?

Yeah you could do that. I guess the main question is how exactly and why you'd want to do that. The main benefit of a GP is its generally good uncertainty quantification - so you can generate predictive uncertainty estimates. How are you intending to utilize that (and if not, why do you want to use a GP)?

[isaactian]

I'm starting from a point where I have an autoencoder extracting some features from a 3D Mesh, which I then separately train a GP on using SKLearn to predict desired target variables (in this case, body composition metrics from DXA scans). This pipeline performs really well with good precision and accuracy.

We wanted to investigate if it was possible to connect these two networks end to end and train them jointly instead of separately. We initialized both submodules to the separately learned weights. When I train like this however the network does not predict the training or test targets accurately at all (I get negative correlation with training data so it's not overfitting, just blowing up into nonsense). I thought maybe there was some problem with how the mll loss was backpropagated to the feature extractor layers in this construction.

I tried it the way specified in the documentation where I nest the feature extractor network inside the GP and it crashes due to running out of GPU memory. I guess it's trying to build some obscenely large covariance matrix accounting for all of the parameters of the mesh convolutional layers? Basically the goal is to turn two separately trained modules (convolutional AE and GPR from latent features) into a single end to end network. If you have any alternative recommendations on how to get this to work I'd greatly appreciate it.

[Turakar]

I guess the feature extractor in GP approach is better than the GP in feature extractor approach, as forward() can then also return the confidence estimates as a MultivariateNormalDistribution. Also, if you reverse the order, why should the memory usage decrease? An exact GP requires the computation of the matrix-vector-product with the kernel matrix, and in a simple implementation, this requires the full kernel matrix which has a size of n^2.