feat(ara): add optional row-norm preservation

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), '

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],

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