8000 Normalize loss value by the number of observations only when using SGD-based GP hparam optimizers (so a consistent learning rate can be used regardless of the number of observations). Previously, this was baked into the loss function definition; now that we're using L-BFGS-B by default, this is unnecessary and possibly confusing. by copybara-service[bot] · Pull Request #602 · google/vizier · GitHub
[go: up one dir, main page]
Skip to content

Normalize loss value by the number of observations only when using SGD-based GP hparam optimizers (so a consistent learning rate can be used regardless of the number of observations). Previously, this was baked into the loss function definition; now that we're using L-BFGS-B by default, this is unnecessary and possibly confusing. #602

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Mar 31, 2023
Merged

Normalize loss value by the number of observations only when using SGD-based GP hparam optimizers (so a consistent learning rate can be used regardless of the number of observations). Previously, this was baked into the loss function definition; now that we're using L-BFGS-B by default, this is unnecessary and possibly confusing. #602

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 20 additions & 4 deletions vizier/_src/algorithms/designers/gp_bandit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,14 +216,29 @@ def _find_best_model_params(self) -> chex.ArrayTree:
# Run ARD.
setup = lambda rng: self._model.init(rng, self._features)['params']
constraints = sp.get_constraints(self._model.coroutine)
loss_fn = jax.jit(loss_fn)
ard_loss_fn = self._get_loss_fn(loss_fn)

logging.info('Optimizing the loss function...')
self._rng, ard_rng = jax.random.split(self._rng, 2)
best_model_params, _ = self._ard_optimizer(
setup, loss_fn, ard_rng, constraints=constraints
setup, ard_loss_fn, ard_rng, constraints=constraints
)
return best_model_params

def _get_loss_fn(self, loss_fn):
if isinstance(self._ard_optimizer, optimizers.OptaxTrainWithRandomRestarts):

def ard_loss_fn(x):
loss_val, metrics = loss_fn(x)
# For SGD, normalize the loss so we can use the same learning rate
# regardless of the number of examples (see
# `OptaxTrainWithRandomRestarts` docstring).
return loss_val / self._features.shape[0], metrics

else:
ard_loss_fn = loss_fn
return jax.jit(ard_loss_fn)

def _compute_state(self):
"""Compute the designer's state.

Expand Down Expand Up @@ -251,9 +266,10 @@ def _compute_state(self):
# Logging for debugging purposes.
logging.info('ARD duration: %s', (datetime.datetime.now() - ard_start_time))
logging.info('Best model parameters: %s', best_model_params)
ard_loss_fn = self._get_loss_fn(loss_fn)
if self._use_vmap:
loss_fn = jax.vmap(loss_fn)
ard_all_losses = loss_fn(best_model_params)[0]
ard_loss_fn = jax.vmap(ard_loss_fn)
ard_all_losses = ard_loss_fn(best_model_params)[0]
logging.info('All losses: %s', ard_all_losses)
ard_best_loss = ard_all_losses.flatten()[0].item()
logging.info('ARD best loss: %s', ard_best_loss)
Expand Down
9 changes: 3 additions & 6 deletions vizier/_src/jax/models/hebo_gp_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,12 +59,9 @@ def loss_fn(params):
gp, mutables = model.apply({'params': params},
features,
mutable=['losses', 'predictive'])
# Normalize so we can use the same learning rate regardless of
# how many examples we have.
loss = (
-gp.log_prob(labels)
+ jax.tree_util.tree_reduce(jnp.add, mutables['losses'])
) / features.shape[0]
loss = -gp.log_prob(labels) + jax.tree_util.tree_reduce(
jnp.add, mutables['losses']
)
return loss, dict()

return model, loss_fn
Expand Down
7 changes: 3 additions & 4 deletions vizier/_src/jax/models/tuned_gp_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,10 +77,9 @@ def loss_fn(params):
gp, mutables = model.apply({'params': params},
features,
mutable=['losses', 'predictive'])
# Normalize so we can use the same learning rate regardless of
# how many examples we have.
loss = (-gp.log_prob(labels) + jax.tree_util.tree_reduce(
jax.numpy.add, mutables['losses'])) / features.shape[0]
loss = -gp.log_prob(labels) + jax.tree_util.tree_reduce(
jax.numpy.add, mutables['losses']
)
return loss, dict()

return model, loss_fn
Expand Down
15 changes: 15 additions & 0 deletions vizier/_src/jax/optimizers/optimizers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,21 @@ def __call__(
class OptaxTrainWithRandomRestarts(Optimizer):
"""Wraps an Optax optimizer.

It's recommended to use this optimizer with a loss function that normalizes by
the number of observations. The unnormalized loss function for parameters `p`
is typically of the form

```None
-gp_likelihood(observations | p) + regularization(p)
```

where regularization may be a negative prior log probability. The likelihood
term is approximately proportional to the number of observations. As the
number of observations changes over the course of a study, dividing the loss
by this number helps ensure that loss values are roughly of the same order of
magnitude, such that a constant learning rate may be used for gradient-based
optimizers. Vizier library models make this adjustment automatically.

Attributes:
optimizer: Optax optimizer such as `optax.adam(1e-2)`.
epochs: Number of train epochs.
Expand Down
0