Add hyperparamter sweep to sweep.py; Fall back to eig when eigh fails

elk/run.py

@@ -41,7 +41,7 @@ class Run(ABC, Serializable):
 
  concatenated_layer_offset: int = 0
  debug: bool = False
- min_gpu_mem: int | None = None
+ min_gpu_mem: int | None = None # in bytes
  num_gpus: int = -1
  out_dir: Path | None = None
  disable_cache: bool = field(default=False, to_dict=False)

elk/training/eigen_reporter.py

@@ -1,6 +1,7 @@
 """An ELK reporter network."""
 
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Optional
 
 import torch
@@ -28,6 +29,7 @@ class EigenReporterConfig(ReporterConfig):
  neg_cov_weight: float = 0.5
 
  num_heads: int = 1
+ save_reporter_stats: bool = False
 
  def __post_init__(self):
  if not (0 <= self.neg_cov_weight <= 1):
@@ -67,9 +69,9 @@ class EigenReporter(Reporter):
 
  config: EigenReporterConfig
 
- intercluster_cov_M2: Tensor # variance
- intracluster_cov: Tensor # invariance
- contrastive_xcov_M2: Tensor # negative covariance
+ intercluster_cov_M2: Tensor | None  # variance
+ intracluster_cov: Tensor | None  # invariance
+ contrastive_xcov_M2: Tensor | None  # negative covariance
  n: Tensor
  class_means: Tensor | None
  weight: Tensor
@@ -147,13 +149,33 @@ def consistency(self) -> Tensor:
 
  def clear(self) -> None:
  """Clear the running statistics of the reporter."""
+ assert (
+ self.contrastive_xcov_M2 is not None
+ and self.intercluster_cov_M2 is not None
+ and self.intracluster_cov is not None
+ ), "Covariance matrices have been deleted"
  self.contrastive_xcov_M2.zero_()
  self.intracluster_cov.zero_()
  self.intercluster_cov_M2.zero_()
  self.n.zero_()
 
+ def delete_stats(self) -> None:
+ """Delete the running covariance matrices.
+
+ This is useful for saving memory when we're done training the reporter.
+ """
+ self.contrastive_xcov_M2 = None
+ self.intercluster_cov_M2 = None
+ self.intracluster_cov = None
+
  @torch.no_grad()
  def update(self, hiddens: Tensor) -> None:
+ assert (
+ self.contrastive_xcov_M2 is not None
+ and self.intercluster_cov_M2 is not None
+ and self.intracluster_cov is not None
+ ), "Covariance matrices have been deleted"
+
  (n, _, k, d) = hiddens.shape
 
  # Sanity checks
@@ -206,6 +228,11 @@ def update(self, hiddens: Tensor) -> None:
  def fit_streaming(self, truncated: bool = False) -> float:
  """Fit the probe using the current streaming statistics."""
  inv_weight = 1 - self.config.neg_cov_weight
+ assert (
+ self.contrastive_xcov_M2 is not None
+ and self.intercluster_cov_M2 is not None
+ and self.intracluster_cov is not None
+ ), "Covariance matrices have been deleted"
  A = (
  self.config.var_weight * self.intercluster_cov
  - inv_weight * self.intracluster_cov
@@ -285,3 +312,9 @@ def closure():
  return float(loss)
 
  opt.step(closure)
+
+ def save(self, path: Path | str):
+ # TODO: this method will save separate JSON and PT files
+ if not self.config.save_reporter_stats:
+ self.delete_stats()
+ super().save(path)

elk/training/train.py

@@ -105,8 +105,7 @@ def apply_to_layer(
  raise ValueError(f"Unknown reporter config type: {type(self.net)}")
 
  # Save reporter checkpoint to disk
- with open(reporter_dir / f"layer_{layer}.pt", "wb") as file:
- torch.save(reporter, file)
+ reporter.save(reporter_dir / f"layer_{layer}.pt")
 
  # Fit supervised logistic regression model
  if self.supervised != "none":