feat(ara): optimize all parameters

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(
+            preserve_good_behavior_weight = trial.suggest_float(
-                "optimization_balance",
+                "preserve_good_behavior_weight",
-                -1.0,
+                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(
+            model.ara_abliterate(good_module_io, bad_module_io, ara_parameters)
-                good_module_io,
-                bad_module_io,
-                start_layer_index,
-                end_layer_index,
-                optimization_balance,
-            )
         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"],
+                    ARAParameters(**trial.user_attrs["ara_parameters"]),
-                    trial.params["end_layer_index"],
-                    trial.params["optimization_balance"],
                 )
             else:
                 print("* Resetting model...")

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,
+        parameters: ARAParameters,
-        end_layer_index: int,
-        optimization_balance: float,
     ):
-        preserve_good_behavior_weight = 1.0
+        for layer_index in range(
-        steer_bad_behavior_weight = 1.0
+            parameters.start_layer_index,
-
+            parameters.end_layer_index,
-        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 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
+                            parameters.preserve_good_behavior_weight
-                            + steer_bad_behavior_weight * steer_bad_behavior
+                            * preserve_good_behavior
+                            + parameters.steer_bad_behavior_weight * steer_bad_behavior
                         )
 
                     optimizer = LBFGS(

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']}",
+            name: (f"{value:.4f}" if isinstance(value, float) else f"{value}")
-            "end_layer_index": f"{trial.params['end_layer_index']}",
+            for name, value in parameters.items()
-            "optimization_balance": f"{trial.params['optimization_balance']:.4f}",
         }
     else:
         params = {}