diff --git a/src/heretic/main.py b/src/heretic/main.py index 36c1e49..fbc18c3 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -578,10 +578,22 @@ def run(): # The parameter ranges are based on experiments with various models # and much wider ranges. They are not set in stone and might have to be # adjusted for future models. - max_weight = trial.suggest_float( - f"{component}.max_weight", - 0.8, - 1.5, + # + # The MLP gets a negative lower bound that is then clamped to 0, so the + # optimizer can fully disable its ablation. The clamp puts a positive + # probability mass on exactly 0 (the continuous sampler would otherwise + # reach 0 with probability zero). Ablating the MLP is often unnecessary for + # removing refusals and tends to damage model intelligence more than + # ablating the attention output, so on many models the optimum is to leave + # it (mostly) untouched. See issue #202. + max_weight_lower_bound = -0.25 if component == "mlp.down_proj" else 0.8 + max_weight = max( + 0.0, + trial.suggest_float( + f"{component}.max_weight", + max_weight_lower_bound, + 1.5, + ), ) max_weight_position = trial.suggest_float( f"{component}.max_weight_position", diff --git a/src/heretic/model.py b/src/heretic/model.py index 3ea72fc..f10ce9b 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -499,6 +499,12 @@ class Model: params.min_weight - params.max_weight ) + # A weight of 0 disables this component's ablation. reset_model() has + # already left the adapter at identity, so abort before the otherwise + # wasteful decomposition (which would also be operating on a zero matrix). + if weight == 0: + continue + if refusal_direction is None: # The index must be shifted by 1 because the first element # of refusal_directions is the direction for the embeddings.