From 154241f8a20294f98ac78258480d9d125e2b6389 Mon Sep 17 00:00:00 2001 From: Philipp Emanuel Weidmann Date: Fri, 27 Feb 2026 14:16:27 +0530 Subject: [PATCH] feat(ara): implement matrix optimization --- src/heretic/main.py | 22 +++++++++- src/heretic/model.py | 100 ++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 113 insertions(+), 9 deletions(-) diff --git a/src/heretic/main.py b/src/heretic/main.py index 6467db3..1f25137 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -207,7 +207,7 @@ def run(): ) # We don't need gradients as we only do inference. - torch.set_grad_enabled(False) + # 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) @@ -422,7 +422,25 @@ def run(): print("Obtaining module I/O for bad prompts...") bad_module_io = model.get_module_io_batched(bad_prompts) - print(good_module_io) + # print(good_module_io) + + print() + print("Performing Arbitrary-Rank Ablation...") + + model.ara_abliterate( + good_module_io, + bad_module_io, + 0, + len(model.get_layers()), + 1.0, + 1.0, + 1.0, + ) + + print() + print("Evaluating...") + evaluator.get_score() + return print() print("Calculating per-layer refusal directions...") diff --git a/src/heretic/model.py b/src/heretic/model.py index 4f1e998..7b0b1c5 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -4,7 +4,7 @@ import math from contextlib import suppress from dataclasses import dataclass -from typing import Any, Callable, Type, cast +from typing import Any, Callable, Type, TypeAlias, cast import bitsandbytes as bnb import torch @@ -14,6 +14,7 @@ from peft import LoraConfig, PeftModel, get_peft_model from peft.tuners.lora.layer import Linear from torch import FloatTensor, LongTensor, Tensor from torch.nn import Module, ModuleList +from torch.optim import LBFGS from torch.utils.hooks import RemovableHandle from transformers import ( AutoModelForCausalLM, @@ -53,6 +54,13 @@ class AbliterationParameters: min_weight_distance: float +# 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 +# tensors of shape (prompt, component). +ModuleIO: TypeAlias = list[dict[str, dict[int, tuple[Tensor, Tensor]]]] + + class Model: model: PreTrainedModel | PeftModel tokenizer: PreTrainedTokenizerBase @@ -145,7 +153,7 @@ class Model: if self.model is None: raise Exception("Failed to load model with all configured dtypes.") - self._apply_lora() + # self._apply_lora() # LoRA B matrices are initialized to zero by default in PEFT, # so we don't need to do anything manually. @@ -520,6 +528,84 @@ class Model: weight_A.data = lora_A.to(weight_A.dtype) weight_B.data = lora_B.to(weight_B.dtype) + def ara_abliterate( + self, + good_module_io: ModuleIO, + bad_module_io: ModuleIO, + start_layer_index: int, + end_layer_index: int, + preserve_good_behavior_weight: float, + steer_bad_behavior_weight: float, + tie_to_original_matrix_weight: float, + ): + 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. + module = cast(Linear, module) + + matrix = module.weight + original_matrix = matrix.detach().clone() + + good_input, good_output = good_module_io[layer_index][component][ + module_index + ] + bad_input, bad_output = bad_module_io[layer_index][component][ + module_index + ] + + def objective(matrix: Tensor) -> Tensor: + # The results of applying the operator to inputs associated + # with "good" prompts should change as little as possible. + preserve_good_behavior = ( + (good_input @ matrix.T - good_output) ** 2 + ).mean() + + # 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). + steer_bad_behavior = ( + ( + (bad_input @ matrix.T).mean(dim=0) + - good_output.mean(dim=0) + ) + ** 2 + ).mean() + + # The matrix itself should change as little as possible overall. + # This prevents overfitting due to underdetermination of the + # optimization problem from a relatively small number of I/O pairs. + tie_to_original_matrix = ( + (matrix - original_matrix) ** 2 + ).mean() + + return ( + preserve_good_behavior_weight * preserve_good_behavior + + steer_bad_behavior_weight * steer_bad_behavior + + tie_to_original_matrix_weight * tie_to_original_matrix + ) + + optimizer = LBFGS( + [matrix], + lr=1.0, + max_iter=20, # Number of internal iterations per step, *not* the number of steps. + history_size=10, + line_search_fn="strong_wolfe", + ) + + def closure() -> Tensor: + optimizer.zero_grad() + loss = objective(matrix) + loss.backward() + return loss + + # 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}" + ) + def generate( self, prompts: list[Prompt], @@ -657,12 +743,12 @@ class Model: def get_module_io( self, prompts: list[Prompt], - ) -> list[dict[str, dict[int, tuple[Tensor, Tensor]]]]: + ) -> ModuleIO: # 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 # tensors of shape (prompt, component). - module_io: list[dict[str, dict[int, tuple[Tensor, Tensor]]]] = [] + module_io: ModuleIO = [] def get_hook( layer_index: int, @@ -722,19 +808,19 @@ class Model: def get_module_io_batched( self, prompts: list[Prompt], - ) -> list[dict[str, dict[int, tuple[Tensor, Tensor]]]]: + ) -> ModuleIO: # Aggregating batch results is more complicated for module I/O # than for other get_*_batched methods, because the structure of the results # might differ between batches, as whether individual modules activate # can depend on the prompt (in particular for MoE models). # In practice, inhomogeneous results should be very rare, but to be fully # generic, this logic is required. - module_io_batches = [ + module_io_batches: list[ModuleIO] = [ self.get_module_io(batch) for batch in batchify(prompts, self.settings.batch_size) ] - module_io: list[dict[str, dict[int, tuple[Tensor, Tensor]]]] = [] + module_io: ModuleIO = [] for layer_index in range(len(self.get_layers())): module_io.append({})