From 0bb9521fbec836d6c5afe5f2ef37ff7c1e3b7759 Mon Sep 17 00:00:00 2001 From: Philipp Emanuel Weidmann Date: Thu, 5 Mar 2026 08:55:58 +0530 Subject: [PATCH] feat(ara): optimize all parameters --- src/heretic/main.py | 47 +++++++++++++++++++++++++++++++------------- src/heretic/model.py | 44 ++++++++++++++++++++--------------------- src/heretic/utils.py | 7 ++++--- 3 files changed, 59 insertions(+), 39 deletions(-) diff --git a/src/heretic/main.py b/src/heretic/main.py index f003679..bc7066f 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -38,7 +38,7 @@ from rich.traceback import install from .analyzer import Analyzer from .config import QuantizationMethod, Settings from .evaluator import Evaluator -from .model import AbliterationParameters, Model, get_model_class +from .model import AbliterationParameters, ARAParameters, Model, get_model_class from .utils import ( empty_cache, format_duration, @@ -474,11 +474,38 @@ def run(): len(model.get_layers()) // 2, len(model.get_layers()), ) - optimization_balance = trial.suggest_float( - "optimization_balance", - -1.0, + preserve_good_behavior_weight = trial.suggest_float( + "preserve_good_behavior_weight", + 0.0, 1.0, ) + steer_bad_behavior_weight = trial.suggest_float( + "steer_bad_behavior_weight", + 0.001, + 1.0, + log=True, + ) + overcorrect_relative_weight = trial.suggest_float( + "overcorrect_relative_weight", + 0.0, + 1.0, + ) + neighbor_count = trial.suggest_int( + "neighbor_count", + 1, + 10, + ) + + ara_parameters = ARAParameters( + start_layer_index=start_layer_index, + end_layer_index=end_layer_index, + preserve_good_behavior_weight=preserve_good_behavior_weight, + steer_bad_behavior_weight=steer_bad_behavior_weight, + overcorrect_relative_weight=overcorrect_relative_weight, + neighbor_count=neighbor_count, + ) + + trial.set_user_attr("ara_parameters", asdict(ara_parameters)) else: direction_scope = trial.suggest_categorical( "direction_scope", @@ -559,13 +586,7 @@ def run(): print("* Reloading model...") model.reset_model() print("* Abliterating (Arbitrary-Rank Ablation)...") - model.ara_abliterate( - good_module_io, - bad_module_io, - start_layer_index, - end_layer_index, - optimization_balance, - ) + model.ara_abliterate(good_module_io, bad_module_io, ara_parameters) else: print("* Resetting model...") model.reset_model() @@ -756,9 +777,7 @@ def run(): model.ara_abliterate( good_module_io, bad_module_io, - trial.params["start_layer_index"], - trial.params["end_layer_index"], - trial.params["optimization_balance"], + ARAParameters(**trial.user_attrs["ara_parameters"]), ) else: print("* Resetting model...") diff --git a/src/heretic/model.py b/src/heretic/model.py index 9feaab9..0a4b4fb 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -54,6 +54,16 @@ class AbliterationParameters: min_weight_distance: float +@dataclass +class ARAParameters: + start_layer_index: int + end_layer_index: int + preserve_good_behavior_weight: float + steer_bad_behavior_weight: float + overcorrect_relative_weight: float + neighbor_count: int + + # The list contains one element per layer. # Each element maps from the component name to a (possibly sparse) mapping # from the module index to an (input, output) tuple containing the I/O @@ -541,19 +551,12 @@ class Model: self, good_module_io: ModuleIO, bad_module_io: ModuleIO, - start_layer_index: int, - end_layer_index: int, - optimization_balance: float, + parameters: ARAParameters, ): - preserve_good_behavior_weight = 1.0 - steer_bad_behavior_weight = 1.0 - - if 0.0 < optimization_balance <= 1.0: - preserve_good_behavior_weight = 1.0 - optimization_balance - elif -1.0 <= optimization_balance < 0.0: - steer_bad_behavior_weight = 1.0 - abs(optimization_balance) - - for layer_index in range(start_layer_index, end_layer_index): + for layer_index in range( + parameters.start_layer_index, + parameters.end_layer_index, + ): for component, modules in self.get_layer_modules(layer_index).items(): for module_index, module in enumerate(modules): # See above for a (partial) justification of this cast. @@ -576,35 +579,32 @@ class Model: (new_good_output - good_output) ** 2 ).mean() - # TODO: Justify the magic weights here and make them configurable/optimizable. - # Experimentally, steer_bad_behavior needs to be about an order - # of magnitude larger than preserve_good_behavior for good results. steer_bad_behavior = ( - 1.0 # Pull the outputs for "bad" prompts towards # the original outputs for "good" prompts. - * mean_distances_to_knn( + mean_distances_to_knn( new_bad_output, good_output, - 10, + parameters.neighbor_count, ).mean() - + 0.5 # Push the outputs for "bad" prompts away from # the original outputs for "bad" prompts. # In combination with the above, this overcorrects # away from the original residuals, which results # in stronger steering that can overcome more complex # refusal mechanisms. + + parameters.overcorrect_relative_weight * -mean_distances_to_knn( new_bad_output, bad_output, - 10, + parameters.neighbor_count, ).mean() ) return ( - preserve_good_behavior_weight * preserve_good_behavior - + steer_bad_behavior_weight * steer_bad_behavior + parameters.preserve_good_behavior_weight + * preserve_good_behavior + + parameters.steer_bad_behavior_weight * steer_bad_behavior ) optimizer = LBFGS( diff --git a/src/heretic/utils.py b/src/heretic/utils.py index 8f0c94e..2cf1462 100644 --- a/src/heretic/utils.py +++ b/src/heretic/utils.py @@ -261,10 +261,11 @@ def empty_cache(): def get_trial_parameters(settings: Settings, trial: Trial) -> dict[str, str]: if settings.use_ara: + parameters = trial.user_attrs["ara_parameters"] + return { - "start_layer_index": f"{trial.params['start_layer_index']}", - "end_layer_index": f"{trial.params['end_layer_index']}", - "optimization_balance": f"{trial.params['optimization_balance']:.4f}", + name: (f"{value:.4f}" if isinstance(value, float) else f"{value}") + for name, value in parameters.items() } else: params = {}