OptimizationGradientError: x elements of the 80 element gradient array gradf are NaN.

[Friedemannn]

Hey there,

I'm working on a BO Loop optimizing 2 physical parameters of a PIC simulation.

Up until now it went well. As input, i use a relatively well manually optimized PIC parameter file, of which two parameters then are varied (in bounds) by the BO Loop. They are also 10 initial simulations where the two parameters are varied randomly before the 40 iterations BO loop with q=4 starts (so 11 initial and 160 BO loop simulations).

The problems arise when the initial parameter file is chosen not so carefully. Then i more often then not run into the following error (sometimes after a few, sometimes after a lot of iterations):

in bo_loop candidates, _ = optimize_acqf(
...
botorch.exceptions.errors.OptimizationGradientError: 8 elements of the 80 element gradient array `gradf` are NaN. This often indicates numerical issues.

My code looks like this:

torch.set_default_dtype(torch.double)
...
    for i in range(num_loops):
        t_model1 = time.time()
        model = SingleTaskGP(x_trainT, y_scaled, outcome_transform=Standardize(1))
        mll = ExactMarginalLogLikelihood(model.likelihood, model)
        fit_gpytorch_mll(mll)

        # Optimize acquisition
        acq_func = qLogExpectedImprovement(model=model, best_f = torch.max(y_scaled),sampler=SobolQMCNormalSampler(sample_shape=torch.Size([1024]))) 
        candidates, _ = optimize_acqf(
            acq_function=acq_func,
            bounds=torch.tensor([[0.0] * num_leaf_keys(bounds_dict), [1.0] * num_leaf_keys(bounds_dict)]),
            q=batch_size,
            num_restarts=num_restarts,
            raw_samples=raw_samples,
        )

After that PIC simulations are run based on the candidates, with the results the objective y_scaled is calculated.
If no electron beam was created in the simulation (or smth else goes wrong) the y_scaled value ist set to 0, so there can be a lot of 0s. But this error also seems to happen if there arent even that many 0s.

What i tried to do:

• using qExpectedImprovement instead of qLogExpectedImprovement bc i thought maybe the log part cant handle the zeros.
• Setting y_scaled to 1E-8 instead of 0

Is there something I can do botorch-wise to make the optimize_acqf() more robust against this?

Thank you!

[saitcakmak]

Hi @Friedemannn. Wrapping the loop in with torch.autograd.set_detect_anomaly(True): could help pinpoint where the gradient instabilities come from. Another idea that comes to mind is to try qLogNoisyExpectedImprovement rather than qLogExpectedImprovemment, which would handle noisy observations better -- though I don't know if that would have any effect on gradient issues necessarily.

[Friedemannn]

Hi Sait,
thanks for the input!
When running with torch.autograd.set_detect_anomaly(True): the loop stops earlier and I get:

  File "site-packages/torch/autograd/graph.py", line 824, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Function 'LinalgEighBackward0' returned nan values in its 0th output.

This doesn't really mean much more to me than the first error message. Do you have a guess or explanation what's going wrong?
Are there too many similar objective values for the optimizer? Or are there too many similar x values or a combination of both?

The weird thing is also there are PIC parameter files, based on which the full Loop will run through with a lot of 0 objective values and then there are some for which I get the errors above directly after the initial simulations or after one or two loop iterations with not even that many 0 values. With the same two parameters/variables being optimized.

Maybe some more info can help:

For Botorch the objective values are scaled with y_scale = 1e13 (The order of the objective is expected to be around 1e-8 and 1e-4)

"NUM_RESTARTS": 10,
"RAW_SAMPLES": 512,
"MC_SAMPLES": 128

The results/logs for the parameter file with which I get the errors above look like this:
with detect_anomaly:
paramterFile1_detect_anomaly_results.txt

without detect_anomaly:
parameterFile1_results.txt

Here are also the results of a run with a parameteter file, where basically all objective values were 0 but the loop nevertheless ran:
parameterFile2_results.txt

I can also upload the model.pt file for that run if thats helpful.

[hvarfner]

Hi @Friedemannn ,

I took a quick look at the parameter file and saw that all the values for recommended points were tiny, in the 0.01 range. I see that the bounds of your search space are specified as [0, 1]^D, but how do the X's look going into your GP? Is it intentional to have the initial samples be in this small range, but still have the bounds be [0, 1]?

Since the values cover such a small range of the entire search space, the GP will not be able to fit the rest of the search space at all. You should do some space-filling with sobol for the second batch too. Then, you should hopefully see the BO working.

[hvarfner]

I looked at the second parameters file, too. I see that

 bounds=torch.tensor([[0.0] * num_leaf_keys(bounds_dict), [1.0] * num_leaf_keys(bounds_dict)]),
            q=batch_size,

but most of the data in parameterFile2_results.txt are exactly

[0.0002 , 0.0028 ],
which appears to be an artificially imposed upper bound. How come the evaluations max out at this value? Do you round them down?

If [0.0002 , 0.0028 ], is indeed your upper bound, you need normalize the input accordingly. The best way to do this is with:

import torch
from botorch.sampling import draw_sobol_samples
from botorch.models import SingleTaskGP
from botorch.models.transforms import Normalize
from botorch.acquisition import qLogExpectedImprovement
from botorch.optim import optimize_acqf

# --- User-defined parameters and objective ---

NUM_INIT = 11
batch_size = 10
num_restarts = 10
raw_samples = 1024


# Define bounds here 
# TODO make sure these are what you want to optimize over
bounds = torch.tensor([[0.0, 0.0], [0.002, 0.0028]])
# Dummy objective function for illustration
def your_objective(X):
    # Example: sum of squares (replace with your actual objective)
    return (X ** 2).sum(dim=-1, keepdim=True)

# Generate initial data - TODO make sure your initial data is within these bounds, or at least close
train_X = draw_sobol_samples(n=NUM_INIT, bounds=bounds, q=1).squeeze(-2)


train_Y = your_objective(train_X)
# Normalize inputs and fit GP model
input_transform = Normalize(d=2, bounds=bounds) # BOUNDS GO HERE

# The outcome transform is added by default, no need to add it!
model = SingleTaskGP(train_X=train_X, train_Y=train_Y, input_transform=input_transform)

# The SobolMCSampler is added by default
acq_func = qLogExpectedImprovement(model=model, best_f=torch.max(train_Y))
# Optimize acquisition function to get new candidates
candidates, _ = optimize_acqf(
    acq_function=acq_func,
    bounds=bounds, # BOUNDS GO HERE
    q=batch_size,
    num_restarts=num_restarts,
    raw_samples=raw_samples,
)
# Evaluate objective at new candidates
new_Y = your_objective(candidates)

Right now, the GP actually wants to go to higher values since you've specified the bounds as [0, 1], and so you end up with a lot of repeat evaluations at the exact same point. After a while, you run into numerical issues.

[hvarfner]

So @Friedemannn ,

Check out the TODOs regarding the bounds and I think you should be fine! If I'm wrong about something here, let me know!

[Friedemannn]

Hi @hvarfner thank you for the detailed answer!

I think I'm already doing most/all of the things in your example. (My raw_samples=512, don't know why I set it to that though.)

Maybe the original code snippet was a bit to short:

def bo_loop(...):
    ...
    # Initialize tensors for BO
    x_trainT = torch.from_numpy(x_train).reshape(-1, num_leaf_keys(bounds_dict))
    objectiveT = torch.from_numpy(np.abs(y_train)).reshape(-1, 1)
    bounds3DT = torch.from_numpy(
        np.array([get_value_by_path(bounds_dict, path) for path in get_leaf_paths(bounds_dict)])
    ).T #these are the [0.0002    , 0.0028] bounds you saw

    x_trainT = normalize(x_trainT, bounds3DT) #normalizing training data using [0.0002    , 0.0028]
    
    # This is a version with two parameters that are optimized and both conicidentally have the bounds [0.0002   , 0.0028] (one is
    # the laser focus and the other one is the start point of a gas)
   

    # Scale objective so that standardization is numerically stable
    y_scaled = objectiveT * y_scale
    objective_list = y_train

    #begin BO loop
    for i in range(num_loops):
        model = SingleTaskGP(x_trainT, y_scaled, outcome_transform=Standardize(1))
        mll = ExactMarginalLogLikelihood(model.likelihood, model)
        fit_gpytorch_mll(mll)

        # Optimize acquisition
        acq_func = qLogExpectedImprovement(model=model, best_f = torch.max(y_scaled),sampler=SobolQMCNormalSampler(sample_shape=torch.Size([1024])))
        candidates, _ = optimize_acqf(
            acq_function=acq_func,
            bounds=torch.tensor([[0.0] * num_leaf_keys(bounds_dict), [1.0] * num_leaf_keys(bounds_dict)]),
            q=batch_size,
            num_restarts=num_restarts,
            raw_samples=raw_samples,  
        )


        # Append new candidate(s)
        x_trainT = torch.cat([x_trainT, candidates])
        new_x_train = unnormalize(candidates, bounds3DT).detach()

        # Write new x values into parameter.json to later call with fbpic (PIC Simulation library i use)
        ...

        # Run simulations with new x value(s)
        ...

        # Wait for all simulations to complete
        ...

        # Load sim results and calculate objective for new x value(s)
        ...
            objective = np.sqrt(np.abs(charge))/(mad/median_energy)
            objective_list = np.append(objective_list, objective)
            objectiveT = torch.cat([objectiveT, torch.tensor([[objective]])])
            y_scaled = torch.cat([y_scaled, torch.tensor([[objective * y_scale]])])

    # Save BoTorch model
    torch.save(model.state_dict(), os.path.join(run_dir, "model.pt"))

    return run_dir

For the initial data, I basically do the same, but just randomly generate values in the bounds of the two parameters that are optimized and then run the simulations to calculate the objective with the sim results.

With this I should end up with the same x-values as you, normalized between 0 and 1 using the bounds [0.0002 , 0.0028], right?

If needed, I can also upload the whole file.

[Friedemannn]

So i think I solved the OptimizationGradientError .... I'm not a 100% sure because I wasnt able to rerun the parameter files yet that this issue was originally about.

But after trying a bunch of random stuff chatgpt suggested, I had success by using a minimum noise contraint for the likelihood and running the fit with cholesky jitter.

def bo_loop:

    min_noise = float(loop_dict.get("min_noise", 1e-6))
    cholesky_jitter = float(loop_dict.get("cholesky_jitter", 1e-6))

...

        likelihood = gpytorch.likelihoods.GaussianLikelihood(
            noise_constraint=gpytorch.constraints.GreaterThan(min_noise)
        )
        model = SingleTaskGP(
            x_trainT,
            y_scaled,
            likelihood=likelihood,
            outcome_transform=Standardize(1),
        )
        mll = ExactMarginalLogLikelihood(model.likelihood, model)
        with gpytorch.settings.cholesky_jitter(cholesky_jitter):
            fit_gpytorch_mll(mll)

It's nice that it seems to work now, but tbh I have no clue what's happening under the hood. Can anyone give me some recommendations for reading material on this?

[Friedemannn]

Here are some before and after results:

This is a run from before the noise floor and cholesky jitter:

While this surprisingly ran without the OptimizationGradientError it shows similar behavior of having a bunch of candidates at the min and max of y and x, which you pointed out before @hvarfner

This is a run after the addidtion of the noise:

This looks much better, but I'm confused of how the bo loop manages to avoid the 0 objective values close to the maximum that were present before. I think this is weird, because the values of the dots are just the objective values calculated from the sim results, without botorch interfering in any way. So one would expect atleast a few of the same zero dots as well.
All in all this kind of linear behavior is also expected from a physics standpoint, i think.

This is a rerun after the addidtion of the noise, nothing has changed to the one from above. I just wanted to rerun to see if its consistent:

While we arrive at similar results, the bo_loop behavior is different, somehow its at the same time more explorative and not. There are a bunch of candidates close to the maximum, but also more spread out ones than before.

Can anyone explain this behavior?