Improving the outer loop: Bayesian optimization where we change the objective function

[bwpriest]

We are beginning to explore models that change the structure of the data as we optimize, such as the anisotropic distortion and deep embedding kernels. We are also interested in exploring the effects of mini-batching. However, our current optimization chassis does not allow for this. Let's discuss how to generalize the optimization chassis, starting with the Bayesian optimization routine, to support this behavior.

[bwpriest]

I think that we should start with a simple refactor of optimize_from_tensors that resamples batches periodically and computes new tensors. Here is some pseudocode:

def optimize_minibatching(
    muygps: MuyGPS,
    train_features: mm.ndarray,
    train_responses: mm.ndarray,
    nhbr_lookup: NN_Wrapper,
    loss_method: str = "lool",
    obj_method: str = "loo_crossval",
    sigma_method: Optional[str] = "analytic",
    verbose: bool = False,
    **kwargs,
):
    #  unchanging component parts of the obj_fn
    loss_fn = get_loss_func(loss_method)
    kernel_fn = muygps.kernel.get_opt_fn()
    mean_fn = muygps.get_opt_mean_fn()
    var_fn = muygps.get_opt_var_fn()
    sigma_sq_fn = make_sigma_sq_optim(sigma_method, muygps)

    #  bayes_opt kwargs boilerplate
    x0_names, x0, bounds = _get_opt_lists(muygps, verbose=verbose)

    x0_map = {n: x0[i] for i, n in enumerate(x0_names)}
    bounds_map = {n: bounds[i] for i, n in enumerate(x0_names)}

    optimizer_kwargs, maximize_kwargs = _bayes_get_kwargs(
        verbose=verbose, **kwargs
    )

    # The list of points to probe
    to_probe = [x0_map]

    for # some number of iterations or other loop condition:
        #  sample a new batch of points
        batch_indices, batch_nn_indices = sample_batch(
            nbrs_lookup, batch_count, sampler.train_count
        )
        #  collect tensors
        (
            batch_crosswise_diffs,
            batch_pairwise_diffs,
            batch_targets,
            batch_nn_targets,
        ) = make_train_tensors(
            batch_indices,
            batch_nn_indices,
            train_features,
            train_responses,
        )
        #  ignore hierarchical stuff for now 
        batch_features=None
        
        #  make the obj_fn
        obj_fn = make_obj_fn(
            obj_method,
            loss_method,
            loss_fn,
            kernel_fn,
            mean_fn,
            var_fn,
            sigma_sq_fn,
            batch_pairwise_diffs,
            batch_crosswise_diffs,
            batch_nn_targets,
            batch_targets,
            batch_features=batch_features,
        )

        #  make the optimizer
        optimizer = BayesianOptimization(
            f=obj_fn,
            pbounds=bounds_map,
            **optimizer_kwargs,
        )
    
        #  probe all traversed points
        for point in to_probe:
            optimizer.probe(point, lazy=True)
    
        #  do the maximization for some number of iterations
        optimizer.maximize(**maximize_kwargs)

        #  figure out how to add all explored points to `to_probe` 

    return _new_muygps(muygps, x0_names, bounds, optimizer.max["params"])

We can explore how to improve upon this once we've done some experimentation.

[akilandrews]

Here is a link to the fork containing the new numpy bayes opt chassis feature optimization-chassis. I created a test using the univariate regression tutorial here MuyGPyS/optimize/experiment/univariate_regression_with_opt_loop.ipynb. This test uses the new chassis and runs bayes optimizer once and successfully produces the same outcomes as the original tutorial. Next is to code for multiple epochs/batches by probing all previous points. Can we brainstorm at our next meeting on how best to test outcomes of multiple iterations and how to compute optimal sigma squared?

[akilandrews]

Ben, I have committed the baseline implementation of feature optimization-chassis and after I complete the unit test the fork should be ready for a pull request. Running the MuyGPyS/optimize/experiment/univariate_regression_with_opt_loop.ipynb notebook outputs:

1. Using a GP prior with 5001 observations, 501 train(10%), and sample batch size = train data size = 501 (uniform random, no replacement). Learning nu with num_epochs = 1 yielded:

   epoch probe point
   ––––– –––––––––––
   0     0.4935585846939505
   loss_method 'mse' nu 1.6109672778498811 sigma_sq [0.26936798]

   epoch probe point
   ––––– –––––––––––
   0     0.4935585846939505
   loss_method 'lool' nu 1.8273947060113152 sigma_sq [0.42540414]

2. Using a GP prior with 5001 observations, 501 train(10%), and sample batch size 100 (uniform random, no replacement). Learning nu with num_epochs = 5 yielded:

   epoch probe point
   ––––– –––––––––––
   0     0.4935585846939505
   1     1.6109672778498811
   2     1.9120931160699903
   3     1.5814296058960151
   4     1.5814296058960151
   loss_method 'mse' nu 1.5814296058960151 sigma_sq [0.24238581]

   epoch probe point
   ––––– –––––––––––
   0     0.4935585846939505
   1     1.8174031964102773
   2     1.8098387058257817
   3     1.8174031964102773
   4     1.8617823763498118
   loss_method 'lool' nu 1.8636020080018265 sigma_sq [0.47303099]

3. Inference yielded:

   bayes mse results:
   RMSE: 0.0006093192113339855
   mean diagonal variance: 3.0570610527356334e-06
   mean confidence interval size: 0.006827740687157497
   coverage: 1.0

   bayes lool results:
   RMSE: 0.0009119032442528957
   mean diagonal variance: 3.4548173618663717e-06
   mean confidence interval size: 0.007285391329334468
   coverage: 0.9991111111111111

[bwpriest]

@akilandrews please open a pull request targeting LLNL/develop, and we can iterate from there.

[akilandrews]

Done. And now I'll move on to the next tasks which include looking into updating the objective function for a BayesianOptimization instance and also creating a Anisotropic GP prior similar to the class UnivariateSampler for our simple case experiments.

[akilandrews]

@bwpriest I added the flake8 linter to my Visual Code. I ignored E501(line length) and E203(space before colon), the latter conflicted with black formatter. The final set of errors appear to be associated with my changes to MuyGPyS/_src/optimize/chassis/init.py. I removed the violating includes that caused fails on mpi, torch, and docs test. I added those directly to MuyGPyS/optimize/chassis.py. I reran my experiment univariate notebook test successfully and the Abseil unit test optimize.py:
Ran 15 tests in 776.448s

OK

Can we run another develop test on the pull request to check for remaining errors? Thanks!

[akilandrews]

Optimization loop experiments from last week using similar workflow from "Anisotropic Metric Tutorial":

1. Validation with 1 epoch reproducing Anisotropic tutorial values:
   0, {'length_scale0': 0.09718538289820054, 'length_scale1': 0.42218698955030526}
   BayesianOptimization finds an optimal:
   length_scale0 is 0.18700151892994976
   length_scale1 is 0.978106563967657
   sigma_sq is 0.8694942754569579

2. Test with 5 epochs, batch size 55, train features 277:
   0, {'length_scale0': 0.09718538289820054, 'length_scale1': 0.42218698955030526}
   1, {'length_scale0': 0.09718538289820054, 'length_scale1': 0.42218698955030526}
   2, {'length_scale0': 0.09718538289820054, 'length_scale1': 0.42218698955030526}
   3, {'length_scale0': 0.14563425651928213, 'length_scale1': 0.8904163963242121}
   4, {'length_scale0': 0.09718538289820054, 'length_scale1': 0.42218698955030526}
   BayesianOptimization finds an optimal:
   length_scale0 is 0.09362294362728572 (pe 6.4%)
   length_scale1 is 0.7709756106408624 (pe 54%)
   sigma_sq is 0.8971665900856329

3. Repeat (2) and scale training features using learned length scale hyperparameters:
   0, {'length_scale0': 0.09718538289820054, 'length_scale1': 0.42218698955030526}
   1, {'length_scale0': 0.19506841348277287, 'length_scale1': 0.961173054024932}
   2, {'length_scale0': 0.1756736822955853, 'length_scale1': 0.9826535449833492}
   3, {'length_scale0': 0.1756736822955853, 'length_scale1': 0.9826535449833492}
   4, {'length_scale0': 0.09147817411853672, 'length_scale1': 0.5273274264558154}
   BayesianOptimization finds an optimal:
   length_scale0 is 0.09147817411853672 (pe 9.0%)
   length_scale1 is 0.5273274264558154 (pe 5.5%)
   sigma_sq is 0.773202746190502

[akilandrews]

Min, here are the results from the experiment we brainstormed. I went back and refactored to use a single initial probe point or all previous probe points. The results improved. I would like to discuss the improvement, also to ensure I have the correct treatments, and a quick code review to make sure the workflow in the opt loop chassis is correct.

Treatment 1 (baseline) - brute force workflow following Anisotropic Metric Tutorial using single optimization step with a single initial probe point:

• train size 277, batch size 277, num epochs 1
• Matched all tutorial values!

Treatment 2a (opt loop) each loop iteration create new optimizer instance, initialize with new objective function, and do not probe previous points:

• train size 277, batch size 55, num epochs 5
• BayesianOptimization finds an optimal:
• length_scale0 is 0.1259 p.e. 26%
• length_scale1 is 1.0000 p.e. 100%
• sigma_sq is 0.8418308639089281
• RMSE: 0.07329470871496807
• mean diagonal variance: 0.006280181480898618
• mean confidence interval size: 0.29292794259142485
• coverage: 0.9723141739392116

Treatment 2b (opt loop) each loop iteration create new optimizer instance, initialize with new objective function, and probe previous points:

• train size 277, batch size 55, num epochs 5
• BayesianOptimization finds an optimal:
• length_scale0 is 0.0915 p.e. 9%
• length_scale1 is 0.5273 p.e. 5%
• sigma_sq is 0.773202746190502
• RMSE: 0.07329470871496807
• mean diagonal variance: 0.0057682056762067896
• mean confidence interval size: 0.28073403842136685
• coverage: 0.9662955160999097

Treatment 2c (opt loop) each loop iteration reuse optimizer instance, set BayesianOptimization.TargetSpace.target_function to new objective function, and do not probe previous points:

• train size 277, batch size 277, num epochs 5
• BayesianOptimization finds an optimal:
• length_scale0 is 0.1167 p.e. 17%
• length_scale1 is 0.6670 p.e. 33%
• sigma_sq is 0.8243962233351186
• RMSE: 0.07329470871496807
• mean diagonal variance: 0.006150116509950241
• mean confidence interval size: 0.28987874738309716
• coverage: 0.9702076436954559

[bwpriest]

Here's an idea for a brute force comparison:

def make_brute_force_obj_fn(train_features, batch_indices, targets):
    def brute_force_obj_fn(**kwargs):
        #  make nearest neighbors object using train_features / length_scales (from kwargs)

        #  query NN_Wrapper for batch_nn_indices
        
        # (...)  = make_train_tensors(...)
    
        # make kernels

        # make predictions

        return loss_fn(predictions, targets, variances, sigma_sq)
    return brute_force_obj_fn

[akilandrews]

@bwpriest Below are the results after merging in the new 2d univariate sampler and bug fix. I also increased the dataset size and ran on a compute node at UNM CARC for evaluating loop execution time:

treatment	time (s)	train size	batch size	num epochs	pts probed	opt steps	l.s.0 (p.e.)	l.s.1 (p.e.)	sigma sq
single step	53.131050	1000	1000	1	6	30	0.1431 (43%)	0.7245 (45%)	2.6585
single step+	346.979755	1000	1000	1	26	150	0.1172 (17%)	0.6025 (20%)	1.5988
no probe	143.653934	1000	200	5	26	150	0.1221 (22%)	0.5839 (17%)	1.5732
probe all	152.327901	1000	200	5	40	150	0.0940 ( 6%)	0.5106 ( 2%)	0.9837
reuse opt	342.288601	1000	1000	5	26	150	0.1150 (15%)	0.6017 (20%)	1.5444

treatment	RMSE	mean diag var	mean c.i. size	coverage
single step	0.04136	0.00177	0.14078	0.94578
single step+	0.04118	0.00187	0.14503	0.95222
no probe	0.04124	0.00180	0.14282	0.95056
probe all	0.04107	0.00204	0.15118	0.95822
reuse opt	0.04119	0.00187	0.14467	0.95100