Skip to content
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

Jump to bottom

Refactor reward-modeling likelihood #200

Merged
merged 4 commits into from
Jun 23, 2024
Merged

Refactor reward-modeling likelihood #200

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ecd5220
Refactor reward-modeling likelihood
runame Jun 15, 2024
15f7c79
Merge branch 'grid-search' into reward-modeling-refactor
runame Jun 23, 2024
48f4f8e
Merge branch 'grid-search' into reward-modeling-refactor
wiseodd Jun 23, 2024
d6ff3b2
Merge branch 'reward-modeling-refactor' of github.com:aleximmer/Lapla…
wiseodd Jun 23, 2024
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
57 changes: 19 additions & 38 deletions laplace/baselaplace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ def __init__(
):
if likelihood not in ["classification", "regression", "reward_modeling"]:
raise ValueError(f"Invalid likelihood type {likelihood}")

self.likelihood = likelihood
self.model = model

# Only do Laplace on params that require grad
Expand All @@ -109,13 +109,6 @@ def __init__(
if sigma_noise != 1 and likelihood != "regression":
raise ValueError("Sigma noise != 1 only available for regression.")

self.reward_modeling = likelihood == "reward_modeling"
if self.reward_modeling:
# For fitting only. After it's done, self.likelihood = 'regression', see self.fit()
self.likelihood = "classification"
else:
self.likelihood = likelihood

self.sigma_noise = sigma_noise
self.temperature = temperature
self.enable_backprop = enable_backprop
Expand Down Expand Up @@ -148,21 +141,21 @@ def __init__(
self.n_outputs = None
self.n_data = 0

# Useful for reward modeling where it behaves like classification during Hessian
# computation and prior-prec optimization, and behaves like regression during
# prediction
self._fitting: bool = True

@property
def _device(self):
return next(self.model.parameters()).device

@property
def backend(self):
if self._backend is None:
likelihood = (
"classification"
if self.likelihood == "reward_modeling"
else self.likelihood
)
self._backend = self._backend_cls(
self.model,
self.likelihood,
likelihood,
dict_key_x=self.dict_key_x,
dict_key_y=self.dict_key_y,
**self._backend_kwargs,
Expand Down Expand Up @@ -310,7 +303,6 @@ def optimize_prior_precision_base(
link_approx="probit",
n_samples=100,
verbose=False,
cv_loss_with_var=False,
progress_bar=False,
):
"""Optimize the prior precision post-hoc using the `method`
Expand Down Expand Up @@ -340,9 +332,6 @@ def optimize_prior_precision_base(
If torchmetrics.Metric, running loss is computed (efficient). The default
depends on the likelihood: `RunningNLLMetric()` for classification and
reward modeling, running `MeanSquaredError()` for regression.
cv_loss_with_var: bool, default=False
if true, `loss` takes three arguments `loss(output_mean, output_var, target)`,
otherwise, `loss` takes two arguments `loss(output_mean, target)`
log_prior_prec_min : float, default=-4
lower bound of gridsearch interval.
log_prior_prec_max : float, default=4
Expand All @@ -361,8 +350,11 @@ def optimize_prior_precision_base(
whether to show a progress bar; updated at every batch-Hessian computation.
Useful for very large model and large amount of data, esp. when `subset_of_weights='all'`.
"""
if self.reward_modeling:
self.likelihood = "classification"
likelihood = (
"classification"
if self.likelihood == "reward_modeling"
else self.likelihood
)

if method == "marglik":
self.prior_precision = init_prior_prec
Expand Down Expand Up @@ -402,7 +394,7 @@ def optimize_prior_precision_base(
if loss is None:
loss = (
tm.MeanSquaredError(num_outputs=self.n_outputs).to(self._device)
if self.likelihood == "regression"
if likelihood == "regression"
else RunningNLLMetric().to(self._device)
)

Expand All @@ -413,7 +405,6 @@ def optimize_prior_precision_base(
pred_type=pred_type,
link_approx=link_approx,
n_samples=n_samples,
loss_with_var=cv_loss_with_var,
progress_bar=progress_bar,
)
else:
Expand All @@ -430,7 +421,6 @@ def _gridsearch(
pred_type,
link_approx="probit",
n_samples=100,
loss_with_var=False,
progress_bar=False,
):
assert callable(loss) or isinstance(loss, tm.Metric)
Expand All @@ -449,7 +439,6 @@ def _gridsearch(
pred_type=pred_type,
link_approx=link_approx,
n_samples=n_samples,
loss_with_var=loss_with_var,
dict_key_y=self.dict_key_y,
)

Expand Down Expand Up @@ -576,9 +565,6 @@ def fit(self, train_loader, override=True, progress_bar=False):
self.loss = 0
self.n_data = 0

if self.reward_modeling:
self.likelihood = "classification"

self.model.eval()

self.mean = parameters_to_vector(self.params)
Expand Down Expand Up @@ -628,7 +614,6 @@ def fit(self, train_loader, override=True, progress_bar=False):
self.H += H_batch

self.n_data += N
self._fitting = False

@property
def scatter(self):
Expand Down Expand Up @@ -825,15 +810,16 @@ def __call__(
):
raise ValueError("Invalid random generator (check type and device).")

if self.reward_modeling:
self.likelihood = "classification" if fitting else "regression"
likelihood = self.likelihood
if likelihood == "reward_modeling":
likelihood = "classification" if fitting else "regression"

if pred_type == "glm":
f_mu, f_var = self._glm_predictive_distribution(
x, joint=joint and self.likelihood == "regression"
x, joint=joint and likelihood == "regression"
)
# regression
if self.likelihood == "regression":
if likelihood == "regression":
return f_mu, f_var
# classification
if link_approx == "mc":
Expand Down Expand Up @@ -871,7 +857,7 @@ def __call__(
alpha = (1 - 2 / K + f_mu.exp() / K**2 * sum_exp) / f_var_diag
return torch.nan_to_num(alpha / alpha.sum(dim=1).unsqueeze(-1), nan=1.0)
else:
if self.likelihood == "regression":
if likelihood == "regression":
samples = self._nn_predictive_samples(x, n_samples, **model_kwargs)
return samples.mean(dim=0), samples.var(dim=0)
else: # classification; the average is computed online
Expand Down Expand Up @@ -1050,7 +1036,6 @@ def optimize_prior_precision(
link_approx="probit",
n_samples=100,
verbose=False,
cv_loss_with_var=False,
progress_bar=False,
):
assert pred_type in ["glm", "nn"]
Expand All @@ -1069,7 +1054,6 @@ def optimize_prior_precision(
link_approx,
n_samples,
verbose,
cv_loss_with_var,
progress_bar,
)

Expand Down Expand Up @@ -1480,9 +1464,6 @@ def fit(self, train_loader, override=True):
# LowRankLA cannot be updated since eigenvalue representation not additive
raise ValueError("LowRank LA does not support updating.")

if self.reward_modeling:
self.likelihood = "classification"

self.model.eval()
self.mean = parameters_to_vector(self.model.parameters())

Expand Down
1 change: 0 additions & 1 deletion laplace/utils/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ def validate(
pred_type="glm",
link_approx="probit",
n_samples=100,
loss_with_var=False,
dict_key_y="labels",
) -> float:
laplace.model.eval()
Expand Down
Loading