Support categoricals in alternating optimization

[TobyBoyne]

Motivation

See issue #2864

Have you read the Contributing Guidelines on pull requests?

Yes

Test Plan

There are already tests in place for mixed integer feature spaces. I have created a new test case for mixed categorical feature spaces. Once the overall approach has been reviewed, I will also add tests for the new functions that have been introduced

Related PRs

None yet - happy to update the docs.

[TobyBoyne]

Some questions I have about my implementation:

1. What should be the expected behaviour if the search space is not mixed (eg. everything is continuous, or everything is categorical)? I have added in a few checks that were necessary for spaces that are categorical+continuous, but no discrete, but I'm not sure which cases need to be well defined.
2. Should I introduce any additional keys to options to allow for more fine-grained control?
3. What should be the fallback when there are too many categories? For discrete ordinal features, continuous relaxation is in place.
4. I found some issues when using this with the OneHotToNumeric transform (see the FIXME comment in optimize_mixed. I would expect get_X_baseline to return the untransformed (ie. one-hot) inputs, but it seems to return the transformed

[Balandat]

What should be the expected behaviour if the search space is not mixed (eg. everything is continuous, or everything is categorical)? I have added in a few checks that were necessary for spaces that are categorical+continuous, but no discrete, but I'm not sure which cases need to be well defined.

I think if either it's fully continuous or fully discrete we should just raise an error. We have other acquisition optimizers for that.

[Balandat]

Should I introduce any additional keys to options to allow for more fine-grained control?

Which ones would those be? In general I think unless we have a need for them we can keep things simple and add them later if needed.

What should be the fallback when there are too many categories? For discrete ordinal features, continuous relaxation is in place.

Good question. If we can't enumerate them all, then a natural thing to do would be to sample a subset in each dimension in each step rather than computing all.

[Balandat]

Hmm interesting. I am not super familiar with this part of the code but it seems that this handling here is not working the way it should: https://github.com/pytorch/botorch/blame/d247a33486c22bc5c6f16107db3e71f1e3bf3acd/botorch/optim/utils/acquisition_utils.py#L126-L129

Looking through the code, it seems that _has_transformed_inputs and _has_transformed_inputs is used very sparingly. This requires some more thorough audit, I think.

The solution you have here will not work if the model doesn't actually have an input_transform attribute.

[TobyBoyne]

My solution was very much a temporary band-aid for something I was trying to run. I'm happy to remove this line, as I agree that it won't work if there are no input_transforms (and I think it might apply some input transforms twice). I will leave a TODO in the code noting that this sometimes breaks with bad input transforms.

[Balandat]

You should check if opt_inputs.acq_function.model has an input_transform attribute and do nothing if it doesn't.

Have you tried this with some other input transforms? Is this only an issue with OneHotToNumeric?

[TobyBoyne]

I've investigated further and it seems to not be a problem on BoTorch's end - when I make models only in BoTorch, I can't recreate the error with any kind of transform (OneHotToNumeric included). It's only when I'm using external code that I seem to get these issues.

I'll keep looking into this separately, but the issue is not with this PR (or even with BoTorch).

[Balandat]

cc @saitcakmak who may be interested in this feature

[TobyBoyne]

I think if either it's fully continuous or fully discrete we should just raise an error. We have other acquisition optimizers for that.

I agree, although making this change does break this existing test, and therefore might break expected behaviour for some BoTorch users?

botorch/test/optim/test_optimize_mixed.py

Lines 535 to 547 in d247a33

	# Only continuous parameters should fallback to optimize_acqf.
	with mock.patch(
	f"{OPT_MODULE}._optimize_acqf", wraps=_optimize_acqf
	) as wrapped_optimize:
	optimize_acqf_mixed_alternating(
	acq_function=acqf,
	bounds=bounds,
	discrete_dims=[],
	options=options,
	q=1,
	raw_samples=20,
	num_restarts=2,
	)

[Balandat]

I agree, although making this change does break this existing test, and therefore might break expected behaviour for some BoTorch users?

Good point - let's just fall back to to the standard _optimize_acqf in the same way in this case:

botorch/botorch/optim/optimize_mixed.py

Lines 679 to 680 in d247a33

	if len(discrete_dims) == 0:
	return _optimize_acqf(opt_inputs=opt_inputs)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (d247a33) to head (0d8e128).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##              main     #2866   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          211       211           
  Lines        19353     19396   +43     
=========================================
+ Hits         19353     19396   +43     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Balandat]

This looks great, thanks a lot for the high quality contribution!

I left a few nits inline; the main gap is to add a test for the case of too many categories (where we sample instead).

[TobyBoyne]

I think that should be everything, @Balandat! Thanks for the quick and thorough review!

[facebook-github-bot]

@Balandat has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[saitcakmak]

Looks great! Thanks for implementing, this has been on my wish list for a while :)

[facebook-github-bot]

@Balandat merged this pull request in a268631.