mlr-org / mlr3mbo

Flexible Bayesian Optimization in R
https://mlr3mbo.mlr-org.com
25 stars 1 forks source link

perf: speed up surrogate predictions #173

Open sumny opened 1 week ago

sumny commented 1 week ago

Acquisition function optimization wit a batch size of 1, i.e., as used with any standard sequential numeric optimizer like DIRECT or L-BFGS-B is currently embarassingly slow due to 1) bbotk overhead and especially 2 ) mlr3 overhead (both overhead results from many assertions, checks, etc. that are triggered whenever the prediction method of the surrogate is used and evaluations are logged into the archive and the overhead of the batched evaluation mechanism of bbotk instances and objectives).

For batch sizes larger than 1, this is less problematic.

This PR tries to at least partially improve the surrogate predict overhead arising from 2) for a single predict call by not relying on predict_newdata but skipping some checks and task constructions and directly using predict of the learner(s) wrapped in the SurrogateLearner or SurrogateLearnerCollection.

To further improve upon this, the only option is likely to move away from wrapping LearnerRegr as surrogates but implementing them directly via the base Surrogate class to skip all the mlr3 assertions and checks.

Benchmark showing a median improvement of a factor of around 1.7 compared to current main branch (https://github.com/mlr-org/mlr3mbo/commit/012b60c0f607abee241d8aedaf4b00d264378a74). Note, however, that this is still embarrassingly slow as can be seen when comparing to the time required to make a direct prediction without mlr3 overhead below, where we still observe an overhead of a factor of roughly 15.

fun = function(xs) {
  list(y = xs$x ^ 2)
}
domain = ps(x = p_dbl(lower = -10, upper = 10))
codomain = ps(y = p_dbl(tags = "minimize"))
objective = ObjectiveRFun$new(fun = fun, domain = domain, codomain = codomain)

instance = OptimInstanceBatchSingleCrit$new(
  objective = objective,
  terminator = trm("evals", n_evals = 5))

xdt = generate_design_random(instance$search_space, n = 4)$data

instance$eval_batch(xdt)

learner = default_gp()

surrogate = srlrn(learner, archive = instance$archive)

surrogate$update()

microbenchmark::microbenchmark({surrogate$predict(data.table(x = 1))}, times = 1000L)

old https://github.com/mlr-org/mlr3mbo/commit/012b60c0f607abee241d8aedaf4b00d264378a74

Unit: milliseconds
                                         expr     min       lq    mean   median       uq      max neval
 {     surrogate$predict(data.table(x = 1)) } 15.7425 16.35645 17.4369 16.58944 16.98181 47.88926  1000

new (this PR)

Unit: milliseconds
                                         expr      min       lq     mean   median       uq      max neval
 {     surrogate$predict(data.table(x = 1)) } 8.731744 9.190959 9.905819 9.329769 9.585994 103.8177  1000

direct prediction without mlr3 overhead:

microbenchmark::microbenchmark({predict(surrogate$learner$model, newdata=data.frame(x = 1), type = "SK", se.compute = TRUE)}, times = 1000L, unit = "milliseconds")
Unit: milliseconds
                                                                                                           expr      min        lq      mean   median        uq      max neval
 {     predict(surrogate$learner$model, newdata = data.frame(x = 1),          type = "SK", se.compute = TRUE) } 0.568835 0.6204085 0.7982923 0.625719 0.6349295 166.9255  1000
sumny commented 1 week ago

This PR now also includes an example (SurrogateGP.R) how to maintain surrogate models without directly relying on mlr3 routines to further reduce overhead.

surrogate = SurrogateGP$new(archive = instance$archive)
surrogate$param_set$set_values(
  covtype = "matern5_2",
  optim.method = "gen",
  control = list(trace = FALSE),
  nugget.stability = 10^-8
)

surrogate$update()

microbenchmark::microbenchmark({surrogate$predict(data.table(x = 1))}, times = 1000L, "milliseconds")
Unit: milliseconds
                                         expr      min      lq     mean   median       uq      max neval
 {     surrogate$predict(data.table(x = 1)) } 1.030535 1.07618 1.116669 1.101158 1.122628 4.567874  1000

Likely this is the way to go to at least replace default_gp() and default_rf() in default_surrogate() or have something like default_efficient_surrogate().

sumny commented 1 week ago

Todos: