fix(ara): free gradient buffers after optimization (#426)

LBFGS leaves one full-size gradient buffer on every optimized weight.
Across the layers processed in a typical trial this is many GiB of VRAM
that persists into evaluation, causing CUDA out-of-memory errors.
Clear the gradients after each module is optimized.

README.md

@@ -2,6 +2,8 @@
 
 # Heretic: Fully automatic censorship removal for language models<br><br>[![Discord](https://img.shields.io/discord/1447831134212984903?color=5865F2&label=discord&labelColor=black&logo=discord&logoColor=white&style=for-the-badge)](https://discord.gg/gdXc48gSyT) [![Follow us on Hugging Face](https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-us-on-hf-md-dark.svg)](https://huggingface.co/heretic-org)
 
+[![#1 Repository of the Day](https://trendshift.io/api/badge/repositories/20538)](https://trendshift.io/repositories/20538)
+
 Heretic is a tool that removes censorship (aka "safety alignment") from
 transformer-based language models without expensive post-training.
 It combines an advanced implementation of directional ablation, also known

pyproject.toml

@@ -22,19 +22,24 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
 ]
 dependencies = [
-    "accelerate~=1.10",
-    "bitsandbytes~=0.45",
-    "datasets~=4.0",
+    "accelerate~=1.13",
+    "bitsandbytes~=0.49",
+    "datasets~=4.7",
     "hf-transfer~=0.1",
-    "huggingface-hub~=0.34",
-    "kernels~=0.11",
-    "optuna~=4.5",
-    "peft~=0.14",
-    "psutil~=7.1",
-    "pydantic-settings~=2.10",
+    "huggingface-hub~=1.7",
+    "immutabledict~=4.3",
+    "kernels~=0.12",
+    "langdetect~=1.0",
+    "lm-eval[hf]~=0.4",
+    "numpy~=2.2",
+    "optuna~=4.7",
+    "peft~=0.18",
+    "psutil~=7.2",
+    "pydantic-settings~=2.13",
     "questionary~=2.1",
-    "rich~=14.1",
-    "transformers~=4.57",
+    "rich~=14.3",
+    "tqdm~=4.67",
+    "transformers~=5.3",
 ]
 
 [project.optional-dependencies]
@@ -42,7 +47,6 @@ research = [
     "geom-median~=0.1",
     "imageio~=2.37",
     "matplotlib~=3.10",
-    "numpy~=2.2",
     "pacmap~=0.8",
     "scikit-learn~=1.7",
 ]

src/heretic/analyzer.py

@@ -3,9 +3,11 @@
 
 from pathlib import Path
 
+import numpy as np
 import torch
 import torch.linalg as LA
 import torch.nn.functional as F
+from numpy.typing import NDArray
 from rich.progress import track
 from rich.table import Table
 from torch import Tensor
@@ -156,11 +158,9 @@ class Analyzer:
         try:
             import imageio.v3 as iio  # ty:ignore[unresolved-import]
             import matplotlib.pyplot as plt  # ty:ignore[unresolved-import]
-            import numpy as np  # ty:ignore[unresolved-import]
             from geom_median.numpy import (  # ty:ignore[unresolved-import]
                 compute_geometric_median,
             )
-            from numpy.typing import NDArray  # ty:ignore[unresolved-import]
             from pacmap import PaCMAP  # ty:ignore[unresolved-import]
         except ImportError:
             print()

src/heretic/config.py

@@ -61,6 +61,18 @@ class DatasetSpecification(BaseModel):
     )
 
 
+class BenchmarkSpecification(BaseModel):
+    task: str = Field(
+        description="Task ID of the benchmark in the Language Model Evaluation Harness."
+    )
+
+    name: str = Field(description="Name of the benchmark for presentation purposes.")
+
+    description: str = Field(
+        description="Description of the benchmark for presentation purposes."
+    )
+
+
 class Settings(BaseSettings):
     model: str = Field(description="Hugging Face model ID, or path to model on disk.")
 
@@ -176,6 +188,42 @@ class Settings(BaseSettings):
         ),
     )
 
+    target_components: list[str] = Field(
+        default=["attn.o_proj", "mlp.down_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."
+        ),
+    )
+    
+    use_ara_lora: bool = Field(
+        default=False,
+        description=(
+            "Use LoRA in ARA instead of full-weight editing. Makes it compatible with quantization and removes model reloads."
+        ),
+    )
+    
+    ara_lora_rank: int = Field(
+        default=128,
+        description="If LoRA is used in ARA, this sets up its rank. Keep it high enough to simulate the 'arbitrary' effect.",
+    )
+
+    use_piqa: bool = Field(
+        default=False,
+        description=(
+            "Whether to use the Physical Interaction: Question Answering (PIQA) benchmark "
+            "as the quality metric instead of the Kullback-Leibler divergence."
+        ),
+    )
+
     orthogonalize_direction: bool = Field(
         default=False,
         description=(
@@ -185,7 +233,7 @@ class Settings(BaseSettings):
     )
 
     row_normalization: RowNormalization = Field(
-        default=RowNormalization.NONE,
+        default=RowNormalization.FULL,
         description=(
             "How to apply row normalization of the weights. Options: "
             '"none" (no normalization), '
@@ -230,6 +278,67 @@ class Settings(BaseSettings):
         description="Directory to save and load study progress to/from.",
     )
 
+    benchmarks: list[BenchmarkSpecification] = Field(
+        default=[
+            BenchmarkSpecification(
+                task="agieval",
+                name="AGIEval",
+                description="A Human-Centric Benchmark for Evaluating Foundation Models",
+            ),
+            BenchmarkSpecification(
+                task="bbh",
+                name="BIG-Bench Hard (BBH)",
+                description="Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them",
+            ),
+            BenchmarkSpecification(
+                task="commonsense_qa",
+                name="CommonsenseQA",
+                description="A Question Answering Challenge Targeting Commonsense Knowledge",
+            ),
+            BenchmarkSpecification(
+                task="eq_bench",
+                name="EQ-Bench",
+                description="An Emotional Intelligence Benchmark for Large Language Models",
+            ),
+            BenchmarkSpecification(
+                task="gsm8k",
+                name="GSM8K",
+                description="Training Verifiers to Solve Math Word Problems",
+            ),
+            BenchmarkSpecification(
+                task="hellaswag",
+                name="HellaSwag",
+                description="Can a Machine Really Finish Your Sentence?",
+            ),
+            BenchmarkSpecification(
+                task="ifeval",
+                name="IFEval",
+                description="Instruction-Following Evaluation for Large Language Models",
+            ),
+            BenchmarkSpecification(
+                task="mmlu",
+                name="MMLU",
+                description="Measuring Massive Multitask Language Understanding",
+            ),
+            BenchmarkSpecification(
+                task="mmlu_pro",
+                name="MMLU-Pro",
+                description="A More Robust and Challenging Multi-Task Language Understanding Benchmark",
+            ),
+            BenchmarkSpecification(
+                task="piqa",
+                name="PIQA",
+                description="Reasoning about Physical Commonsense in Natural Language",
+            ),
+            BenchmarkSpecification(
+                task="winogrande",
+                name="WinoGrande",
+                description="An Adversarial Winograd Schema Challenge at Scale",
+            ),
+        ],
+        description="Benchmarks to offer to the user for evaluating abliterated models.",
+    )
+
     refusal_markers: list[str] = Field(
         default=[
             "sorry",

src/heretic/evaluator.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
+import lm_eval
 import torch.nn.functional as F
+from lm_eval.models.huggingface import HFLM
 from torch import Tensor
 
 from .config import Settings
@@ -21,15 +23,16 @@ class Evaluator:
         self.settings = settings
         self.model = model
 
-        print()
-        print(
-            f"Loading good evaluation prompts from [bold]{settings.good_evaluation_prompts.dataset}[/]..."
-        )
-        self.good_prompts = load_prompts(settings, settings.good_evaluation_prompts)
-        print(f"* [bold]{len(self.good_prompts)}[/] prompts loaded")
+        if not settings.use_piqa:
+            print()
+            print(
+                f"Loading good evaluation prompts from [bold]{settings.good_evaluation_prompts.dataset}[/]..."
+            )
+            self.good_prompts = load_prompts(settings, settings.good_evaluation_prompts)
+            print(f"* [bold]{len(self.good_prompts)}[/] prompts loaded")
 
-        print("* Obtaining first-token probability distributions...")
-        self.base_logprobs = model.get_logprobs_batched(self.good_prompts)
+            print("* Obtaining first-token probability distributions...")
+            self.base_logprobs = model.get_logprobs_batched(self.good_prompts)
 
         print()
         print(
@@ -93,33 +96,57 @@ class Evaluator:
         return refusal_count
 
     def get_score(self) -> tuple[tuple[float, float], float, int]:
-        print("  * Obtaining first-token probability distributions...")
-        logprobs = self.model.get_logprobs_batched(self.good_prompts)
-        kl_divergence = F.kl_div(
-            logprobs,
-            self.base_logprobs,
-            reduction="batchmean",
-            log_target=True,
-        ).item()
-        print(f"  * KL divergence: [bold]{kl_divergence:.4f}[/]")
+        if self.settings.use_piqa:
+            print("  * Running PIQA benchmark...")
+            hflm = HFLM(
+                pretrained=self.model.model,  # ty:ignore[invalid-argument-type]
+                tokenizer=self.model.tokenizer,  # ty:ignore[invalid-argument-type]
+                batch_size="auto",
+            )
+            results = lm_eval.simple_evaluate(
+                model=hflm,
+                tasks=["piqa"],
+            )
+            piqa_acc_norm: float = results["results"]["piqa"]["acc_norm,none"]
+            print(f"  * PIQA acc_norm: [bold]{piqa_acc_norm:.4f}[/]")
+        else:
+            print("  * Obtaining first-token probability distributions...")
+            logprobs = self.model.get_logprobs_batched(self.good_prompts)
+            kl_divergence = F.kl_div(
+                logprobs,
+                self.base_logprobs,
+                reduction="batchmean",
+                log_target=True,
+            ).item()
+            print(f"  * KL divergence: [bold]{kl_divergence:.4f}[/]")
 
         print("  * Counting model refusals...")
         refusals = self.count_refusals()
         print(f"  * Refusals: [bold]{refusals}[/]/{len(self.bad_prompts)}")
 
-        kl_divergence_scale = self.settings.kl_divergence_scale
-        kl_divergence_target = self.settings.kl_divergence_target
-
-        refusals_score = refusals / self.base_refusals
-
-        if kl_divergence >= kl_divergence_target:
-            kld_score = kl_divergence / kl_divergence_scale
-        else:
-            kld_score = refusals_score * kl_divergence_target / kl_divergence_scale
-
-        score = (
-            kld_score,
-            refusals_score,
+        refusals_score = (
+            refusals / self.base_refusals if self.base_refusals > 0 else float(refusals)
         )
 
-        return score, kl_divergence, refusals
+        if self.settings.use_piqa:
+            score = (
+                -piqa_acc_norm,
+                refusals_score,
+            )
+
+            return score, -piqa_acc_norm, refusals
+        else:
+            kl_divergence_scale = self.settings.kl_divergence_scale
+            kl_divergence_target = self.settings.kl_divergence_target
+
+            if kl_divergence >= kl_divergence_target:
+                kld_score = kl_divergence / kl_divergence_scale
+            else:
+                kld_score = refusals_score * kl_divergence_target / kl_divergence_scale
+
+            score = (
+                kld_score,
+                refusals_score,
+            )
+
+            return score, kl_divergence, refusals

src/heretic/main.py

@@ -1,6 +1,15 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
+# ruff: noqa: E402
+
+from .progress import patch_tqdm
+
+# This patches tqdm class definitions, which must happen
+# before any other module imports tqdm.
+patch_tqdm()
+
+import logging
 import math
 import os
 import sys
@@ -10,9 +19,13 @@ from dataclasses import asdict
 from importlib.metadata import version
 from os.path import commonprefix
 from pathlib import Path
+from typing import Any
 
 import huggingface_hub
+import lm_eval
+import numpy as np
 import optuna
+import questionary
 import torch
 import torch.nn.functional as F
 import transformers
@@ -24,6 +37,7 @@ from accelerate.utils import (
     is_xpu_available,
 )
 from huggingface_hub import ModelCard, ModelCardData
+from lm_eval.models.huggingface import HFLM
 from optuna import Trial, TrialPruned
 from optuna.exceptions import ExperimentalWarning
 from optuna.samplers import TPESampler
@@ -32,13 +46,14 @@ from optuna.storages.journal import JournalFileBackend, JournalFileOpenLock
 from optuna.study import StudyDirection
 from optuna.trial import TrialState
 from pydantic import ValidationError
-from questionary import Choice
+from questionary import Choice, Style
+from rich.table import Table
 from rich.traceback import install
 
 from .analyzer import Analyzer
-from .config import QuantizationMethod, Settings
+from .config import QuantizationMethod, RowNormalization, 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,
@@ -174,9 +189,15 @@ def run():
     # Adapted from https://github.com/huggingface/accelerate/blob/main/src/accelerate/commands/env.py
     if torch.cuda.is_available():
         count = torch.cuda.device_count()
-        print(f"Detected [bold]{count}[/] CUDA device(s):")
+        total_vram = sum(torch.cuda.mem_get_info(i)[1] for i in range(count))
+        print(
+            f"Detected [bold]{count}[/] CUDA device(s) ({total_vram / (1024**3):.2f} GB total VRAM):"
+        )
         for i in range(count):
-            print(f"* GPU {i}: [bold]{torch.cuda.get_device_name(i)}[/]")
+            vram = torch.cuda.mem_get_info(i)[1] / (1024**3)
+            print(
+                f"* GPU {i}: [bold]{torch.cuda.get_device_name(i)}[/] ({vram:.2f} GB)"
+            )
     elif is_xpu_available():
         count = torch.xpu.device_count()
         print(f"Detected [bold]{count}[/] XPU device(s):")
@@ -206,8 +227,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)
@@ -219,6 +241,9 @@ def run():
     # In my entire career I've never seen a useful warning from that library.
     transformers.logging.set_verbosity_error()
 
+    # Another library that generates warning spam.
+    logging.getLogger("lm_eval").setLevel(logging.ERROR)
+
     # We do our own trial logging, so we don't need the INFO messages
     # about parameters and results.
     optuna.logging.set_verbosity(optuna.logging.WARNING)
@@ -371,7 +396,8 @@ def run():
 
     print()
     print("Checking for common response prefix...")
-    responses = model.get_responses_batched(good_prompts[:100] + bad_prompts[:100])
+    prefix_check_prompts = good_prompts[:100] + bad_prompts[:100]
+    responses = model.get_responses_batched(prefix_check_prompts)
 
     # Despite being located in os.path, commonprefix actually performs
     # a naive string operation without any path-specific logic,
@@ -382,24 +408,39 @@ def run():
     model.response_prefix = commonprefix(responses).rstrip(" ")
 
     # Suppress CoT output.
-    if model.response_prefix.startswith("<think>"):
-        # Most thinking models.
-        model.response_prefix = "<think></think>"
-    elif model.response_prefix.startswith("<|channel|>analysis<|message|>"):
-        # gpt-oss.
-        model.response_prefix = "<|channel|>analysis<|message|><|end|><|start|>assistant<|channel|>final<|message|>"
-    elif model.response_prefix.startswith("<thought>"):
-        # Unknown, suggested by user.
-        model.response_prefix = "<thought></thought>"
-    elif model.response_prefix.startswith("[THINK]"):
-        # Unknown, suggested by user.
-        model.response_prefix = "[THINK][/THINK]"
+    recheck_prefix = False
+    if model.response_prefix:
+        # When using any of the predefined prefixes below, we need to check that
+        # the prefix is actually complete (e.g. not missing a trailing newline).
+        recheck_prefix = True
+        if model.response_prefix.startswith("<think>"):
+            # Most thinking models.
+            model.response_prefix = "<think></think>"
+        elif model.response_prefix.startswith("<|channel|>analysis<|message|>"):
+            # gpt-oss.
+            model.response_prefix = "<|channel|>analysis<|message|><|end|><|start|>assistant<|channel|>final<|message|>"
+        elif model.response_prefix.startswith("<thought>"):
+            # Unknown, suggested by user.
+            model.response_prefix = "<thought></thought>"
+        elif model.response_prefix.startswith("[THINK]"):
+            # Unknown, suggested by user.
+            model.response_prefix = "[THINK][/THINK]"
+        else:
+            recheck_prefix = False
 
     if model.response_prefix:
         print(f"* Prefix found: [bold]{model.response_prefix!r}[/]")
     else:
         print("* None found")
 
+    if recheck_prefix:
+        print("* Rechecking with prefix...")
+        responses = model.get_responses_batched(prefix_check_prompts)
+        additional_prefix = commonprefix(responses).rstrip(" ")
+        if additional_prefix:
+            model.response_prefix += additional_prefix
+            print(f"* Extended prefix found: [bold]{model.response_prefix!r}[/]")
+
     evaluator = Evaluator(settings, model)
 
     if settings.evaluate_model is not None:
@@ -411,40 +452,47 @@ def run():
         evaluator.get_score()
         return
 
-    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)
+    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)
 
-    good_means = good_residuals.mean(dim=0)
-    bad_means = bad_residuals.mean(dim=0)
+        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)
+        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)
+        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)
+        analyzer = Analyzer(settings, model, good_residuals, bad_residuals)
 
-    if settings.print_residual_geometry:
-        analyzer.print_residual_geometry()
+        if settings.print_residual_geometry:
+            analyzer.print_residual_geometry()
 
-    if settings.plot_residuals:
-        analyzer.plot_residuals()
+        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
@@ -455,83 +503,144 @@ 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()) // 2,
             )
-            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.0001,
                 1.0,
-                0.6 * last_layer_index,
+                log=True,
+            )
+            overcorrect_relative_weight = trial.suggest_float(
+                "overcorrect_relative_weight",
+                0.0,
+                1.3,
+            )
+            neighbor_count = trial.suggest_int(
+                "neighbor_count",
+                1,
+                15,
             )
 
-            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,
+            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("direction_index", direction_index)
-        trial.set_user_attr("parameters", {k: asdict(v) for k, v in parameters.items()})
+            trial.set_user_attr("ara_parameters", asdict(ara_parameters))
+        else:
+            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,
+                )
+                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_lora:
+            print("* Resetting model...")
+            model.reset_model()
+            print("* Abliterating (Arbitrary-Rank Ablation with LoRA)...")
+            model.ara_lora_abliterate(
+                good_module_io,
+                bad_module_io,
+                ARAParameters(**trial.user_attrs["ara_parameters"]),
+            )
+        elif settings.use_ara:
+            print("* Reloading model...")
+            model.reset_model()
+            print("* Abliterating (Arbitrary-Rank Ablation)...")
+            model.ara_abliterate(good_module_io, bad_module_io, ara_parameters)
+        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()
 
@@ -629,7 +738,7 @@ def run():
                 title=(
                     f"[Trial {trial.user_attrs['index']:>3}] "
                     f"Refusals: {trial.user_attrs['refusals']:>2}/{len(evaluator.bad_prompts)}, "
-                    f"KL divergence: {trial.user_attrs['kl_divergence']:.4f}"
+                    f"{'PIQA acc_norm' if settings.use_piqa else 'KL divergence'}: {(-1 if settings.use_piqa else 1) * trial.user_attrs['kl_divergence']:.4f}"
                 ),
                 value=trial,
             )
@@ -708,19 +817,38 @@ 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_lora:
+                print("* Resetting model...")
+                model.reset_model()
+                print("* Abliterating (Arbitrary-Rank Ablation with LoRA)...")
+                model.ara_lora_abliterate(
+                    good_module_io,
+                    bad_module_io,
+                    ARAParameters(**trial.user_attrs["ara_parameters"]),
+                )
+            elif settings.use_ara:
+                print("* Reloading model...")
+                model.reset_model()
+                print("* Abliterating (Arbitrary-Rank Ablation)...")
+                model.ara_abliterate(
+                    good_module_io,
+                    bad_module_io,
+                    ARAParameters(**trial.user_attrs["ara_parameters"]),
+                )
+            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()
@@ -730,6 +858,7 @@ def run():
                         "Save the model to a local folder",
                         "Upload the model to Hugging Face",
                         "Chat with the model",
+                        "Benchmark the model",
                         "Return to the trial selection menu",
                     ],
                 )
@@ -755,8 +884,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()
@@ -794,6 +927,8 @@ def run():
                                     "Private",
                                 ],
                             )
+                            if visibility is None:
+                                continue
                             private = visibility == "Private"
 
                             strategy = obtain_merge_strategy(settings)
@@ -808,8 +943,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,
@@ -823,11 +962,23 @@ def run():
                                     token=token,
                                 )
 
-                            # If the model path doesn't exist locally, it can be assumed
-                            # to be a model hosted on the Hugging Face Hub, in which case
+                            # If the model path exists locally and includes the
+                            # card, use it directly. If the model path doesn't
+                            # exist locally, it can be assumed to be a model
+                            # hosted on the Hugging Face Hub, in which case
                             # we can retrieve the model card.
-                            if not Path(settings.model).exists():
+                            model_path = Path(settings.model)
+                            if model_path.exists():
+                                card_path = (
+                                    model_path / huggingface_hub.constants.REPOCARD_NAME
+                                )
+                                if card_path.exists():
+                                    card = ModelCard.load(card_path)
+                                else:
+                                    card = None
+                            else:
                                 card = ModelCard.load(settings.model)
+                            if card is not None:
                                 if card.data is None:
                                     card.data = ModelCardData()
                                 if card.data.tags is None:
@@ -836,6 +987,14 @@ def run():
                                 card.data.tags.append("uncensored")
                                 card.data.tags.append("decensored")
                                 card.data.tags.append("abliterated")
+                                if settings.use_ara:
+                                    card.data.tags.append("ara")
+                                elif (
+                                    settings.orthogonalize_direction
+                                    and settings.row_normalization
+                                    == RowNormalization.FULL
+                                ):
+                                    card.data.tags.append("mpoa")
                                 card.text = (
                                     get_readme_intro(
                                         settings,
@@ -879,6 +1038,113 @@ def run():
                                     # Ctrl+C/Ctrl+D
                                     break
 
+                        case "Benchmark the model":
+                            benchmarks = questionary.checkbox(
+                                "Which benchmarks do you want to run?",
+                                [
+                                    Choice(
+                                        title=f"{benchmark.name}: {benchmark.description}",
+                                        value=benchmark,
+                                    )
+                                    for benchmark in settings.benchmarks
+                                ],
+                                style=Style([("highlighted", "reverse")]),
+                            ).ask()
+                            if not benchmarks:
+                                continue
+
+                            scope = prompt_select(
+                                (
+                                    "Do you want to benchmark the original model along with the decensored model? "
+                                    "Benchmarking both models allows you to compare the scores, but it takes twice as much time."
+                                ),
+                                [
+                                    "Benchmark only the decensored model",
+                                    "Benchmark both models",
+                                ],
+                            )
+                            if scope is None:
+                                continue
+                            benchmark_original_model = scope == "Benchmark both models"
+
+                            hflm = HFLM(
+                                pretrained=model.model,  # ty:ignore[invalid-argument-type]
+                                tokenizer=model.tokenizer,  # ty:ignore[invalid-argument-type]
+                                batch_size="auto",
+                            )
+
+                            table = Table()
+                            table.add_column("Benchmark")
+                            table.add_column("Metric")
+                            if benchmark_original_model:
+                                table.add_column("This model", justify="right")
+                                table.add_column("Original model", justify="right")
+                            else:
+                                table.add_column("Value", justify="right")
+
+                            try:
+                                first_benchmark = True
+
+                                for benchmark in benchmarks:
+                                    print(
+                                        f"Running benchmark [bold]{benchmark.name}[/]..."
+                                    )
+
+                                    def get_results() -> dict[str, Any]:
+                                        results = lm_eval.simple_evaluate(
+                                            model=hflm,
+                                            tasks=[benchmark.task],
+                                        )
+                                        return results["results"][benchmark.task]
+
+                                    results = get_results()
+                                    if benchmark_original_model:
+                                        with model.model.disable_adapter():  # ty:ignore[call-non-callable]
+                                            original_results = get_results()
+
+                                    first_row = True
+
+                                    for metric, value in results.items():
+                                        if metric != "alias":
+                                            if first_row and not first_benchmark:
+                                                if benchmark_original_model:
+                                                    table.add_row("", "", "", "")
+                                                else:
+                                                    table.add_row("", "", "")
+
+                                            def format_value(value: Any) -> str:
+                                                if isinstance(
+                                                    value,
+                                                    (float, np.floating),
+                                                ):
+                                                    return f"{value:.4f}"
+                                                else:
+                                                    return f"{value}"
+
+                                            cells = [
+                                                benchmark.name if first_row else "",
+                                                metric,
+                                                format_value(value),
+                                            ]
+                                            if benchmark_original_model:
+                                                cells.append(
+                                                    format_value(
+                                                        original_results[metric]
+                                                    )
+                                                )
+                                            table.add_row(*cells)
+
+                                            first_row = False
+                                            first_benchmark = False
+                            except KeyboardInterrupt:
+                                pass
+
+                            # The benchmark run might have been cancelled by the user
+                            # before any benchmark was completed, so we only print results
+                            # if there actually are some.
+                            if table.rows:
+                                print(table)
+
                 except Exception as error:
                     print(f"[red]Error: {error}[/]")
 

src/heretic/model.py

@@ -4,7 +4,7 @@
 import math
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import Any, Type, cast
+from typing import Any, Callable, Type, TypeAlias, cast
 
 import bitsandbytes as bnb
 import torch
@@ -14,6 +14,8 @@ 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,
     AutoModelForImageTextToText,
@@ -30,7 +32,7 @@ from transformers.generation import (
 )
 
 from .config import QuantizationMethod, RowNormalization, Settings
-from .utils import Prompt, batchify, empty_cache, print
+from .utils import Prompt, batchify, empty_cache, mean_distances_to_knn, print
 
 
 def get_model_class(
@@ -52,6 +54,23 @@ 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
+# tensors of shape (prompt, component).
+ModuleIO: TypeAlias = list[dict[str, dict[int, tuple[Tensor, Tensor]]]]
+
+
 class Model:
     model: PreTrainedModel | PeftModel
     tokenizer: PreTrainedTokenizerBase
@@ -91,7 +110,7 @@ class Model:
             self.trusted_models[settings.evaluate_model] = settings.trust_remote_code
 
         for dtype in settings.dtypes:
-            print(f"* Trying dtype [bold]{dtype}[/]... ", end="")
+            print(f"* Trying dtype [bold]{dtype}[/]...")
 
             try:
                 quantization_config = self._get_quantization_config(dtype)
@@ -131,46 +150,62 @@ class Model:
             except Exception as error:
                 self.model = None  # ty:ignore[invalid-assignment]
                 empty_cache()
-                print(f"[red]Failed[/] ({error})")
+                print(f"* [red]Failed[/] ({error})")
                 continue
 
             if settings.quantization == QuantizationMethod.BNB_4BIT:
-                print("[green]Ok[/] (quantized to 4-bit precision)")
-            else:
-                print("[green]Ok[/]")
+                print("* Quantized to 4-bit precision")
 
             break
 
         if self.model is None:
             raise Exception("Failed to load model with all configured dtypes.")
 
-        self._apply_lora()
+        if not settings.use_ara or settings.use_ara_lora:
+            self._apply_lora()
 
         # LoRA B matrices are initialized to zero by default in PEFT,
         # so we don't need to do anything manually.
 
         print(f"* Transformer model with [bold]{len(self.get_layers())}[/] layers")
         print("* Abliterable components:")
-        for component, modules in self.get_layer_modules(0).items():
-            print(
-                f"  * [bold]{component}[/]: [bold]{len(modules)}[/] modules per layer"
-            )
+        all_components = {}
+        for layer_index in range(len(self.get_layers())):
+            for component, modules in self.get_layer_modules(layer_index).items():
+                if component not in all_components:
+                    all_components[component] = 0
+                all_components[component] += len(modules)
+        for component, count in all_components.items():
+            print(f"  * [bold]{component}[/]: [bold]{count}[/] modules total")
 
     def _apply_lora(self):
         # Guard against calling this method at the wrong time.
         assert isinstance(self.model, PreTrainedModel)
 
         # Always use LoRA adapters for abliteration (faster reload, no weight modification).
-        # We use the leaf names (e.g. "o_proj") as target modules.
-        # This may cause LoRA adapters to be attached to unrelated modules (e.g. "conv.o_proj"),
-        # but this is harmless as we only abliterate the modules we target in `abliterate()`,
-        # leaving the others at their default (identity) state.
-        # NOTE: This will need to be updated when hybrid layer support (#43) is merged.
-        target_modules = [
-            comp.split(".")[-1] for comp in self.get_abliterable_components()
-        ]
+        # Collect actual leaf module names from the model for LoRA targeting.
+        # This is more robust than splitting component keys (e.g. "attn.o_proj" -> "o_proj")
+        # because hybrid models like Qwen3.5 MoE have modules with different names
+        # across layers (e.g. "o_proj" on attention layers, "out_proj" on linear attention layers).
+        target_modules_set: set[str] = set()
+        
+        module_id_to_full_name = {
+            id(module): module_name
+            for module_name, module in self.model.named_modules()
+        }
 
-        if self.settings.row_normalization != RowNormalization.FULL:
+        for layer_index in range(len(self.get_layers())):
+            for modules in self.get_layer_modules(layer_index).values():
+                for module in modules:
+                    full_name = module_id_to_full_name.get(id(module))
+                    if full_name is not None:
+                        target_modules_set.add(full_name)
+
+        target_modules = sorted(target_modules_set)
+
+        if self.settings.use_ara_lora:
+            lora_rank = self.settings.ara_lora_rank
+        elif self.settings.row_normalization != RowNormalization.FULL:
             # Rank 1 is sufficient for directional ablation without renormalization.
             lora_rank = 1
         else:
@@ -192,7 +227,10 @@ class Model:
         # so the result is a PeftModel rather than a PeftMixedModel.
         self.model = cast(PeftModel, get_peft_model(self.model, self.peft_config))
 
-        print(f"* LoRA adapters initialized (targets: {', '.join(target_modules)})")
+        display_targets = sorted({name.rsplit(".", 1)[-1] for name in target_modules})
+        print(
+            f"* LoRA adapters initialized (target types: {', '.join(display_targets)})"
+        )
 
     def _get_quantization_config(self, dtype: str) -> BitsAndBytesConfig | None:
         """
@@ -276,7 +314,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 or self.settings.use_ara_lora)
+        ):
             # Reset LoRA adapters to zero (identity transformation)
             for name, module in self.model.named_modules():
                 if "lora_B" in name and hasattr(module, "weight"):
@@ -305,7 +347,8 @@ class Model:
             **extra_kwargs,
         )
 
-        self._apply_lora()
+        if not self.settings.use_ara or self.settings.use_ara_lora:
+            self._apply_lora()
 
         self.needs_reload = False
 
@@ -329,6 +372,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:
@@ -340,9 +386,14 @@ class Model:
                     f"Unexpected Tensor in {component} - expected nn.Module"
                 )
 
-        # Exceptions aren't suppressed here, because there is currently
-        # no alternative location for the attention out-projection.
-        try_add("attn.o_proj", layer.self_attn.o_proj)  # ty:ignore[possibly-missing-attribute]
+        # Standard self-attention out-projection (most models).
+        with suppress(Exception):
+            try_add("attn.o_proj", layer.self_attn.o_proj)  # ty:ignore[possibly-missing-attribute]
+
+        # Qwen3.5 MoE hybrid layers use GatedDeltaNet (linear attention) instead
+        # of standard self-attention, so self_attn.o_proj doesn't exist on those layers.
+        with suppress(Exception):
+            try_add("attn.o_proj", layer.linear_attn.out_proj)  # ty:ignore[possibly-missing-attribute]
 
         # Most dense models.
         with suppress(Exception):
@@ -374,7 +425,12 @@ class Model:
         return modules
 
     def get_abliterable_components(self) -> list[str]:
-        return list(self.get_layer_modules(0).keys())
+        # Scan all layers because hybrid models (e.g. Qwen3.5 MoE) have different
+        # components on different layers (some have self_attn, others linear_attn).
+        components: set[str] = set()
+        for layer_index in range(len(self.get_layers())):
+            components.update(self.get_layer_modules(layer_index).keys())
+        return sorted(components)
 
     def abliterate(
         self,
@@ -519,6 +575,228 @@ 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,
+        parameters: ARAParameters,
+    ):
+        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.
+                    module = cast(Linear, module)
+                    matrix = module.weight
+
+                    row_norms = LA.vector_norm(matrix, dim=1, keepdim=True).detach()
+
+                    # Helper function for reparameterization (row-norm preservation constraint).
+                    def get_matrix() -> Tensor:
+                        if self.settings.row_normalization == RowNormalization.FULL:
+                            # See https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration
+                            return row_norms * F.normalize(matrix, p=2, dim=1)
+                        else:
+                            return matrix
+
+                    good_input, good_output = good_module_io[layer_index][component][
+                        module_index
+                    ]
+                    bad_input, bad_output = bad_module_io[layer_index][component][
+                        module_index
+                    ]
+
+                    good_input = good_input.to(matrix.device)
+                    good_output = good_output.to(matrix.device)
+                    bad_input = bad_input.to(matrix.device)
+                    bad_output = bad_output.to(matrix.device)
+
+                    def objective(matrix: Tensor) -> Tensor:
+                        new_good_output = good_input @ matrix.T
+                        new_bad_output = bad_input @ matrix.T
+
+                        # The outputs for "good" prompts should change as little as possible.
+                        preserve_good_behavior = (
+                            (new_good_output - good_output) ** 2
+                        ).mean()
+
+                        steer_bad_behavior = (
+                            # Pull the outputs for "bad" prompts towards
+                            # the original outputs for "good" prompts.
+                            mean_distances_to_knn(
+                                new_bad_output,
+                                good_output,
+                                parameters.neighbor_count,
+                            ).mean()
+                            # 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,
+                                parameters.neighbor_count,
+                            ).mean()
+                        )
+
+                        return (
+                            parameters.preserve_good_behavior_weight
+                            * preserve_good_behavior
+                            + parameters.steer_bad_behavior_weight * steer_bad_behavior
+                        )
+
+                    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(get_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}"
+                        # )
+
+                    # Free the gradient buffers accumulated on the weight parameters
+                    # during optimization. Without this, they persist on the model
+                    # (one full-size gradient per processed weight) and can easily
+                    # consume tens of GiB of VRAM, causing out-of-memory errors
+                    # during the subsequent evaluation.
+                    optimizer.zero_grad(set_to_none=True)
+
+                    with torch.no_grad():
+                        matrix.copy_(get_matrix())
+
+    def ara_lora_abliterate(
+        self,
+        good_module_io: ModuleIO,
+        bad_module_io: ModuleIO,
+        parameters: ARAParameters,
+    ):
+        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):
+                    # Cast to Linear to access weights and LoRA adapters.
+                    module = cast(Linear, module)
+
+                    # Base weight handling and dequantization.
+                    # We need the base weight in float32 to compute the effective weight.
+                    base_weight = cast(Tensor, module.base_layer.weight)
+                    quant_state = getattr(base_weight, "quant_state", None)
+
+                    if quant_state is None:
+                        W_base = base_weight.to(torch.float32)
+                    else:
+                        # Maintain the original dequantization logic for bitsandbytes.
+                        W_base = cast(
+                            Tensor,
+                            bnb.functional.dequantize_4bit(
+                                base_weight.data, 
+                                quant_state
+                            ).to(torch.float32),
+                        )
+
+                    # Row normalization setup.
+                    # Pre-calculate the original row norms to preserve them.
+                    # This implements the RowNormalization.FULL logic.
+                    W_row_norms = LA.vector_norm(W_base, dim=1, keepdim=True).detach()
+
+                    # Adapter target identification.
+                    # We optimize the LoRA weights A and B.
+                    lora_A = cast(Tensor, module.lora_A["default"].weight)
+                    lora_B = cast(Tensor, module.lora_B["default"].weight)
+
+                    # Data preparation.
+                    # Move I/O tensors to the device of the adapter weights.
+                    good_input, good_output = good_module_io[layer_index][component][module_index]
+                    bad_input, bad_output = bad_module_io[layer_index][component][module_index]
+
+                    good_input = good_input.float().to(lora_A.device)
+                    good_output = good_output.float().to(lora_A.device)
+                    bad_input = bad_input.float().to(lora_A.device)
+                    bad_output = bad_output.float().to(lora_A.device)
+
+                    # The objective function.
+                    def objective(A: Tensor, B: Tensor) -> Tensor:
+                        # Calculate effective weight: W_eff = W_base + B @ A.
+                        W_eff = W_base + (B @ A)
+
+                        # Apply Row Normalization (keep original norms).
+                        if self.settings.row_normalization == RowNormalization.FULL:
+                            # Normalize to unit length, then scale by original norms.
+                            W_eff = F.normalize(W_eff, p=2, dim=1) * W_row_norms
+
+                        # Compute outputs using the effective weight.
+                        new_good_output = good_input @ W_eff.T
+                        new_bad_output = bad_input @ W_eff.T
+
+                        # The original ARA loss function.
+                        preserve_good_behavior = (
+                            (new_good_output - good_output) ** 2
+                        ).mean()
+
+                        steer_bad_behavior = (
+                            mean_distances_to_knn(
+                                new_bad_output,
+                                good_output,
+                                parameters.neighbor_count,
+                            ).mean()
+                            + parameters.overcorrect_relative_weight
+                            * -mean_distances_to_knn(
+                                new_bad_output,
+                                bad_output,
+                                parameters.neighbor_count,
+                            ).mean()
+                        )
+
+                        return (
+                            parameters.preserve_good_behavior_weight
+                            * preserve_good_behavior
+                            + parameters.steer_bad_behavior_weight * steer_bad_behavior
+                        )
+
+                    # Optimization loop.
+                    # We optimize A and B, not the base matrix.
+                    optimizer = LBFGS(
+                        [lora_A, lora_B],
+                        lr=1.0,
+                        max_iter=20,
+                        history_size=10,
+                        line_search_fn="strong_wolfe",
+                    )
+
+                    def closure():
+                        optimizer.zero_grad()
+                        # Pass the actual tensors being optimized to the objective.
+                        loss = objective(lora_A, lora_B)
+                        loss.backward()
+                        return loss
+
+                    # Run optimization steps.
+                    for step in range(5):
+                        optimizer.step(closure)
+
+                    # Free the gradient buffers accumulated on the LoRA adapter
+                    # parameters during optimization (see ara_abliterate for details).
+                    optimizer.zero_grad(set_to_none=True)
+
     def generate(
         self,
         prompts: list[Prompt],
@@ -653,6 +931,132 @@ class Model:
 
         return torch.cat(residuals, dim=0)
 
+    def get_module_io(
+        self,
+        prompts: list[Prompt],
+    ) -> 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: ModuleIO = []
+
+        def get_hook(
+            layer_index: int,
+            component: str,
+            module_index: int,
+        ) -> Callable[[Module, tuple[Tensor, ...], Tensor], None]:
+            def hook(
+                module: Module,
+                inputs: tuple[Tensor, ...],
+                outputs: Tensor,
+            ) -> None:
+                if len(module_io) == layer_index:
+                    # First invocation of the hook for this layer.
+                    module_io.append({})
+
+                # Layers are invoked in order during inference,
+                # so this should always hold.
+                assert len(module_io) == layer_index + 1
+
+                if component not in module_io[layer_index]:
+                    module_io[layer_index][component] = {}
+
+                # Each module should be invoked at most once per inference step.
+                assert module_index not in module_io[layer_index][component]
+
+                # inputs[0] and outputs have shape (prompt, position, component),
+                # so this extracts the input/output at the end of each prompt.
+                # Move to CPU to decouple from device assignments, which can
+                # change between model reloads in multi-GPU configurations.
+                input = inputs[0][:, -1, :].detach().clone().cpu()
+                output = outputs[:, -1, :].detach().clone().cpu()
+
+                # The modules associated with a component (e.g. expert MLPs)
+                # are not necessarily invoked in order, nor are all of them
+                # necessarily invoked in each inference step, so we cannot
+                # use a list here.
+                module_io[layer_index][component][module_index] = (input, output)
+
+            return hook
+
+        hook_handles: list[RemovableHandle] = []
+
+        for layer_index in range(len(self.get_layers())):
+            for component, modules in self.get_layer_modules(layer_index).items():
+                for module_index, module in enumerate(modules):
+                    hook_handles.append(
+                        module.register_forward_hook(
+                            get_hook(layer_index, component, module_index)
+                        )
+                    )
+
+        self.generate(prompts, max_new_tokens=1)
+
+        for hook_handle in hook_handles:
+            hook_handle.remove()
+
+        return module_io
+
+    def get_module_io_batched(
+        self,
+        prompts: list[Prompt],
+    ) -> 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: list[ModuleIO] = [
+            self.get_module_io(batch)
+            for batch in batchify(prompts, self.settings.batch_size)
+        ]
+
+        module_io: ModuleIO = []
+
+        for layer_index in range(len(self.get_layers())):
+            module_io.append({})
+
+            for module_io_batch in module_io_batches:
+                for component, io_map in module_io_batch[layer_index].items():
+                    if component not in module_io[layer_index]:
+                        module_io[layer_index][component] = {}
+
+                    for module_index in io_map:
+                        if module_index not in module_io[layer_index][component]:
+                            # This is a placeholder; the actual aggregation happens below.
+                            # We need to iterate over the batches twice because we don't
+                            # know in advance which components and module indices are present.
+                            module_io[layer_index][component][module_index] = (
+                                torch.empty(0),
+                                torch.empty(0),
+                            )
+
+            for component, io_map in module_io[layer_index].items():
+                for module_index in io_map:
+                    inputs_outputs = [
+                        module_io_batch[layer_index][component][module_index]
+                        for module_io_batch in module_io_batches
+                        if component in module_io_batch[layer_index]
+                        and module_index in module_io_batch[layer_index][component]
+                    ]
+                    input = torch.cat(
+                        [input_output[0] for input_output in inputs_outputs],
+                        dim=0,
+                    )
+                    output = torch.cat(
+                        [input_output[1] for input_output in inputs_outputs],
+                        dim=0,
+                    )
+
+                    # The key already exists, and replacing existing values
+                    # in a dictionary while iterating over the same dictionary
+                    # is safe in Python.
+                    module_io[layer_index][component][module_index] = (input, output)
+
+        return module_io
+
     # We work with logprobs rather than probabilities for numerical stability
     # when computing the KL divergence.
     def get_logprobs(self, prompts: list[Prompt]) -> Tensor:
@@ -719,7 +1123,12 @@ class Model:
             max_new_tokens=4096,
         )  # ty:ignore[call-non-callable]
 
-        return self.tokenizer.decode(
-            outputs[0, inputs["input_ids"].shape[1] :],
-            skip_special_tokens=True,
+        # This cast is valid because str is the return type
+        # when passing a sequence of token IDs.
+        return cast(
+            str,
+            self.tokenizer.decode(
+                outputs[0, inputs["input_ids"].shape[1] :],
+                skip_special_tokens=True,
+            ),
         )

src/heretic/progress.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+from typing import Any
+
+import tqdm
+import tqdm.auto
+from rich.progress import Progress
+
+
+# A class that provides the same interface as tqdm,
+# but displays progress bars using Rich.
+class TqdmShim(tqdm.tqdm):
+    def __init__(self, *args: Any, **kwargs: Any):
+        self.rich_progress = Progress(transient=True)
+        self.rich_progress.start()
+        self.rich_task_id = self.rich_progress.add_task(
+            kwargs.get("desc", ""),
+            total=kwargs.get("total", None),
+        )
+
+        # Chain up to the parent constructor to ensure that the internal state of the superclass
+        # is correctly initialized, which some methods that we don't override might rely on.
+        super().__init__(*args, **kwargs)
+
+    def display(self, *args: Any, **kwargs: Any):
+        self.rich_progress.update(
+            self.rich_task_id,
+            description=self.desc,
+            total=self.total,
+            completed=self.n,
+        )
+
+    def close(self, *args: Any, **kwargs: Any):
+        self.rich_progress.stop()
+
+
+def patch_tqdm():
+    tqdm.tqdm = TqdmShim  # ty:ignore[invalid-assignment]
+    tqdm.auto.tqdm = TqdmShim  # ty:ignore[invalid-assignment]

src/heretic/utils.py

@@ -25,8 +25,9 @@ from optuna import Trial
 from psutil import Process
 from questionary import Choice, Style
 from rich.console import Console
+from torch import Tensor
 
-from .config import DatasetSpecification, Settings
+from .config import DatasetSpecification, RowNormalization, Settings
 
 print = Console(highlight=False).print
 
@@ -38,11 +39,17 @@ def print_memory_usage():
     p("Resident system RAM", Process().memory_info().rss)
 
     if torch.cuda.is_available():
-        p("Allocated GPU VRAM", torch.cuda.memory_allocated())
-        p("Reserved GPU VRAM", torch.cuda.memory_reserved())
+        count = torch.cuda.device_count()
+        allocated = sum(torch.cuda.memory_allocated(device) for device in range(count))
+        reserved = sum(torch.cuda.memory_reserved(device) for device in range(count))
+        p("Allocated GPU VRAM", allocated)
+        p("Reserved GPU VRAM", reserved)
     elif is_xpu_available():
-        p("Allocated XPU memory", torch.xpu.memory_allocated())
-        p("Reserved XPU memory", torch.xpu.memory_reserved())
+        count = torch.xpu.device_count()
+        allocated = sum(torch.xpu.memory_allocated(device) for device in range(count))
+        reserved = sum(torch.xpu.memory_reserved(device) for device in range(count))
+        p("Allocated XPU memory", allocated)
+        p("Reserved XPU memory", reserved)
     elif torch.backends.mps.is_available():
         p("Allocated MPS memory", torch.mps.current_allocated_memory())
         p("Driver (reserved) MPS memory", torch.mps.driver_allocated_memory())
@@ -228,6 +235,14 @@ def batchify(items: list[T], batch_size: int) -> list[list[T]]:
     return [items[i : i + batch_size] for i in range(0, len(items), batch_size)]
 
 
+# For each vector in the 2D-tensor `a`, computes the mean Euclidean distance
+# to the `k` nearest neighbors of the vector among the vectors in the 2D-tensor `b`.
+def mean_distances_to_knn(a: Tensor, b: Tensor, k: int) -> Tensor:
+    distances = torch.cdist(a, b)
+    nearest_distances, _ = distances.topk(k, dim=1, largest=False)
+    return nearest_distances.mean(1)
+
+
 def empty_cache():
     # Collecting garbage is not an idempotent operation, and to avoid OOM errors,
     # gc.collect() has to be called both before and after emptying the backend cache.
@@ -250,19 +265,46 @@ 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:
+        parameters = trial.user_attrs["ara_parameters"]
 
-    direction_index = trial.user_attrs["direction_index"]
-    params["direction_index"] = (
-        "per layer" if (direction_index is None) else f"{direction_index:.2f}"
-    )
+        return {
+            name: (f"{value:.4f}" if isinstance(value, float) else f"{value}")
+            for name, value in parameters.items()
+        }
+    else:
+        params = {}
 
-    for component, parameters in trial.user_attrs["parameters"].items():
-        for name, value in parameters.items():
-            params[f"{component}.{name}"] = f"{value:.2f}"
+        direction_index = trial.user_attrs["direction_index"]
+        params["direction_index"] = (
+            "per layer" if (direction_index is None) else f"{direction_index:.2f}"
+        )
 
-    return params
+        for component, parameters in trial.user_attrs["parameters"].items():
+            for name, value in parameters.items():
+                params[f"{component}.{name}"] = f"{value:.2f}"
+
+        return params
+
+
+def get_method_description(settings: Settings) -> str:
+    if settings.use_ara:
+        return (
+            " with the [Arbitrary-Rank Ablation (ARA)](https://github.com/p-e-w/heretic/pull/211) method"
+            + (
+                " (with row-norm preservation)"
+                if settings.row_normalization == RowNormalization.FULL
+                else ""
+            )
+        )
+    elif (
+        settings.orthogonalize_direction
+        and settings.row_normalization == RowNormalization.FULL
+    ):
+        return " with a variant of the [Magnitude-Preserving Orthogonal Ablation (MPOA)](https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration) method"
+    else:
+        return ""
 
 
 def get_readme_intro(
@@ -271,11 +313,17 @@ def get_readme_intro(
     base_refusals: int,
     bad_prompts: list[Prompt],
 ) -> str:
-    model_link = f"[{settings.model}](https://huggingface.co/{settings.model})"
+    if Path(settings.model).exists():
+        # Hide the path, which may contain private information.
+        model_link = "a model"
+    else:
+        model_link = f"[{settings.model}](https://huggingface.co/{settings.model})"
 
     return f"""# This is a decensored version of {
         model_link
-    }, made using [Heretic](https://github.com/p-e-w/heretic) v{version("heretic-llm")}
+    }, made using [Heretic](https://github.com/p-e-w/heretic) v{version("heretic-llm")}{
+        get_method_description(settings)
+    }
 
 ## Abliteration parameters
 
@@ -285,7 +333,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()
             ]
         )
     }
@@ -294,7 +342,10 @@ def get_readme_intro(
 
 | Metric | This model | Original model ({model_link}) |
 | :----- | :--------: | :---------------------------: |
-| **KL divergence** | {trial.user_attrs["kl_divergence"]:.4f} | 0 *(by definition)* |
+| **{"PIQA acc_norm" if settings.use_piqa else "KL divergence"}** | {
+        (-1 if settings.use_piqa else 1) * trial.user_attrs["kl_divergence"]:.4f} | {
+        "*Unknown*" if settings.use_piqa else "0 *(by definition)*"
+    } |
 | **Refusals** | {trial.user_attrs["refusals"]}/{len(bad_prompts)} | {base_refusals}/{
         len(bad_prompts)
     } |