default meanonly norm

EleutherAI · Nov 14, 2023 · 6e07819 · 6e07819
1 parent e751405
commit 6e07819
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/ccs/training/ccs_reporter.py b/ccs/training/ccs_reporter.py
@@ -43,7 +43,7 @@ class CcsConfig(FitterConfig):
     function 1.0*consistency_squared + 0.5*prompt_var.
     """
     loss_dict: dict[str, float] = field(default_factory=dict, init=False)
-    norm: Literal["leace", "burns"] = "leace"
+    norm: Literal["leace", "burns", "meanonly"] = "meanonly"
     num_layers: int = 1
     """The number of layers in the MLP."""
     pre_ln: bool = False
@@ -209,6 +209,8 @@ def fit(self, hiddens: Tensor) -> float:
 
         if self.config.norm == "burns":
             self.norm = BurnsNorm()
+        elif self.config.norm == "meanonly":
+            self.norm = BurnsNorm(scale=False)
         else:
             fitter = LeaceFitter(d, 2 * v, dtype=x_neg.dtype, device=x_neg.device)
             fitter.update(