From f7a456bd0c72b406915b8ba7957a0a2348a5d51b Mon Sep 17 00:00:00 2001
From: Philipp Emanuel Weidmann <pew@worldwidemann.com>
Date: Tue, 31 Mar 2026 15:55:27 +0530
Subject: [PATCH] feat(ara): add optional row-norm preservation

---
 src/heretic/config.py |  4 ++--
 src/heretic/model.py  | 15 ++++++++++++++-
 src/heretic/utils.py  |  9 ++++++++-
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/heretic/config.py b/src/heretic/config.py
index 1796e7c..77de3c5 100644
--- a/src/heretic/config.py
+++ b/src/heretic/config.py
@@ -205,7 +205,7 @@ class Settings(BaseSettings):
     )
 
     use_piqa: bool = Field(
-        default=True,
+        default=False,
         description=(
             "Whether to use the Physical Interaction: Question Answering (PIQA) benchmark "
             "as the quality metric instead of the Kullback-Leibler divergence."
@@ -221,7 +221,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), '
diff --git a/src/heretic/model.py b/src/heretic/model.py
index 1feb63d..108e9df 100644
--- a/src/heretic/model.py
+++ b/src/heretic/model.py
@@ -585,6 +585,16 @@ class Model:
                     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
                     ]
@@ -644,7 +654,7 @@ class Model:
 
                     def closure() -> Tensor:
                         optimizer.zero_grad()
-                        loss = objective(matrix)
+                        loss = objective(get_matrix())
                         loss.backward()
                         return loss
 
@@ -655,6 +665,9 @@ class Model:
                         #    f"\\[{layer_index}/{component}/{module_index}] Step: {step}, Loss: {loss.item():.6f}"
                         # )
 
+                    with torch.no_grad():
+                        matrix.copy_(get_matrix())
+
     def generate(
         self,
         prompts: list[Prompt],
diff --git a/src/heretic/utils.py b/src/heretic/utils.py
index 69d6468..96896ad 100644
--- a/src/heretic/utils.py
+++ b/src/heretic/utils.py
@@ -290,7 +290,14 @@ def get_trial_parameters(settings: Settings, trial: Trial) -> dict[str, str]:
 
 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"
+        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