Vizier output transforms

[CompRhys]

Motivation

I wanted to experiment with the vizier output transforms in botorch - in particular the infeasible transform. Potentially misplaced effort as it appears that within the meta ecosystem Ax actually handles this sort of stuff and so these may not actually end up being able to used very much but worth upstreaming.

takes transforms from: https://arxiv.org/abs/2408.11527

Test Plan

Added tests for all functionality, not sure I am at 100% coverage

Related PRs

this branch takes off from #2636 so that would need to be merged first.

[saitcakmak]

Thanks for adding these. I've been curious to try out the Vizier transforms. These transforms try to ensure we model the good regions better, which seems better than trying to model everything equally well, when the goal is to optimize a function. Have you tried these on any benchmarks yet? It'd be good to see how they work e2e.

None of these transforms implement the untransform_posterior method, which is necessary to make sure the transforms work as an outcome_transform in BoTorch models. Without it, Model.posterior will error out. Would it be possible to implement untransform_posterior, likely returning a TransformedPosterior object.

None of the transforms seem to support Yvar either. We use BoTorch through Ax and try to keep the same methods regardless of whether the noise is observed. What does Vizier do with observation noise? Is it simply ignored, or is this just a limitation of current implementation?

[saitcakmak]

We tend to use self.assertEqual(...), self.assertTrue/False(...) etc from TestCase rather than using bare assert in tests

[saitcakmak]

It'd be great to support Yvar as well. Could just set it to the mean of all Yvars? That would only minimally affect the model predictions and let the code run.

[saitcakmak]

BoTorch models do not support NaN inputs. If we didn't need to deal with NaN and somehow avoided unique, could we avoid the double for loop here?

[CompRhys]

I think this is the salient point, for me which indicates that I should probably have done this in Ax. Would Ax error out if a metric returned NaN before applying transforms or only after?

[saitcakmak]

I believe the error would come after, though I would have to double check. IIRC, the NaN will propagate down to BoTorch and raise an error there.

I am planning a major refactor of data handling within Ax ModelBridge layer & transforms. The goal is to reduce looping to make it more efficient & scalable. Though if you want to add these in Ax now, it doesn't have to be blocked on the refactor.

[saitcakmak]

If we didn't care about uniqueness, we could get the rank / order using torch.argsort. Imo, the duplicate observations should be pretty rare in practice and probably are safe to ignore. They're probably mostly introduced by the infeasibility warper.

[saitcakmak]

I guess we could also find the median using the middle index if we simply sorted Y from the start.

[CompRhys]

I think it's worth having this generality incase other packages like BoFire want to make use of these methods, I tried to stick as close to the reference implementation as possible but expanding to handle multi-dimension and batch_size

[saitcakmak]

This is quite a complicated process to untransform. I find it quite difficult to review since I don't know what the algorithm is supposed to do.

[CompRhys]

The HalfRankTransform splits the data at the median and then warps the bottom half onto a gaussian whose std matches the std of the top half of the distribution. On inverting there are 3 cases, a) we can just look it up as it's very close to a known value from the fit, b) we take a in support linear interpolation, c) outside of the support we take a linear extrapolation based on the entire data range.

[CompRhys]

Thanks for adding these. I've been curious to try out the Vizier transforms. These transforms try to ensure we model the good regions better, which seems better than trying to model everything equally well, when the goal is to optimize a function. Have you tried these on any benchmarks yet? It'd be good to see how they work e2e.

I implemented these here before looking into what I would need to do to connect them to my benchmark testcase that I have been prototyping in Ax. I assumed that the best case would be to implement it lower in the stack and then bubble up but seemingly I should have implemented at the level of Ax (especially as the harder part here was handling the batch_size) hence the comment about potentially misplaced effort. It may be that these shouldn't be merged here if the details can't be hashed out.

None of these transforms implement the untransform_posterior method, which is necessary to make sure the transforms work as an outcome_transform in BoTorch models. Without it, Model.posterior will error out. Would it be possible to implement untransform_posterior, likely returning a TransformedPosterior object.

This should be able to be done based on the logic in Standardize and just using the sample transform. In general the bit here that makes things annoying is the batch_size.

None of the transforms seem to support Yvar either. We use BoTorch through Ax and try to keep the same methods regardless of whether the noise is observed. What does Vizier do with observation noise? Is it simply ignored, or is this just a limitation of current implementation?

I think the vizier approach is similar to performing the transformations at the layer of Ax. Their equivalent to Metric only deals with a single float observation. As such and from the docs I believe it's ignored. I don't think that we can implement Yvar for either the LogWarping or the HalfRank transform

[saitcakmak]

Ax vs BoTorch

This is something I've also been thinking about quite a bit lately, as part of reworking what Ax uses by default (though this has been rather focused on input transforms so far). Outcome transforms may be slightly simpler in Ax. You technically only need to be able to transform, though we also want some rough untransform method for reporting & cross validation purposes. In BoTorch, you need the transform to be differentiable, since the untransformed posterior is used to compute the acquisition value which needs to be differentiated through during optimization. Ax transforms are applied once as a pre-processing step but BoTorch transforms are applied at each posterior call. There is also the aspect of learnable / tunable transforms in BoTorch but that doesn't apply to outcome transforms.

If I were to implement Vizier transforms, I'd probably try them in Ax, since there are fewer restrictions.

In general the bit here that makes things annoying is the batch_size.

Do you often use batched models? If not, you don't need to support batch_size arg and different treatment for each batch. You could use the same parameters for all batches.

I think the vizier approach is similar to performing the transformations at the layer of Ax.

Agreed. It is more of a pre-processing step.

Their equivalent to Metric only deals with a single float observation. As such and from the docs I believe it's ignored. I don't think that we can implement Yvar for either the LogWarping or the HalfRank transform

Cool, this is good to know.

[CompRhys]

Will move forward with a PR at Ax as that seems like the faster way to make progress on my usecase

[saitcakmak]

Sounds good. StandardizeY or PowerTransformY should be good examples to follow