vh/heretic
1
0
Fork 0
mirror of https://github.com/p-e-w/heretic.git synced 2026-09-19 02:26:17 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat(ara): optimize all parameters

Browse Source
This commit is contained in:
Philipp Emanuel Weidmann
2026-03-05 08:55:58 +05:30
parent 992fb3a4b3
commit 0bb9521fbe
3 changed files with 59 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/heretic/main.py
+33 -14
View File
@@ -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...")
src/heretic/model.py
+22 -22
View File
@@ -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(
src/heretic/utils.py
+4 -3
View File
@@ -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 = {}