Merge pull request #288 from EleutherAI/fix-platt-scaling-ccs

add platt scaling for burns + fix for leace

elk/training/ccs_reporter.py

@@ -9,6 +9,7 @@
 import torch.nn as nn
 from concept_erasure import LeaceFitter
 from torch import Tensor
+from typing_extensions import override
 
 from ..parsing import parse_loss
 from ..utils.typing import assert_type
@@ -88,6 +89,7 @@ def __init__(
         num_variants: int = 1,
     ):
         super().__init__()
+
         self.config = cfg
         self.in_features = in_features
         self.num_variants = num_variants
@@ -128,6 +130,15 @@ def __init__(
                 )
             )
 
+    @override
+    def parameters(self, recurse=True):
+        parameters = super(CcsReporter, self).parameters(recurse=recurse)
+        for param in parameters:
+            # exclude the platt scaling parameters
+            # kind of a hack for now, we should find probably a cleaner way
+            if param is not self.scale and param is not self.bias:
+                yield param
+
     def reset_parameters(self):
         """Reset the parameters of the probe.
 
@@ -164,15 +175,9 @@ def reset_parameters(self):
     def forward(self, x: Tensor) -> Tensor:
         """Return the credence assigned to the hidden state `x`."""
         assert self.norm is not None, "Must call fit() before forward()"
-
         raw_scores = self.probe(self.norm(x)).squeeze(-1)
-        if self.config.norm == "leace":
-            return raw_scores.mul(self.scale).add(self.bias).squeeze(-1)
-
-        elif self.config.norm == "burns":
-            return raw_scores
-        else:
-            raise ValueError(f"Unknown normalization {self.config.norm}.")
+        platt_scaled_scores = raw_scores.mul(self.scale).add(self.bias).squeeze(-1)
+        return platt_scaled_scores
 
     def loss(self, logit0: Tensor, logit1: Tensor) -> Tensor:
         """Return the loss of the reporter on the contrast pair (x0, x1).
@@ -248,6 +253,7 @@ def fit(self, hiddens: Tensor) -> float:
             raise RuntimeError("Got NaN/infinite loss during training")
 
         self.load_state_dict(best_state)
+
         return best_loss
 
     def train_loop_adam(self, x_neg: Tensor, x_pos: Tensor) -> float:

elk/training/train.py

@@ -82,13 +82,8 @@ def apply_to_layer(
 
             reporter = CcsReporter(self.net, d, device=device, num_variants=v)
             train_loss = reporter.fit(first_train_h)
-
-            if not self.net.norm == "burns":
-                (_, v, k, _) = first_train_h.shape
-                reporter.platt_scale(
-                    to_one_hot(repeat(train_gt, "n -> (n v)", v=v), k).flatten(),
-                    rearrange(first_train_h, "n v k d -> (n v k) d"),
-                )
+            labels = repeat(to_one_hot(train_gt, k), "n k -> n v k", v=v)
+            reporter.platt_scale(labels, first_train_h)
 
         elif isinstance(self.net, EigenFitterConfig):
             fitter = EigenFitter(