feat: Constrain-able Abliteration

config.default.toml

@@ -91,6 +91,24 @@ n_trials = 200
 # Number of trials that use random sampling for the purpose of exploration.
 n_startup_trials = 60
 
+# Constraints for the optimization search space (layers and weights).
+[constraints]
+# Fraction of layers (from 0.0 to 1.0) where the direction search starts.
+layer_start_fraction = 0.4
+# Fraction of layers (from 0.0 to 1.0) where the direction search ends.
+layer_end_fraction = 0.9
+
+# Search constraints for Attention components (e.g., o_proj).
+[constraints.attention]
+max_weight_min = 0.8
+max_weight_max = 1.5
+
+# Search constraints for MLP components (e.g., down_proj).
+[constraints.mlp]
+max_weight_min = 0.8
+max_weight_max = 1.5
+
+
 # Directory to save and load study progress to/from.
 study_checkpoint_dir = "checkpoints"
 

src/heretic/config.py

@@ -61,6 +61,36 @@ class DatasetSpecification(BaseModel):
     )
 
 
+class ComponentConstraints(BaseModel):
+    max_weight_min: float = Field(
+        default=0.8,
+        description="Minimum value for the max_weight parameter search range.",
+    )
+    max_weight_max: float = Field(
+        default=1.5,
+        description="Maximum value for the max_weight parameter search range.",
+    )
+
+
+class OptimizationConstraints(BaseModel):
+    layer_start_fraction: float = Field(
+        default=0.4,
+        description="Fraction of layers (from 0.0 to 1.0) where the direction search starts.",
+    )
+    layer_end_fraction: float = Field(
+        default=0.9,
+        description="Fraction of layers (from 0.0 to 1.0) where the direction search ends.",
+    )
+    attention: ComponentConstraints = Field(
+        default_factory=ComponentConstraints,
+        description="Search constraints for Attention components (e.g., o_proj).",
+    )
+    mlp: ComponentConstraints = Field(
+        default_factory=ComponentConstraints,
+        description="Search constraints for MLP components (e.g., down_proj).",
+    )
+
+
 class Settings(BaseSettings):
     model: str = Field(description="Hugging Face model ID, or path to model on disk.")
 
@@ -225,6 +255,11 @@ class Settings(BaseSettings):
         description="Number of trials that use random sampling for the purpose of exploration.",
     )
 
+    constraints: OptimizationConstraints = Field(
+        default_factory=OptimizationConstraints,
+        description="Constraints for the optimization search space (layers and weights).",
+    )
+
     study_checkpoint_dir: str = Field(
         default="checkpoints",
         description="Directory to save and load study progress to/from.",

src/heretic/main.py

@@ -474,8 +474,8 @@ def objective(trial: Trial) -> tuple[float, float]:
         # work with conditional or variable-range parameters.
         direction_index = trial.suggest_float(
             "direction_index",
-            0.4 * last_layer_index,
-            0.9 * last_layer_index,
+            settings.constraints.layer_start_fraction * last_layer_index,
+            settings.constraints.layer_end_fraction * last_layer_index,
         )
 
         if direction_scope == "per layer":
@@ -487,10 +487,15 @@ def objective(trial: Trial) -> tuple[float, float]:
             # 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.
+
+            if "down_proj" in component:
+                constraints = settings.constraints.mlp
+            else:
+                constraints = settings.constraints.attention
             max_weight = trial.suggest_float(
                 f"{component}.max_weight",
-                0.8,
-                1.5,
+                constraints.max_weight_min,
+                constraints.max_weight_max,
             )
             max_weight_position = trial.suggest_float(
                 f"{component}.max_weight_position",