Constrained multi-objective multi-fidelity BO

[Frxljord]

Hi everybody,

I'm trying to build a BO for a problem with 15 input parameters, 2 objectives, 2 constraints. I can evaluate a point in my design space at either a low (0) or high (1) fidelity (2 discrete levels, assume costs 1 and 100 for evaluation). The constrained nature of my problem prevents me from using Multi-fidelity Bayesian optimization with discrete fidelities using KG directly, as KG doesn't support constraints. I would like to optimize a per cost unit qNEHVI and compare with a single fidelity BO also using qNEHVI, but I'm having trouble understanding how to tell the program that the high fidelity should be the target. This is my code for now. I initialize it with only low fidelity data, and even if I set identical costs for both fidelities (or even a lower cost for high fidelity!) the BO only ever evaluates at low fidelity. This hints to the fact that the BO doesn't actually know what the target fidelity is.
Currently, my code is

    def generate_initial_data(self, n_init):
        X_params = draw_sobol_samples(
            bounds=self.problem.bounds, 
            n=n_init, 
            q=1
        ).squeeze(1)
        
        fidelities = torch.zeros(n_init, 1, **self.tkwargs)
        X = torch.cat([X_params, fidelities], dim=1)
        
        print(f"Generating {n_init} LF initial points:")
        
        Y = self.problem(X)
        return X, Y

    def create_multi_fidelity_models(self, X, Y):
        X_norm = normalize(X, self.extended_bounds)
        models = []
        Y_var = torch.full_like(Y, 1e-6)
        
        for i in range(Y.shape[-1]):
            model = SingleTaskMultiFidelityGP(
                train_X=X_norm,
                train_Y=Y[..., i:i+1],
                train_Yvar=Y_var[..., i : i + 1],
                outcome_transform=Standardize(m=1),
                input_transform=Normalize(d=self.total_dim),
                data_fidelities=[-1]
            )
            models.append(model)
        
        model_list = ModelListGP(*models)
        mll = SumMarginalLogLikelihood(model_list.likelihood, model_list)
        
        return mll, model_list

    def get_cost_model(self):
        def cost_func(X):
            fidelity = X[..., -1]
            discrete_fidelity = fidelity > 0.5
            
            costs = torch.zeros_like(fidelity)
            for fid_level, cost in self.fidelity_costs.items():
                mask = discrete_fidelity == fid_level
                costs[mask] = cost
            return costs
        return cost_func

    def generate_candidates(self, model, X_train, batch_size=1):
        X_norm = normalize(X_train, self.extended_bounds)
        
        base_acq = qLogNoisyExpectedHypervolumeImprovement(
            model=model,
            ref_point=self.problem.ref_point,
            X_baseline=X_norm,
            sampler=SobolQMCNormalSampler(sample_shape=torch.Size([128])),
            prune_baseline=True,
            objective=IdentityMCMultiOutputObjective(outcomes=[0, 1]),
            constraints=[
                lambda Y: Y[..., 2],
                lambda Y: Y[..., 3],
            ],
            eta=1e-6
        )
        
        cost_model = self.get_cost_model()
        acq_func = CostNormalizedAcquisition(
            base_acq,
            cost_model=cost_model,
            alpha=1.0,   # qNEHVI per unit cost
            eps=1e-8,
        )
        
        candidates, _ = optimize_acqf_mixed(
            acq_function=acq_func,
            bounds=torch.tensor([[0.0] * self.total_dim, [1.0] * self.total_dim], **self.tkwargs),
            fixed_features_list=[
                {self.param_dim: 0.0},
                {self.param_dim: 1.0},
            ],
            q=batch_size,
            num_restarts=10,
            raw_samples=1024,
            options={"batch_limit": 5, "maxiter": 200},
        )
        
        params_norm = candidates[..., :-1]
        fid_column = candidates[..., -1]

        params_unnorm = unnormalize(params_norm, self.problem.bounds)
        candidates_unnorm = torch.cat([params_unnorm, fid_column.unsqueeze(-1)], dim=-1)

        return candidates_unnorm

and my per-unit acquisition function is

class CostNormalizedAcquisition(MCAcquisitionFunction):
    def __init__(self, base_acq, cost_model, alpha = 1.0, eps = 1e-8):
        super().__init__(model=base_acq.model,
                         objective=base_acq.objective
                         )
        self.base_acq = base_acq
        self.cost_model = cost_model
        self.alpha = float(alpha)
        self.eps = float(eps)

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        base_val = self.base_acq(X)
        cost = self.cost_model(X).to(base_val.dtype)

        if cost.ndim > base_val.ndim:
            cost = cost.squeeze(-1)

        denom = torch.pow(cost + self.eps, self.alpha)
        return base_val / denom

Should I have some cost-based utility to tell the optimizer to balance cost and information gain, and/or use project_to_target_fidelity like in qMultiFidelityKnowledgeGradient?

Any help is appreciated, many thanks :)

[esantorella]

Yeah, you would need to somehow tell the optimizer that the higher fidelity is more reliable. Would MOMF suit your use case?

[Frxljord]

Yes, MOMF suits my use case (as a fallback), although my concern was that I would be unfairly comparing single-fidelity vs. multi-fidelity by using two (sort of?) different acquisition functions for the two methods. Thanks!