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

feat(ara): implement optimization for ARA parameters

Browse Source
This commit is contained in:
Philipp Emanuel Weidmann
2026-03-02 14:36:57 +05:30
parent 154241f8a2
commit b8f4a9c985
4 changed files with 237 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/heretic/config.py
+16
View File
@@ -176,6 +176,22 @@ class Settings(BaseSettings):
),
)
target_components: list[str] = Field(
default=["attn.o_proj"],
description=(
"List of component names to target for abliteration. "
'Currently supported values are "attn.o_proj" and "mlp.down_proj".'
),
)
use_ara: bool = Field(
default=True,
description=(
"Whether to use Arbitrary-Rank Ablation (ARA), an abliteration method based on matrix optimization, "
"instead of traditional directional ablation."
),
)
orthogonalize_direction: bool = Field(
default=False,
description=(
src/heretic/main.py
+178 -136
View File
@@ -206,8 +206,9 @@ def run():
"[bold yellow]No GPU or other accelerator detected. Operations will be slow.[/]"
)
# We don't need gradients as we only do inference.
# torch.set_grad_enabled(False)
if not settings.use_ara:
# We don't need gradients as we only do inference.
torch.set_grad_enabled(False)
# While determining the optimal batch size, we will try many different batch sizes,
# resulting in many computation graphs being compiled. Raising the limit (default = 8)
@@ -411,71 +412,47 @@ def run():
evaluator.get_score()
return
def tensor_shape_repr(self: torch.Tensor):
return f"tensor(shape={tuple(self.shape)}, dtype={self.dtype}, device={self.device})"
if settings.use_ara:
print()
print("Obtaining module I/O for good prompts...")
good_module_io = model.get_module_io_batched(good_prompts)
print("Obtaining module I/O for bad prompts...")
bad_module_io = model.get_module_io_batched(bad_prompts)
else:
print()
print("Calculating per-layer refusal directions...")
print("* Obtaining residuals for good prompts...")
good_residuals = model.get_residuals_batched(good_prompts)
print("* Obtaining residuals for bad prompts...")
bad_residuals = model.get_residuals_batched(bad_prompts)
torch.Tensor.__repr__ = tensor_shape_repr # ty:ignore[invalid-assignment]
good_means = good_residuals.mean(dim=0)
bad_means = bad_residuals.mean(dim=0)
print()
print("Obtaining module I/O for good prompts...")
good_module_io = model.get_module_io_batched(good_prompts)
print("Obtaining module I/O for bad prompts...")
bad_module_io = model.get_module_io_batched(bad_prompts)
refusal_directions = F.normalize(bad_means - good_means, p=2, dim=1)
# print(good_module_io)
if settings.orthogonalize_direction:
# Implements https://huggingface.co/blog/grimjim/projected-abliteration
# Adjust the refusal directions so that only the component that is
# orthogonal to the good direction is subtracted during abliteration.
good_directions = F.normalize(good_means, p=2, dim=1)
projection_vector = torch.sum(refusal_directions * good_directions, dim=1)
refusal_directions = (
refusal_directions - projection_vector.unsqueeze(1) * good_directions
)
refusal_directions = F.normalize(refusal_directions, p=2, dim=1)
print()
print("Performing Arbitrary-Rank Ablation...")
analyzer = Analyzer(settings, model, good_residuals, bad_residuals)
model.ara_abliterate(
good_module_io,
bad_module_io,
0,
len(model.get_layers()),
1.0,
1.0,
1.0,
)
if settings.print_residual_geometry:
analyzer.print_residual_geometry()
print()
print("Evaluating...")
evaluator.get_score()
return
if settings.plot_residuals:
analyzer.plot_residuals()
print()
print("Calculating per-layer refusal directions...")
print("* Obtaining residuals for good prompts...")
good_residuals = model.get_residuals_batched(good_prompts)
print("* Obtaining residuals for bad prompts...")
bad_residuals = model.get_residuals_batched(bad_prompts)
good_means = good_residuals.mean(dim=0)
bad_means = bad_residuals.mean(dim=0)
refusal_directions = F.normalize(bad_means - good_means, p=2, dim=1)
if settings.orthogonalize_direction:
# Implements https://huggingface.co/blog/grimjim/projected-abliteration
# Adjust the refusal directions so that only the component that is
# orthogonal to the good direction is subtracted during abliteration.
good_directions = F.normalize(good_means, p=2, dim=1)
projection_vector = torch.sum(refusal_directions * good_directions, dim=1)
refusal_directions = (
refusal_directions - projection_vector.unsqueeze(1) * good_directions
)
refusal_directions = F.normalize(refusal_directions, p=2, dim=1)
analyzer = Analyzer(settings, model, good_residuals, bad_residuals)
if settings.print_residual_geometry:
analyzer.print_residual_geometry()
if settings.plot_residuals:
analyzer.plot_residuals()
# We don't need the residuals after computing refusal directions.
del good_residuals, bad_residuals, analyzer
empty_cache()
# We don't need the residuals after computing refusal directions.
del good_residuals, bad_residuals, analyzer
empty_cache()
trial_index = 0
start_index = 0
@@ -486,83 +463,126 @@ def run():
trial_index += 1
trial.set_user_attr("index", trial_index)
direction_scope = trial.suggest_categorical(
"direction_scope",
[
"global",
"per layer",
],
)
last_layer_index = len(model.get_layers()) - 1
# Discrimination between "harmful" and "harmless" inputs is usually strongest
# in layers slightly past the midpoint of the layer stack. See the original
# abliteration paper (https://arxiv.org/abs/2406.11717) for a deeper analysis.
#
# Note that we always sample this parameter even though we only need it for
# the "global" direction scope. The reason is that multivariate TPE doesn't
# work with conditional or variable-range parameters.
direction_index = trial.suggest_float(
"direction_index",
0.4 * last_layer_index,
0.9 * last_layer_index,
)
if direction_scope == "per layer":
direction_index = None
parameters = {}
for component in model.get_abliterable_components():
# 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,
if settings.use_ara:
start_layer_index = trial.suggest_int(
"start_layer_index",
0,
len(model.get_layers()) // 3,
)
max_weight_position = trial.suggest_float(
f"{component}.max_weight_position",
0.6 * last_layer_index,
1.0 * last_layer_index,
end_layer_index = trial.suggest_int(
"end_layer_index",
len(model.get_layers()) // 2,
len(model.get_layers()),
)
# For sampling purposes, min_weight is expressed as a fraction of max_weight,
# again because multivariate TPE doesn't support variable-range parameters.
# The value is transformed into the actual min_weight value below.
min_weight = trial.suggest_float(
f"{component}.min_weight",
preserve_good_behavior_weight = trial.suggest_float(
"preserve_good_behavior_weight",
0.0,
1.0,
)
min_weight_distance = trial.suggest_float(
f"{component}.min_weight_distance",
steer_bad_behavior_weight = trial.suggest_float(
"steer_bad_behavior_weight",
0.0,
1.0,
0.6 * last_layer_index,
)
tie_to_original_matrix_weight = trial.suggest_float(
"tie_to_original_matrix_weight",
0.2, # Minimum to prevent "optimizing" away the regularization term.
1.0,
)
else:
direction_scope = trial.suggest_categorical(
"direction_scope",
[
"global",
"per layer",
],
)
parameters[component] = AbliterationParameters(
max_weight=max_weight,
max_weight_position=max_weight_position,
min_weight=(min_weight * max_weight),
min_weight_distance=min_weight_distance,
last_layer_index = len(model.get_layers()) - 1
# Discrimination between "harmful" and "harmless" inputs is usually strongest
# in layers slightly past the midpoint of the layer stack. See the original
# abliteration paper (https://arxiv.org/abs/2406.11717) for a deeper analysis.
#
# Note that we always sample this parameter even though we only need it for
# the "global" direction scope. The reason is that multivariate TPE doesn't
# work with conditional or variable-range parameters.
direction_index = trial.suggest_float(
"direction_index",
0.4 * last_layer_index,
0.9 * last_layer_index,
)
trial.set_user_attr("direction_index", direction_index)
trial.set_user_attr("parameters", {k: asdict(v) for k, v in parameters.items()})
if direction_scope == "per layer":
direction_index = None
parameters = {}
for component in model.get_abliterable_components():
# 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,
)
max_weight_position = trial.suggest_float(
f"{component}.max_weight_position",
0.6 * last_layer_index,
1.0 * last_layer_index,
)
# For sampling purposes, min_weight is expressed as a fraction of max_weight,
# again because multivariate TPE doesn't support variable-range parameters.
# The value is transformed into the actual min_weight value below.
min_weight = trial.suggest_float(
f"{component}.min_weight",
0.0,
1.0,
)
min_weight_distance = trial.suggest_float(
f"{component}.min_weight_distance",
1.0,
0.6 * last_layer_index,
)
parameters[component] = AbliterationParameters(
max_weight=max_weight,
max_weight_position=max_weight_position,
min_weight=(min_weight * max_weight),
min_weight_distance=min_weight_distance,
)
trial.set_user_attr("direction_index", direction_index)
trial.set_user_attr(
"parameters", {k: asdict(v) for k, v in parameters.items()}
)
print()
print(
f"Running trial [bold]{trial_index}[/] of [bold]{settings.n_trials}[/]..."
)
print("* Parameters:")
for name, value in get_trial_parameters(trial).items():
for name, value in get_trial_parameters(settings, trial).items():
print(f" * {name} = [bold]{value}[/]")
print("* Resetting model...")
model.reset_model()
print("* Abliterating...")
model.abliterate(refusal_directions, direction_index, parameters)
if settings.use_ara:
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,
preserve_good_behavior_weight,
steer_bad_behavior_weight,
tie_to_original_matrix_weight,
)
else:
print("* Resetting model...")
model.reset_model()
print("* Abliterating...")
model.abliterate(refusal_directions, direction_index, parameters)
print("* Evaluating...")
score, kl_divergence, refusals = evaluator.get_score()
@@ -739,19 +759,33 @@ def run():
print()
print(f"Restoring model from trial [bold]{trial.user_attrs['index']}[/]...")
print("* Parameters:")
for name, value in get_trial_parameters(trial).items():
for name, value in get_trial_parameters(settings, trial).items():
print(f" * {name} = [bold]{value}[/]")
print("* Resetting model...")
model.reset_model()
print("* Abliterating...")
model.abliterate(
refusal_directions,
trial.user_attrs["direction_index"],
{
k: AbliterationParameters(**v)
for k, v in trial.user_attrs["parameters"].items()
},
)
if settings.use_ara:
print("* Reloading model...")
model.reset_model()
print("* Abliterating (Arbitrary-Rank Ablation)...")
model.ara_abliterate(
good_module_io,
bad_module_io,
trial.params["start_layer_index"],
trial.params["end_layer_index"],
trial.params["preserve_good_behavior_weight"],
trial.params["steer_bad_behavior_weight"],
trial.params["tie_to_original_matrix_weight"],
)
else:
print("* Resetting model...")
model.reset_model()
print("* Abliterating...")
model.abliterate(
refusal_directions,
trial.user_attrs["direction_index"],
{
k: AbliterationParameters(**v)
for k, v in trial.user_attrs["parameters"].items()
},
)
while True:
print()
@@ -786,8 +820,12 @@ def run():
print("Saving LoRA adapter...")
model.model.save_pretrained(save_directory)
else:
print("Saving merged model...")
merged_model = model.get_merged_model()
if settings.use_ara:
print("Saving model...")
merged_model = model.model
else:
print("Saving merged model...")
merged_model = model.get_merged_model()
merged_model.save_pretrained(save_directory)
del merged_model
empty_cache()
@@ -839,8 +877,12 @@ def run():
token=token,
)
else:
print("Uploading merged model...")
merged_model = model.get_merged_model()
if settings.use_ara:
print("Uploading model...")
merged_model = model.model
else:
print("Uploading merged model...")
merged_model = model.get_merged_model()
merged_model.push_to_hub(
repo_id,
private=private,
src/heretic/model.py
+23 -6
View File
@@ -153,7 +153,8 @@ class Model:
if self.model is None:
raise Exception("Failed to load model with all configured dtypes.")
# self._apply_lora()
if not settings.use_ara:
self._apply_lora()
# LoRA B matrices are initialized to zero by default in PEFT,
# so we don't need to do anything manually.
@@ -285,7 +286,11 @@ class Model:
performs full model reload with quantization config.
"""
current_model = getattr(self.model.config, "name_or_path", None)
if current_model == self.settings.model and not self.needs_reload:
if (
current_model == self.settings.model
and not self.needs_reload
and not self.settings.use_ara
):
# Reset LoRA adapters to zero (identity transformation)
for name, module in self.model.named_modules():
if "lora_B" in name and hasattr(module, "weight"):
@@ -314,7 +319,8 @@ class Model:
**extra_kwargs,
)
self._apply_lora()
if not self.settings.use_ara:
self._apply_lora()
self.needs_reload = False
@@ -338,6 +344,9 @@ class Model:
modules = {}
def try_add(component: str, module: Any):
if component not in self.settings.target_components:
return
# Only add if it's a proper nn.Module (PEFT can wrap these with LoRA)
if isinstance(module, Module):
if component not in modules:
@@ -564,6 +573,14 @@ class Model:
# On average, the outputs for "bad" prompts should resemble
# the original outputs for "good" prompts (which steers the
# behavior for "bad" prompts towards that for "good" prompts).
#
# TODO: An alternative formulation could use the mean distance
# of "bad" outputs from the boundary of the core cluster
# of original "good" outputs. This would classify an output
# configuration as optimal as long as all "bad" outputs
# are inside the same cluster as the "good" outputs, even
# if their centroid is different from those of the "good"
# outputs.
steer_bad_behavior = (
(
(bad_input @ matrix.T).mean(dim=0)
@@ -602,9 +619,9 @@ class Model:
# Convergence usually happens within 2-3 steps, so this is more than enough.
for step in range(5):
loss = optimizer.step(closure)
print(
f"\\[{layer_index}/{component}/{module_index}] Step: {step}, Loss: {loss.item():.6f}"
)
# print(
# f"\\[{layer_index}/{component}/{module_index}] Step: {step}, Loss: {loss.item():.6f}"
# )
def generate(
self,
src/heretic/utils.py
+20 -11
View File
@@ -250,19 +250,28 @@ def empty_cache():
gc.collect()
def get_trial_parameters(trial: Trial) -> dict[str, str]:
params = {}
def get_trial_parameters(settings: Settings, trial: Trial) -> dict[str, str]:
if settings.use_ara:
return {
"start_layer_index": f"{trial.params['start_layer_index']}",
"end_layer_index": f"{trial.params['end_layer_index']}",
"preserve_good_behavior_weight": f"{trial.params['preserve_good_behavior_weight']:.4f}",
"steer_bad_behavior_weight": f"{trial.params['steer_bad_behavior_weight']:.4f}",
"tie_to_original_matrix_weight": f"{trial.params['tie_to_original_matrix_weight']:.4f}",
}
else:
params = {}
direction_index = trial.user_attrs["direction_index"]
params["direction_index"] = (
"per layer" if (direction_index is None) else f"{direction_index:.2f}"
)
direction_index = trial.user_attrs["direction_index"]
params["direction_index"] = (
"per layer" if (direction_index is None) else f"{direction_index:.2f}"
)
for component, parameters in trial.user_attrs["parameters"].items():
for name, value in parameters.items():
params[f"{component}.{name}"] = f"{value:.2f}"
for component, parameters in trial.user_attrs["parameters"].items():
for name, value in parameters.items():
params[f"{component}.{name}"] = f"{value:.2f}"
return params
return params
def get_readme_intro(
@@ -285,7 +294,7 @@ def get_readme_intro(
chr(10).join(
[
f"| **{name}** | {value} |"
for name, value in get_trial_parameters(trial).items()
for name, value in get_trial_parameters(settings, trial).items()
]
)
}