How to control batch size in fit_gpytorch_mll

[kheyer]

I'm running into an issue where fit_gpytorch_mll is giving me a cuda out of memory error. Is there a parameter to tune batch size or control memory usage?

train_Y = ...
train_X = ...

model = SingleTaskGP(train_X, train_Y, input_transform=Normalize(d=train_X.shape[-1]).to(device))
mll = ExactMarginalLogLikelihood(model.likelihood, model)
fit_gpytorch_mll(mll)

[Balandat]

How large is the data size in this case? How much memory does your GPU have / is anything else running on that device?

Note that fit_gpytorch_mll by default uses L-BFGS-B to fit the model, so there is no minibatching going on. In general it's not obvious how to do this for an exact GP (this is more an approach you'd use for SVGPs).

[kheyer]

I'm running an iterative process to find "high scoring" items in a dataset

• Train model
• Sample from model
• Sample similar items from database
• Score items
• Repeat

Due to the iterative process, the training dataset grows over time.

I hit a memory error on ~20k inputs on a GPU with 24GB memory. The vectors in train_X are 768 in size.

Is there a better approach in botorch for the sort of thing I am trying to do? Especially given I would want it to scale to ~100k examples

[Balandat]

When you say "Train model" and "Sample from model" you are talking about the BoTorch surrogate model, right? How long does it take to "sample similar items from database" and "score items"?

I hit a memory error on ~20k inputs on a GPU with 24GB memory. The vectors in train_X are 768 in size.

So train_X in your case is of shape 20k x 768?

Is there a better approach in botorch for the sort of thing I am trying to do? Especially given I would want it to scale to ~100k examples

The number of data points you are considering here is outside of the range where a standard exact GP model (as SingleTaskGP is) reasonably scales to (which are usually maybe a few thousand). If you do want to use a GP, you should consider using an approximate GP. We do have a SingleTaskVariationalGP that you could try for this in order to deal with the number of data points.

That said, if the dimensionality of each data point is indeed 768 that's also very large for a GP (or more generally kernel-based methods), and you'd probably need some custom approaches to get this to work well.

In short, it unfortunately does seem like your use case is one that may not be well served by BoTorch and you may want to consider other "active learning" methods that are targeted to this problem and scale better.

[kheyer]

The database in question is a vector database, so it's extremely fast to query nearest neighbors. However, the database is very large. We were using active learning, but the cost of repeated screening of the full database with the proxy model is prohibitive for where we want to go in terms of scale.

We're interested in using some finite set of scored data points to learn a prior distribution over the input embeddings to remove the need to extensive screening of the full database. The goal is to generate candidates by sampling from the learning prior and using those samples for nearest neighbor queries. Have you seen any approaches along those lines?