Optimization with non-linear mean function

[qres]

I am trying to optimize a model with a nonlinear, parametrized mean. E.g. something like $y(x) = f(x,a) + GP(x)$ where we known that the data has some general behavior, $f$ (such as "sin(a*x)") and the residuals are modeled by a zero-mean GP. Here, $a$ is an unknown parameter that has to modeled, too.
Is it possible to optimize such a model in BoTorch? I saw that there is a FullBayesian model based on NUTS for kernel parameters, but I am unsure if it can be used or adjusted for other parameters.

[saitcakmak]

Hi @qres. I think you can handle this case using a custom mean module in the GP. GPyTorch has a LinearMean (https://docs.gpytorch.ai/en/stable/means.html#gpytorch.means.LinearMean), which should support modeling functions like a * x with unknown a and allow training a with other hyper-parameters. To make this into f(a*x) with a known f, you can overwrite the forward call of the model, changing the mean_x = self.mean_module(x) to mean_x = f(self.mean_module(x)).