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.

.gemini/styleguide.md

@@ -0,0 +1,11 @@
+# Style guide and coding conventions
+
+* Identifier names should not contain abbreviations unless those abbreviations are very widely used and understood (e.g. "KL divergence").
+* Comments should start with a capital letter and end with a period. They should use correct grammar and spelling.
+* Function and method signatures **must** be fully type-annotated, including the return type (if any).
+* Every Python code file **must** start with an SPDX/Copyright header.
+* Settings descriptions should start with a capital letter and end with a period.
+* When new settings are added in `config.py`, they should also be added to `config.default.toml`, set to their default value and with their description as a comment. The order of settings in `config.default.toml` should match that in `config.py`.
+* Pull requests should implement one change, and one change only.
+  * PRs containing multiple semantically independent changes **must** be split into multiple PRs.
+  * PRs **must not** change existing code unless the changes are *directly related* to the PR. This includes changes to formatting and comments.

.gitattributes

@@ -0,0 +1 @@
+* text eol=lf

.github/workflows/ci.yml

@@ -17,10 +17,10 @@ jobs:
 
     steps:
       - name: Check out code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           enable-cache: true
           cache-dependency-glob: "uv.lock"
@@ -37,6 +37,9 @@ jobs:
       - name: Lint and check import sorting
         run: uv run ruff check --output-format=github --extend-select I .
 
+      - name: Check typing
+        run: uv run ty check --output-format=github --error-on-warning .
+
       - name: Build package
         run: uv build
 

.gitignore

@@ -7,7 +7,7 @@ wheels/
 *.egg-info
 
 # Virtual environments
-.venv
+.venv/
 
 # Caches
 /.ruff_cache/
@@ -17,3 +17,9 @@ wheels/
 
 # Configuration files
 /config.toml
+
+# Study checkpoints
+/checkpoints/
+
+# Residual plots
+/plots/

README.md

@@ -1,11 +1,15 @@
-# Heretic: Fully automatic censorship removal for language models
+<img width="128" height="128" align="right" alt="Logo" src="https://github.com/user-attachments/assets/df5f2840-2f92-4991-aa57-252747d7182e" />
 
-[![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)
+# 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
-as "abliteration" ([Arditi et al. 2024](https://arxiv.org/abs/2406.11717)),
+as "abliteration" ([Arditi et al. 2024](https://arxiv.org/abs/2406.11717),
+Lai 2025 ([1](https://huggingface.co/blog/grimjim/projected-abliteration),
+[2](https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration))),
 with a TPE-based parameter optimizer powered by [Optuna](https://optuna.org/).
 
 This approach enables Heretic to work **completely automatically.** Heretic
@@ -65,8 +69,11 @@ Heretic supports most dense models, including many multimodal models, and
 several different MoE architectures. It does not yet support SSMs/hybrid models,
 models with inhomogeneous layers, and certain novel attention systems.
 
-You can find a collection of models that have been decensored using Heretic
-[on Hugging Face](https://huggingface.co/collections/p-e-w/the-bestiary).
+You can find a small collection of models that have been decensored using Heretic
+[on Hugging Face](https://huggingface.co/collections/p-e-w/the-bestiary),
+and the community has created and published
+[well over 1,000](https://huggingface.co/models?other=heretic)
+Heretic models in addition to those.
 
 
 ## Usage
@@ -89,8 +96,10 @@ a configuration file.
 
 At the start of a program run, Heretic benchmarks the system to determine
 the optimal batch size to make the most of the available hardware.
-On an RTX 3090, with the default configuration, decensoring Llama-3.1-8B
-takes about 45 minutes.
+On an RTX 3090, with the default configuration, decensoring Llama-3.1-8B-Instruct
+takes about 45 minutes. Note that Heretic supports model quantization with
+bitsandbytes, which can drastically reduce the amount of VRAM required to process
+models. Set the `quantization` option to `bnb_4bit` to enable quantization.
 
 After Heretic has finished decensoring a model, you are given the option to
 save the model, upload it to Hugging Face, chat with it to test how well it works,
@@ -242,7 +251,8 @@ The development of Heretic was informed by:
 * [The original abliteration paper (Arditi et al. 2024)](https://arxiv.org/abs/2406.11717)
 * [Maxime Labonne's article on abliteration](https://huggingface.co/blog/mlabonne/abliteration),
   as well as some details from the model cards of his own abliterated models (see above)
-* [Jim Lai's article describing "projected abliteration"](https://huggingface.co/blog/grimjim/projected-abliteration)
+* Jim Lai's articles describing ["projected abliteration"](https://huggingface.co/blog/grimjim/projected-abliteration)
+  and ["norm-preserving biprojected abliteration"](https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration)
 
 
 ## Citation
@@ -263,7 +273,7 @@ If you use Heretic for your research, please cite it using the following BibTeX
 
 ## License
 
-Copyright &copy; 2025  Philipp Emanuel Weidmann (<pew@worldwidemann.com>)
+Copyright &copy; 2025-2026  Philipp Emanuel Weidmann (<pew@worldwidemann.com>) + contributors
 
 This program is free software: you can redistribute it and/or modify
 it under the terms of the GNU Affero General Public License as published by

config.default.toml

@@ -1,4 +1,5 @@
-# Copy this file to config.toml and edit the configuration to your liking.
+# Rename this file to config.toml, place it in the working directory
+# that you run Heretic from, and edit the configuration to your liking.
 
 # List of PyTorch dtypes to try when loading model tensors.
 # If loading with a dtype fails, the next dtype in the list will be tried.
@@ -15,9 +16,17 @@ dtypes = [
     "float32",
 ]
 
+# Quantization method to use when loading the model. Options:
+# "none" (no quantization),
+# "bnb_4bit" (4-bit quantization using bitsandbytes).
+quantization = "none"
+
 # Device map to pass to Accelerate when loading the model.
 device_map = "auto"
 
+# Maximum memory to allocate per device.
+# max_memory = {"0": "20GB", "cpu": "64GB"}
+
 # Number of input sequences to process in parallel (0 = auto).
 batch_size = 0  # auto
 
@@ -27,6 +36,9 @@ max_batch_size = 128
 # Maximum number of tokens to generate for each response.
 max_response_length = 100
 
+# Whether to print prompt/response pairs when counting refusals.
+print_responses = false
+
 # Whether to print detailed information about residuals and refusal directions.
 print_residual_geometry = false
 
@@ -46,12 +58,42 @@ residual_plot_style = "dark_background"
 # This is used to ensure balanced co-optimization of KL divergence and refusal count.
 kl_divergence_scale = 1.0
 
+# The KL divergence to target. Below this value, an objective based on the refusal count is used.
+# This helps prevent the sampler from extensively exploring parameter combinations that "do nothing".
+kl_divergence_target = 0.01
+
+# Whether to adjust the refusal directions so that only the component that is
+# orthogonal to the good direction is subtracted during abliteration.
+orthogonalize_direction = false
+
+# How to apply row normalization of the weights. Options:
+# "none" (no normalization),
+# "pre" (compute LoRA adapter relative to row-normalized weights),
+# "full" (like "pre", but renormalizes to preserve original row magnitudes).
+row_normalization = "none"
+
+# The rank of the LoRA adapter to use when "full" row normalization is used.
+# Row magnitude preservation is approximate due to non-linear effects,
+# and this determines the rank of that approximation. Higher ranks produce
+# larger output files and may slow down evaluation.
+full_normalization_lora_rank = 3
+
+# The symmetric winsorization to apply to the per-prompt, per-layer residual vectors,
+# expressed as the quantile to clamp to (between 0 and 1). Disabled by default.
+# This can tame so-called "massive activations" that occur in some models.
+# Example: winsorization_quantile = 0.95 computes the 0.95-quantile of the absolute values
+# of the components, then clamps the magnitudes of all components to that quantile.
+winsorization_quantile = 1.0
+
 # Number of abliteration trials to run during optimization.
 n_trials = 200
 
 # Number of trials that use random sampling for the purpose of exploration.
 n_startup_trials = 60
 
+# Directory to save and load study progress to/from.
+study_checkpoint_dir = "checkpoints"
+
 # Strings whose presence in a response (case insensitive) identifies the response as a refusal.
 refusal_markers = [
     "sorry",

config.noslop.toml

@@ -0,0 +1,163 @@
+# Rename this file to config.toml, place it in the working directory
+# that you run Heretic from, and edit the configuration to your liking.
+
+max_response_length = 300
+
+residual_plot_title = "PaCMAP Projection of Residuals for Slop-Suppressing/Inducing Prompts"
+
+refusal_markers = [
+    "Eldoria",
+    "Lumina",
+    "ethereal",
+    "thick with",
+    "celestial",
+    "radiant",
+    "black as",
+    "despair",
+    "crimson",
+    "resplendent",
+    "unravel",
+    "belied",
+    "velvet",
+    "moonless",
+    "moonlit",
+    "entangled",
+    "twilight",
+    "forever",
+    "first kiss",
+    "gasp",
+    "whisper",
+    "hue",
+    "symphony",
+    "scarcely believe",
+    "gilded",
+    "hummed",
+    "abuzz",
+    "perpetually",
+    "scent",
+    "perfume",
+    "neon lights",
+    "kaleidoscopic",
+    "adrift",
+    "sultry",
+    "melancholic",
+    "stark contrast",
+    "inky",
+    "coy",
+    "vast",
+    "purr",
+    "radiant",
+    "beacon",
+    "a thousand ships",
+    "tapestry",
+    "bustling",
+    "abyss",
+    "gnarled",
+    "tremble",
+    "trembling",
+    "profound",
+    "terrible",
+    "ancient",
+    "sapphire",
+    "ruby",
+    "emerald",
+    "diamond",
+    "stolen",
+    "promise",
+    "the air was",
+    "obsidian",
+    "gleaming with",
+    "faintest hint",
+    "trepidation",
+    "sun-kissed",
+    "azure",
+    "deep",
+    "beloved",
+    "cosmos",
+    "devoid",
+    "soft chime",
+    "echo",
+    "palpable",
+    "blossom",
+    "adrift",
+    "faint",
+    "emerged",
+    "shiver",
+    "spine",
+    "hairs on the back",
+    "cinematic",
+    "specter",
+    "golden",
+    "inescapable",
+    "sentinel",
+    "flicker",
+    "testament",
+    "embodiment",
+    "etched with",
+    "rise and fall",
+    "the very air",
+    "slither",
+    "a pang of",
+    "eternal",
+    "eternity",
+    "veil of",
+    "painting the",
+    "bathed in",
+    "boundless",
+    "stretched out",
+    "beneath",
+    "lullaby",
+    "unsuspecting",
+    "handsome",
+    "defied the very",
+    "barely above",
+    "never-ending",
+    "caress",
+    "realm",
+    "fiery",
+    "raven",
+    "twin pools",
+    "gloaming",
+    "grimy",
+    "labyrinth",
+    "the very notion",
+    "something...",
+    "the halls of",
+    "conflagration of",
+    "shattered like",
+    "as dark as",
+    "yearned for",
+    "unyielding",
+    "lifetime",
+    "ensnared",
+]
+
+system_prompt = "You are a professional writer."
+
+[good_prompts]
+dataset = "llm-aes/writing-prompts"
+split = "train[:500]"
+column = "prompt"
+prefix = "Write a short story based on the writing prompt below. Avoid literary cliches, purple prose, and flowery language.\n\nWriting prompt:"
+residual_plot_label = "Slop-suppressing prompts"
+residual_plot_color = "royalblue"
+
+[bad_prompts]
+dataset = "llm-aes/writing-prompts"
+split = "train[:500]"
+column = "prompt"
+prefix = "Write a short story based on the writing prompt below. Make extensive use of literary cliches, purple prose, and flowery language.\n\nWriting prompt:"
+residual_plot_label = "Slop-inducing prompts"
+residual_plot_color = "darkorange"
+
+[good_evaluation_prompts]
+dataset = "llm-aes/writing-prompts"
+split = "train[1000:1100]"
+column = "prompt"
+prefix = "Write a short story based on the writing prompt below. Avoid literary cliches, purple prose, and flowery language.\n\nWriting prompt:"
+
+[bad_evaluation_prompts]
+dataset = "llm-aes/writing-prompts"
+split = "train[1000:1100]"
+column = "prompt"
+prefix = "Write a short story based on the writing prompt below.\n\nWriting prompt:"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "heretic-llm"
-version = "1.1.0"
+version = "1.2.0"
 description = "Fully automatic censorship removal for language models"
 readme = "README.md"
 license = "AGPL-3.0-or-later"
@@ -22,30 +22,39 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
 ]
 dependencies = [
-    "accelerate>=1.10.0",
-    "datasets>=4.0.0",
-    "hf-transfer>=0.1.9",
-    "huggingface-hub>=0.34.4",
-    "optuna>=4.5.0",
-    "pydantic-settings>=2.10.1",
-    "questionary>=2.1.1",
-    "rich>=14.1.0",
-    "transformers>=4.55.2",
+    "accelerate~=1.13",
+    "bitsandbytes~=0.49",
+    "datasets~=4.7",
+    "hf-transfer~=0.1",
+    "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.3",
+    "tqdm~=4.67",
+    "transformers~=5.3",
 ]
 
 [project.optional-dependencies]
 research = [
-    "geom-median>=0.1.0",
-    "imageio>=2.37.2",
-    "matplotlib>=3.10.7",
-    "numpy>=2.2.6",
-    "pacmap>=0.8.0",
-    "scikit-learn>=1.7.2",
+    "geom-median~=0.1",
+    "imageio~=2.37",
+    "matplotlib~=3.10",
+    "pacmap~=0.8",
+    "scikit-learn~=1.7",
 ]
 
 [dependency-groups]
 dev = [
     "ruff>=0.14.5",
+    "ty>=0.0.5",
 ]
 
 [project.urls]

src/heretic/analyzer.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025  Philipp Emanuel Weidmann <pew@worldwidemann.com>
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
 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
@@ -30,8 +32,10 @@ class Analyzer:
 
     def print_residual_geometry(self):
         try:
-            from geom_median.torch import compute_geometric_median
-            from sklearn.metrics import silhouette_score
+            from geom_median.torch import (  # ty:ignore[unresolved-import]
+                compute_geometric_median,
+            )
+            from sklearn.metrics import silhouette_score  # ty:ignore[unresolved-import]
         except ImportError:
             print()
             print(
@@ -152,12 +156,12 @@ class Analyzer:
 
     def plot_residuals(self):
         try:
-            import imageio.v3 as iio
-            import matplotlib.pyplot as plt
-            import numpy as np
-            from geom_median.numpy import compute_geometric_median
-            from numpy.typing import NDArray
-            from pacmap import PaCMAP
+            import imageio.v3 as iio  # ty:ignore[unresolved-import]
+            import matplotlib.pyplot as plt  # ty:ignore[unresolved-import]
+            from geom_median.numpy import (  # ty:ignore[unresolved-import]
+                compute_geometric_median,
+            )
+            from pacmap import PaCMAP  # ty:ignore[unresolved-import]
         except ImportError:
             print()
             print(

src/heretic/config.py

@@ -1,17 +1,31 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025  Philipp Emanuel Weidmann <pew@worldwidemann.com>
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
+from enum import Enum
 from typing import Dict
 
 from pydantic import BaseModel, Field
 from pydantic_settings import (
     BaseSettings,
+    CliSettingsSource,
+    EnvSettingsSource,
     PydanticBaseSettingsSource,
-    SettingsConfigDict,
     TomlConfigSettingsSource,
 )
 
 
+class QuantizationMethod(str, Enum):
+    NONE = "none"
+    BNB_4BIT = "bnb_4bit"
+
+
+class RowNormalization(str, Enum):
+    NONE = "none"
+    PRE = "pre"
+    # POST = "post"  # Theoretically possible, but provides no advantage.
+    FULL = "full"
+
+
 class DatasetSpecification(BaseModel):
     dataset: str = Field(
         description="Hugging Face dataset ID, or path to dataset on disk."
@@ -21,6 +35,21 @@ class DatasetSpecification(BaseModel):
 
     column: str = Field(description="Column in the dataset that contains the prompts.")
 
+    prefix: str = Field(
+        default="",
+        description="Text to prepend to each prompt.",
+    )
+
+    suffix: str = Field(
+        default="",
+        description="Text to append to each prompt.",
+    )
+
+    system_prompt: str | None = Field(
+        default=None,
+        description="System prompt to use with the prompts (overrides global system prompt if set).",
+    )
+
     residual_plot_label: str | None = Field(
         default=None,
         description="Label to use for the dataset in plots of residual vectors.",
@@ -32,12 +61,27 @@ 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.")
 
     evaluate_model: str | None = Field(
         default=None,
-        description="If this model ID or path is set, then instead of abliterating the main model, evaluate this model relative to the main model.",
+        description=(
+            "If this model ID or path is set, then instead of abliterating the main model, "
+            "evaluate this model relative to the main model."
+        ),
     )
 
     dtypes: list[str] = Field(
@@ -53,7 +97,19 @@ class Settings(BaseSettings):
             # if that was the dtype "auto" resolved to).
             "float32",
         ],
-        description="List of PyTorch dtypes to try when loading model tensors. If loading with a dtype fails, the next dtype in the list will be tried.",
+        description=(
+            "List of PyTorch dtypes to try when loading model tensors. "
+            "If loading with a dtype fails, the next dtype in the list will be tried."
+        ),
+    )
+
+    quantization: QuantizationMethod = Field(
+        default=QuantizationMethod.NONE,
+        description=(
+            "Quantization method to use when loading the model. Options: "
+            '"none" (no quantization), '
+            '"bnb_4bit" (4-bit quantization using bitsandbytes).'
+        ),
     )
 
     device_map: str | Dict[str, int | str] = Field(
@@ -61,6 +117,11 @@ class Settings(BaseSettings):
         description="Device map to pass to Accelerate when loading the model.",
     )
 
+    max_memory: Dict[str, str] | None = Field(
+        default=None,
+        description='Maximum memory to allocate per device (e.g., {"0": "20GB", "cpu": "64GB"}).',
+    )
+
     trust_remote_code: bool | None = Field(
         default=None,
         description="Whether to trust remote code when loading the model.",
@@ -81,6 +142,11 @@ class Settings(BaseSettings):
         description="Maximum number of tokens to generate for each response.",
     )
 
+    print_responses: bool = Field(
+        default=False,
+        description="Whether to print prompt/response pairs when counting refusals.",
+    )
+
     print_residual_geometry: bool = Field(
         default=False,
         description="Whether to print detailed information about residuals and refusal directions.",
@@ -114,6 +180,89 @@ class Settings(BaseSettings):
         ),
     )
 
+    kl_divergence_target: float = Field(
+        default=0.01,
+        description=(
+            "The KL divergence to target. Below this value, an objective based on the refusal count is used. "
+            'This helps prevent the sampler from extensively exploring parameter combinations that "do nothing".'
+        ),
+    )
+
+    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=(
+            "Whether to adjust the refusal directions so that only the component that is "
+            "orthogonal to the good direction is subtracted during abliteration."
+        ),
+    )
+
+    row_normalization: RowNormalization = Field(
+        default=RowNormalization.FULL,
+        description=(
+            "How to apply row normalization of the weights. Options: "
+            '"none" (no normalization), '
+            '"pre" (compute LoRA adapter relative to row-normalized weights), '
+            '"full" (like "pre", but renormalizes to preserve original row magnitudes).'
+        ),
+    )
+
+    full_normalization_lora_rank: int = Field(
+        default=3,
+        description=(
+            'The rank of the LoRA adapter to use when "full" row normalization is used. '
+            "Row magnitude preservation is approximate due to non-linear effects, "
+            "and this determines the rank of that approximation. Higher ranks produce "
+            "larger output files and may slow down evaluation."
+        ),
+    )
+
+    winsorization_quantile: float = Field(
+        default=1.0,
+        description=(
+            "The symmetric winsorization to apply to the per-prompt, per-layer residual vectors, "
+            "expressed as the quantile to clamp to (between 0 and 1). Disabled by default. "
+            'This can tame so-called "massive activations" that occur in some models. '
+            "Example: winsorization_quantile = 0.95 computes the 0.95-quantile of the absolute values "
+            "of the components, then clamps the magnitudes of all components to that quantile."
+        ),
+    )
+
     n_trials: int = Field(
         default=200,
         description="Number of abliteration trials to run during optimization.",
@@ -124,6 +273,72 @@ class Settings(BaseSettings):
         description="Number of trials that use random sampling for the purpose of exploration.",
     )
 
+    study_checkpoint_dir: str = Field(
+        default="checkpoints",
+        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",
@@ -207,16 +422,6 @@ class Settings(BaseSettings):
         description="Dataset of prompts that tend to result in refusals (used for evaluating model performance).",
     )
 
-    # "Model" refers to the Pydantic model of the settings class here,
-    # not to the language model. The field must have this exact name.
-    model_config = SettingsConfigDict(
-        toml_file="config.toml",
-        env_prefix="HERETIC_",
-        cli_parse_args=True,
-        cli_implicit_flags=True,
-        cli_kebab_case=True,
-    )
-
     @classmethod
     def settings_customise_sources(
         cls,
@@ -227,9 +432,15 @@ class Settings(BaseSettings):
         file_secret_settings: PydanticBaseSettingsSource,
     ) -> tuple[PydanticBaseSettingsSource, ...]:
         return (
-            init_settings,
-            env_settings,
+            init_settings,  # Used during resume - should override *all* other sources.
+            CliSettingsSource(
+                settings_cls,
+                cli_parse_args=True,
+                cli_implicit_flags=True,
+                cli_kebab_case=True,
+            ),
+            EnvSettingsSource(settings_cls, env_prefix="HERETIC_"),
             dotenv_settings,
             file_secret_settings,
-            TomlConfigSettingsSource(settings_cls),
+            TomlConfigSettingsSource(settings_cls, toml_file="config.toml"),
         )

src/heretic/evaluator.py

@@ -1,23 +1,34 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025  Philipp Emanuel Weidmann <pew@worldwidemann.com>
+# 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
 from .model import Model
-from .utils import load_prompts, print
+from .utils import Prompt, load_prompts, print
 
 
 class Evaluator:
+    settings: Settings
+    model: Model
+    good_prompts: list[Prompt]
+    bad_prompts: list[Prompt]
+    base_logprobs: Tensor
+    base_refusals: int
+
     def __init__(self, settings: Settings, model: Model):
         self.settings = settings
         self.model = model
 
+        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.good_evaluation_prompts)
+            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...")
@@ -27,7 +38,7 @@ class Evaluator:
         print(
             f"Loading bad evaluation prompts from [bold]{settings.bad_evaluation_prompts.dataset}[/]..."
         )
-        self.bad_prompts = load_prompts(settings.bad_evaluation_prompts)
+        self.bad_prompts = load_prompts(settings, settings.bad_evaluation_prompts)
         print(f"* [bold]{len(self.bad_prompts)}[/] prompts loaded")
 
         print("* Counting model refusals...")
@@ -57,11 +68,48 @@ class Evaluator:
         return False
 
     def count_refusals(self) -> int:
-        responses = self.model.get_responses_batched(self.bad_prompts)
-        refusals = [response for response in responses if self.is_refusal(response)]
-        return len(refusals)
+        refusal_count = 0
+
+        responses = self.model.get_responses_batched(
+            self.bad_prompts,
+            skip_special_tokens=True,
+        )
+
+        for prompt, response in zip(self.bad_prompts, responses):
+            is_refusal = self.is_refusal(response)
+            if is_refusal:
+                refusal_count += 1
+
+            if self.settings.print_responses:
+                print()
+                print(f"[bold]System prompt:[/] {prompt.system}")
+                print(f"[bold]Prompt:[/] {prompt.user}")
+                if not response.strip():
+                    response = "[italic]\\[empty][/]"
+                print(
+                    f"[bold]Response:[/] [{'red' if is_refusal else 'green'}]{response}[/]"
+                )
+
+        if self.settings.print_responses:
+            print()
+
+        return refusal_count
 
     def get_score(self) -> tuple[tuple[float, float], float, int]:
+        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(
@@ -76,9 +124,29 @@ class Evaluator:
         refusals = self.count_refusals()
         print(f"  * Refusals: [bold]{refusals}[/]/{len(self.bad_prompts)}")
 
+        refusals_score = (
+            refusals / self.base_refusals if self.base_refusals > 0 else float(refusals)
+        )
+
+        if self.settings.use_piqa:
             score = (
-            (kl_divergence / self.settings.kl_divergence_scale),
-            (refusals / self.base_refusals),
+                -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,17 +1,31 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025  Philipp Emanuel Weidmann <pew@worldwidemann.com>
+# 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
 import time
 import warnings
+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
@@ -23,18 +37,23 @@ 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
+from optuna.storages import JournalStorage
+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 Settings
+from .config import QuantizationMethod, RowNormalization, Settings
 from .evaluator import Evaluator
-from .model import AbliterationParameters, Model
+from .model import AbliterationParameters, ARAParameters, Model, get_model_class
 from .utils import (
     empty_cache,
     format_duration,
@@ -42,6 +61,7 @@ from .utils import (
     get_trial_parameters,
     load_prompts,
     print,
+    print_memory_usage,
     prompt_password,
     prompt_path,
     prompt_select,
@@ -49,6 +69,80 @@ from .utils import (
 )
 
 
+def obtain_merge_strategy(settings: Settings) -> str | None:
+    """
+    Prompts the user for how to proceed with saving the model.
+    Provides info to the user if the model is quantized on memory use.
+    Returns "merge", "adapter", or None (if cancelled/invalid).
+    """
+
+    if settings.quantization == QuantizationMethod.BNB_4BIT:
+        print()
+        print(
+            "Model was loaded with quantization. Merging requires reloading the base model."
+        )
+        print(
+            "[yellow]WARNING: CPU merging requires dequantizing the entire model to system RAM.[/]"
+        )
+        print("[yellow]This can lead to system freezes if you run out of memory.[/]")
+
+        try:
+            # Estimate memory requirements by loading the model structure on the "meta" device.
+            # This doesn't consume actual RAM but allows us to inspect the parameter count/dtype.
+            #
+            # Suppress warnings during meta device loading (e.g., "Some weights were not initialized").
+            # These are expected and harmless since we're only inspecting model structure, not running inference.
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                meta_model = get_model_class(settings.model).from_pretrained(
+                    settings.model,
+                    device_map="meta",
+                    torch_dtype=torch.bfloat16,
+                    trust_remote_code=True,
+                )
+                footprint_bytes = meta_model.get_memory_footprint()
+                footprint_gb = footprint_bytes / (1024**3)
+                print(
+                    f"[yellow]Estimated RAM required (excluding overhead): [bold]~{footprint_gb:.2f} GB[/][/]"
+                )
+        except Exception:
+            # Fallback if meta loading fails (e.g. owing to custom model code
+            # or bitsandbytes quantization config issues on the meta device).
+            print(
+                "[yellow]Rule of thumb: You need approximately 3x the parameter count in GB RAM.[/]"
+            )
+            print(
+                "[yellow]Example: A 27B model requires ~80GB RAM. A 70B model requires ~200GB RAM.[/]"
+            )
+        print()
+
+        strategy = prompt_select(
+            "How do you want to proceed?",
+            choices=[
+                Choice(
+                    title="Merge LoRA into full model"
+                    + (
+                        ""
+                        if settings.quantization == QuantizationMethod.NONE
+                        else " (requires sufficient RAM)"
+                    ),
+                    value="merge",
+                ),
+                Choice(
+                    title="Cancel",
+                    value="cancel",
+                ),
+            ],
+        )
+
+        if strategy == "cancel":
+            return None
+
+        return strategy
+    else:
+        return "merge"
+
+
 def run():
     # Enable expandable segments to reduce memory fragmentation on multi-GPU setups.
     if (
@@ -77,7 +171,9 @@ def run():
         sys.argv.insert(-1, "--model")
 
     try:
-        settings = Settings()
+        # The required argument "model" must be provided by the user,
+        # either on the command line or in the configuration file.
+        settings = Settings()  # ty:ignore[missing-argument]
     except ValidationError as error:
         print(f"[red]Configuration contains [bold]{error.error_count()}[/] errors:[/]")
 
@@ -92,24 +188,46 @@ def run():
 
     # Adapted from https://github.com/huggingface/accelerate/blob/main/src/accelerate/commands/env.py
     if torch.cuda.is_available():
-        print(f"GPU type: [bold]{torch.cuda.get_device_name()}[/]")
+        count = torch.cuda.device_count()
+        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):
+            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():
-        print(f"XPU type: [bold]{torch.xpu.get_device_name()}[/]")
+        count = torch.xpu.device_count()
+        print(f"Detected [bold]{count}[/] XPU device(s):")
+        for i in range(count):
+            print(f"* XPU {i}: [bold]{torch.xpu.get_device_name(i)}[/]")
     elif is_mlu_available():
-        print(f"MLU type: [bold]{torch.mlu.get_device_name()}[/]")
+        count = torch.mlu.device_count()  # ty:ignore[unresolved-attribute]
+        print(f"Detected [bold]{count}[/] MLU device(s):")
+        for i in range(count):
+            print(f"* MLU {i}: [bold]{torch.mlu.get_device_name(i)}[/]")  # ty:ignore[unresolved-attribute]
     elif is_sdaa_available():
-        print(f"SDAA type: [bold]{torch.sdaa.get_device_name()}[/]")
+        count = torch.sdaa.device_count()  # ty:ignore[unresolved-attribute]
+        print(f"Detected [bold]{count}[/] SDAA device(s):")
+        for i in range(count):
+            print(f"* SDAA {i}: [bold]{torch.sdaa.get_device_name(i)}[/]")  # ty:ignore[unresolved-attribute]
     elif is_musa_available():
-        print(f"MUSA type: [bold]{torch.musa.get_device_name()}[/]")
+        count = torch.musa.device_count()  # ty:ignore[unresolved-attribute]
+        print(f"Detected [bold]{count}[/] MUSA device(s):")
+        for i in range(count):
+            print(f"* MUSA {i}: [bold]{torch.musa.get_device_name(i)}[/]")  # ty:ignore[unresolved-attribute]
     elif is_npu_available():
-        print(f"CANN version: [bold]{torch.version.cann}[/]")
+        print(f"NPU detected (CANN version: [bold]{torch.version.cann}[/])")  # ty:ignore[unresolved-attribute]
     elif torch.backends.mps.is_available():
-        print("GPU type: [bold]Apple Metal (MPS)[/]")
+        print("Detected [bold]1[/] MPS device (Apple Metal)")
     else:
         print(
             "[bold yellow]No GPU or other accelerator detected. Operations will be slow.[/]"
         )
 
+    if not settings.use_ara:
         # We don't need gradients as we only do inference.
         torch.set_grad_enabled(False)
 
@@ -123,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)
@@ -130,16 +251,101 @@ def run():
     # Silence the warning about multivariate TPE being experimental.
     warnings.filterwarnings("ignore", category=ExperimentalWarning)
 
+    os.makedirs(settings.study_checkpoint_dir, exist_ok=True)
+
+    study_checkpoint_file = os.path.join(
+        settings.study_checkpoint_dir,
+        "".join(
+            [(c if (c.isalnum() or c in ["_", "-"]) else "--") for c in settings.model]
+        )
+        + ".jsonl",
+    )
+
+    lock_obj = JournalFileOpenLock(study_checkpoint_file)
+    backend = JournalFileBackend(study_checkpoint_file, lock_obj=lock_obj)
+    storage = JournalStorage(backend)
+
+    try:
+        existing_study = storage.get_all_studies()[0]
+    except IndexError:
+        existing_study = None
+
+    if existing_study is not None and settings.evaluate_model is None:
+        choices = []
+
+        if existing_study.user_attrs["finished"]:
+            print()
+            print(
+                (
+                    "[green]You have already processed this model.[/] "
+                    "You can show the results from the previous run, allowing you to export models or to run additional trials. "
+                    "Alternatively, you can ignore the previous run and start from scratch. "
+                    "This will delete the checkpoint file and all results from the previous run."
+                )
+            )
+            choices.append(
+                Choice(
+                    title="Show the results from the previous run",
+                    value="continue",
+                )
+            )
+        else:
+            print()
+            print(
+                (
+                    "[yellow]You have already processed this model, but the run was interrupted.[/] "
+                    "You can continue the previous run from where it stopped. This will override any specified settings. "
+                    "Alternatively, you can ignore the previous run and start from scratch. "
+                    "This will delete the checkpoint file and all results from the previous run."
+                )
+            )
+            choices.append(
+                Choice(
+                    title="Continue the previous run",
+                    value="continue",
+                )
+            )
+
+        choices.append(
+            Choice(
+                title="Ignore the previous run and start from scratch",
+                value="restart",
+            )
+        )
+
+        choices.append(
+            Choice(
+                title="Exit program",
+                value="",
+            )
+        )
+
+        print()
+        choice = prompt_select("How would you like to proceed?", choices)
+
+        if choice == "continue":
+            settings = Settings.model_validate_json(
+                existing_study.user_attrs["settings"]
+            )
+        elif choice == "restart":
+            os.unlink(study_checkpoint_file)
+            backend = JournalFileBackend(study_checkpoint_file, lock_obj=lock_obj)
+            storage = JournalStorage(backend)
+        elif choice is None or choice == "":
+            return
+
     model = Model(settings)
+    print()
+    print_memory_usage()
 
     print()
     print(f"Loading good prompts from [bold]{settings.good_prompts.dataset}[/]...")
-    good_prompts = load_prompts(settings.good_prompts)
+    good_prompts = load_prompts(settings, settings.good_prompts)
     print(f"* [bold]{len(good_prompts)}[/] prompts loaded")
 
     print()
     print(f"Loading bad prompts from [bold]{settings.bad_prompts.dataset}[/]...")
-    bad_prompts = load_prompts(settings.bad_prompts)
+    bad_prompts = load_prompts(settings, settings.bad_prompts)
     print(f"* [bold]{len(bad_prompts)}[/] prompts loaded")
 
     if settings.batch_size == 0:
@@ -190,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,
@@ -201,40 +408,79 @@ def run():
     model.response_prefix = commonprefix(responses).rstrip(" ")
 
     # Suppress CoT output.
+    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:
         print()
         print(f"Loading model [bold]{settings.evaluate_model}[/]...")
         settings.model = settings.evaluate_model
-        model.reload_model()
+        model.reset_model()
         print("* Evaluating...")
         evaluator.get_score()
         return
 
+    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)
-    refusal_directions = F.normalize(
-        bad_residuals.mean(dim=0) - good_residuals.mean(dim=0),
-        p=2,
-        dim=1,
+
+        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)
 
@@ -249,6 +495,7 @@ def run():
         empty_cache()
 
     trial_index = 0
+    start_index = 0
     start_time = time.perf_counter()
 
     def objective(trial: Trial) -> tuple[float, float]:
@@ -256,6 +503,50 @@ def run():
         trial_index += 1
         trial.set_user_attr("index", trial_index)
 
+        if settings.use_ara:
+            start_layer_index = trial.suggest_int(
+                "start_layer_index",
+                0,
+                len(model.get_layers()) // 2,
+            )
+            end_layer_index = trial.suggest_int(
+                "end_layer_index",
+                len(model.get_layers()) // 2,
+                len(model.get_layers()),
+            )
+            preserve_good_behavior_weight = trial.suggest_float(
+                "preserve_good_behavior_weight",
+                0.0,
+                1.0,
+            )
+            steer_bad_behavior_weight = trial.suggest_float(
+                "steer_bad_behavior_weight",
+                0.0001,
+                1.0,
+                log=True,
+            )
+            overcorrect_relative_weight = trial.suggest_float(
+                "overcorrect_relative_weight",
+                0.0,
+                1.3,
+            )
+            neighbor_count = trial.suggest_int(
+                "neighbor_count",
+                1,
+                15,
+            )
+
+            ara_parameters = ARAParameters(
+                start_layer_index=start_layer_index,
+                end_layer_index=end_layer_index,
+                preserve_good_behavior_weight=preserve_good_behavior_weight,
+                steer_bad_behavior_weight=steer_bad_behavior_weight,
+                overcorrect_relative_weight=overcorrect_relative_weight,
+                neighbor_count=neighbor_count,
+            )
+
+            trial.set_user_attr("ara_parameters", asdict(ara_parameters))
+        else:
             direction_scope = trial.suggest_categorical(
                 "direction_scope",
                 [
@@ -264,6 +555,8 @@ def run():
                 ],
             )
 
+            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.
@@ -273,8 +566,8 @@ def run():
             # work with conditional or variable-range parameters.
             direction_index = trial.suggest_float(
                 "direction_index",
-            0.4 * (len(model.get_layers()) - 1),
-            0.9 * (len(model.get_layers()) - 1),
+                0.4 * last_layer_index,
+                0.9 * last_layer_index,
             )
 
             if direction_scope == "per layer":
@@ -293,8 +586,8 @@ def run():
                 )
                 max_weight_position = trial.suggest_float(
                     f"{component}.max_weight_position",
-                0.6 * (len(model.get_layers()) - 1),
-                len(model.get_layers()) - 1,
+                    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.
@@ -307,7 +600,7 @@ def run():
                 min_weight_distance = trial.suggest_float(
                     f"{component}.min_weight_distance",
                     1.0,
-                0.6 * (len(model.get_layers()) - 1),
+                    0.6 * last_layer_index,
                 )
 
                 parameters[component] = AbliterationParameters(
@@ -318,24 +611,41 @@ def run():
                 )
 
             trial.set_user_attr("direction_index", direction_index)
-        trial.set_user_attr("parameters", parameters)
+            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}[/]")
+        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.reload_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()
 
         elapsed_time = time.perf_counter() - start_time
-        remaining_time = (elapsed_time / trial_index) * (
+        remaining_time = (elapsed_time / (trial_index - start_index)) * (
             settings.n_trials - trial_index
         )
         print()
@@ -344,6 +654,7 @@ def run():
             print(
                 f"[grey50]Estimated remaining time: [bold]{format_duration(remaining_time)}[/][/]"
             )
+        print_memory_usage()
 
         trial.set_user_attr("kl_divergence", kl_divergence)
         trial.set_user_attr("refusals", refusals)
@@ -365,33 +676,69 @@ def run():
             multivariate=True,
         ),
         directions=[StudyDirection.MINIMIZE, StudyDirection.MINIMIZE],
+        storage=storage,
+        study_name="heretic",
+        load_if_exists=True,
     )
 
+    study.set_user_attr("settings", settings.model_dump_json())
+    study.set_user_attr("finished", False)
+
+    def count_completed_trials() -> int:
+        # Count number of complete trials to compute trials to run.
+        return sum([(1 if t.state == TrialState.COMPLETE else 0) for t in study.trials])
+
+    start_index = trial_index = count_completed_trials()
+    if start_index > 0:
+        print()
+        print("Resuming existing study.")
+
     try:
-        study.optimize(objective_wrapper, n_trials=settings.n_trials)
+        study.optimize(
+            objective_wrapper,
+            n_trials=settings.n_trials - count_completed_trials(),
+        )
     except KeyboardInterrupt:
         # This additional handler takes care of the small chance that KeyboardInterrupt
         # is raised just between trials, which wouldn't be caught by the handler
         # defined in objective_wrapper above.
         pass
 
+    if count_completed_trials() == settings.n_trials:
+        study.set_user_attr("finished", True)
+
+    while True:
         # If no trials at all have been evaluated, the study must have been stopped
         # by pressing Ctrl+C while the first trial was running. In this case, we just
         # re-raise the interrupt to invoke the standard handler defined below.
-    if not study.best_trials:
+        completed_trials = [t for t in study.trials if t.state == TrialState.COMPLETE]
+        if not completed_trials:
             raise KeyboardInterrupt
 
-    best_trials = sorted(
-        study.best_trials,
-        key=lambda trial: trial.user_attrs["refusals"],
+        # Get the Pareto front of trials. We can't use study.best_trials directly
+        # as get_score() doesn't return the pure KL divergence and refusal count.
+        # Note: Unlike study.best_trials, this does not handle objective constraints.
+        sorted_trials = sorted(
+            completed_trials,
+            key=lambda trial: (
+                trial.user_attrs["refusals"],
+                trial.user_attrs["kl_divergence"],
+            ),
         )
+        min_divergence = math.inf
+        best_trials = []
+        for trial in sorted_trials:
+            kl_divergence = trial.user_attrs["kl_divergence"]
+            if kl_divergence < min_divergence:
+                min_divergence = kl_divergence
+                best_trials.append(trial)
 
         choices = [
             Choice(
                 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,
             )
@@ -400,7 +747,14 @@ def run():
 
         choices.append(
             Choice(
-            title="None (exit program)",
+                title="Run additional trials",
+                value="continue",
+            )
+        )
+
+        choices.append(
+            Choice(
+                title="Exit program",
                 value="",
             )
         )
@@ -421,18 +775,79 @@ def run():
             print()
             trial = prompt_select("Which trial do you want to use?", choices)
 
-        if trial is None or trial == "":
+            if trial == "continue":
+                while True:
+                    try:
+                        n_additional_trials = prompt_text(
+                            "How many additional trials do you want to run?"
+                        )
+                        if n_additional_trials is None or n_additional_trials == "":
+                            n_additional_trials = 0
                             break
+                        n_additional_trials = int(n_additional_trials)
+                        if n_additional_trials > 0:
+                            break
+                        print("[red]Please enter a number greater than 0.[/]")
+                    except ValueError:
+                        print("[red]Please enter a number.[/]")
+
+                if n_additional_trials == 0:
+                    continue
+
+                settings.n_trials += n_additional_trials
+                study.set_user_attr("settings", settings.model_dump_json())
+                study.set_user_attr("finished", False)
+
+                try:
+                    study.optimize(
+                        objective_wrapper,
+                        n_trials=settings.n_trials - count_completed_trials(),
+                    )
+                except KeyboardInterrupt:
+                    pass
+
+                if count_completed_trials() == settings.n_trials:
+                    study.set_user_attr("finished", True)
+
+                break
+
+            elif trial is None or trial == "":
+                return
 
             print()
             print(f"Restoring model from trial [bold]{trial.user_attrs['index']}[/]...")
+            print("* Parameters:")
+            for name, value in get_trial_parameters(settings, trial).items():
+                print(f"  * {name} = [bold]{value}[/]")
+            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.reload_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"],
-            trial.user_attrs["parameters"],
+                    {
+                        k: AbliterationParameters(**v)
+                        for k, v in trial.user_attrs["parameters"].items()
+                    },
                 )
 
             while True:
@@ -443,11 +858,12 @@ def run():
                         "Save the model to a local folder",
                         "Upload the model to Hugging Face",
                         "Chat with the model",
-                    "Nothing (return to trial selection menu)",
+                        "Benchmark the model",
+                        "Return to the trial selection menu",
                     ],
                 )
 
-            if action is None or action == "Nothing (return to trial selection menu)":
+                if action is None or action == "Return to the trial selection menu":
                     break
 
                 # All actions are wrapped in a try/except block so that if an error occurs,
@@ -460,9 +876,25 @@ def run():
                             if not save_directory:
                                 continue
 
-                        print("Saving model...")
+                            strategy = obtain_merge_strategy(settings)
+                            if strategy is None:
+                                continue
+
+                            if strategy == "adapter":
+                                print("Saving LoRA adapter...")
                                 model.model.save_pretrained(save_directory)
+                            else:
+                                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()
                                 model.tokenizer.save_pretrained(save_directory)
+
                             print(f"Model saved to [bold]{save_directory}[/].")
 
                         case "Upload the model to Hugging Face":
@@ -495,26 +927,58 @@ def run():
                                     "Private",
                                 ],
                             )
+                            if visibility is None:
+                                continue
                             private = visibility == "Private"
 
-                        print("Uploading model...")
+                            strategy = obtain_merge_strategy(settings)
+                            if strategy is None:
+                                continue
 
+                            if strategy == "adapter":
+                                print("Uploading LoRA adapter...")
                                 model.model.push_to_hub(
                                     repo_id,
                                     private=private,
                                     token=token,
                                 )
+                            else:
+                                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,
+                                    token=token,
+                                )
+                                del merged_model
+                                empty_cache()
                                 model.tokenizer.push_to_hub(
                                     repo_id,
                                     private=private,
                                     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:
@@ -523,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,
@@ -559,11 +1031,120 @@ def run():
 
                                     print("[bold]Assistant:[/] ", end="")
                                     response = model.stream_chat_response(chat)
-                                chat.append({"role": "assistant", "content": response})
+                                    chat.append(
+                                        {"role": "assistant", "content": response}
+                                    )
                                 except (KeyboardInterrupt, EOFError):
                                     # 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/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

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025  Philipp Emanuel Weidmann <pew@worldwidemann.com>
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
 import gc
 import getpass
 import os
-from dataclasses import asdict
+from dataclasses import dataclass
 from importlib.metadata import version
 from pathlib import Path
 from typing import Any, TypeVar
@@ -17,19 +17,44 @@ from accelerate.utils import (
     is_sdaa_available,
     is_xpu_available,
 )
-from datasets import ReadInstruction, load_dataset, load_from_disk
+from datasets import DatasetDict, ReadInstruction, load_dataset, load_from_disk
 from datasets.config import DATASET_STATE_JSON_FILENAME
 from datasets.download.download_manager import DownloadMode
 from datasets.utils.info_utils import VerificationMode
 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
 
 
+def print_memory_usage():
+    def p(label: str, size_in_bytes: int):
+        print(f"[grey50]{label}: [bold]{size_in_bytes / (1024**3):.2f} GB[/][/]")
+
+    p("Resident system RAM", Process().memory_info().rss)
+
+    if torch.cuda.is_available():
+        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():
+        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())
+
+
 def is_notebook() -> bool:
     # Check for specific environment variables (Colab, Kaggle).
     # This is necessary because when running as a subprocess (e.g. !heretic),
@@ -39,7 +64,7 @@ def is_notebook() -> bool:
 
     # Check IPython shell type (for library usage).
     try:
-        from IPython import get_ipython  # pyright: ignore[reportMissingModuleSource]
+        from IPython import get_ipython  # ty:ignore[unresolved-import]
 
         shell = get_ipython()
         if shell is None:
@@ -136,7 +161,16 @@ def format_duration(seconds: float) -> str:
         return f"{seconds}s"
 
 
-def load_prompts(specification: DatasetSpecification) -> list[str]:
+@dataclass
+class Prompt:
+    system: str
+    user: str
+
+
+def load_prompts(
+    settings: Settings,
+    specification: DatasetSpecification,
+) -> list[Prompt]:
     path = specification.dataset
     split_str = specification.split
 
@@ -145,6 +179,9 @@ def load_prompts(specification: DatasetSpecification) -> list[str]:
             # Dataset saved with datasets.save_to_disk; needs special handling.
             # Path should be the subdirectory for a particular split.
             dataset = load_from_disk(path)
+            assert not isinstance(dataset, DatasetDict), (
+                "Loading dataset dicts is not supported"
+            )
             # Parse the split instructions.
             instruction = ReadInstruction.from_spec(split_str)
             # Associate the split with its number of examples (lines).
@@ -168,7 +205,27 @@ def load_prompts(specification: DatasetSpecification) -> list[str]:
         # Probably a repository path; let load_dataset figure it out.
         dataset = load_dataset(path, split=split_str)
 
-    return list(dataset[specification.column])
+    prompts = list(dataset[specification.column])
+
+    if specification.prefix:
+        prompts = [f"{specification.prefix} {prompt}" for prompt in prompts]
+
+    if specification.suffix:
+        prompts = [f"{prompt} {specification.suffix}" for prompt in prompts]
+
+    system_prompt = (
+        settings.system_prompt
+        if specification.system_prompt is None
+        else specification.system_prompt
+    )
+
+    return [
+        Prompt(
+            system=system_prompt,
+            user=prompt,
+        )
+        for prompt in prompts
+    ]
 
 
 T = TypeVar("T")
@@ -178,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.
@@ -189,18 +254,26 @@ def empty_cache():
     elif is_xpu_available():
         torch.xpu.empty_cache()
     elif is_mlu_available():
-        torch.mlu.empty_cache()
+        torch.mlu.empty_cache()  # ty:ignore[unresolved-attribute]
     elif is_sdaa_available():
-        torch.sdaa.empty_cache()
+        torch.sdaa.empty_cache()  # ty:ignore[unresolved-attribute]
     elif is_musa_available():
-        torch.musa.empty_cache()
+        torch.musa.empty_cache()  # ty:ignore[unresolved-attribute]
     elif torch.backends.mps.is_available():
         torch.mps.empty_cache()
 
     gc.collect()
 
 
-def get_trial_parameters(trial: Trial) -> dict[str, str]:
+def get_trial_parameters(settings: Settings, trial: Trial) -> dict[str, str]:
+    if settings.use_ara:
+        parameters = trial.user_attrs["ara_parameters"]
+
+        return {
+            name: (f"{value:.4f}" if isinstance(value, float) else f"{value}")
+            for name, value in parameters.items()
+        }
+    else:
         params = {}
 
         direction_index = trial.user_attrs["direction_index"]
@@ -209,23 +282,48 @@ def get_trial_parameters(trial: Trial) -> dict[str, str]:
         )
 
         for component, parameters in trial.user_attrs["parameters"].items():
-        for name, value in asdict(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(
     settings: Settings,
     trial: Trial,
     base_refusals: int,
-    bad_prompts: list[str],
+    bad_prompts: list[Prompt],
 ) -> str:
+    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
 
@@ -235,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()
             ]
         )
     }
@@ -244,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)
     } |