Log ensembled metrics

elk/evaluation/evaluate.py

@@ -45,21 +45,32 @@ def apply_to_layer(
  for ds_name, (val_h, val_gt, _) in val_output.items():
  meta = {"dataset": ds_name, "layer": layer}
 
- val_result = evaluate_preds(val_gt, reporter(val_h))
- row_bufs["eval"].append({**meta, **val_result.to_dict()})
+ val_credences = reporter(val_h)
+ for mode in ("none", "partial", "full"):
+ row_bufs["eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(val_gt, val_credences, mode).to_dict(),
+ }
+ )
 
- lr_dir = experiment_dir / "lr_models"
- if not self.skip_supervised and lr_dir.exists():
- with open(lr_dir / f"layer_{layer}.pt", "rb") as f:
- lr_models = torch.load(f, map_location=device)
- if not isinstance(lr_models, list): # backward compatibility
- lr_models = [lr_models]
+  lr_dir = experiment_dir / "lr_models"
+  if not self.skip_supervised and lr_dir.exists():
+  with open(lr_dir / f"layer_{layer}.pt", "rb") as f:
+  lr_models = torch.load(f, map_location=device)
+  if not isinstance(lr_models, list): # backward compatibility
+  lr_models = [lr_models]
 
- for i, model in enumerate(lr_models):
- model.eval()
- lr_result = evaluate_preds(val_gt, model(val_h))
- row_bufs["lr_eval"].append(
- {"inlp_iter": i, **meta, **lr_result.to_dict()}
- )
+ for i, model in enumerate(lr_models):
+ model.eval()
+ row_bufs["lr_eval"].append(
+ {
+ "ensembling": mode,
+ "inlp_iter": i,
+ **meta,
+ **evaluate_preds(val_gt, model(val_h), mode).to_dict(),
+ }
+ )
 
  return {k: pd.DataFrame(v) for k, v in row_bufs.items()}

elk/metrics/eval.py

@@ -1,4 +1,5 @@
 from dataclasses import asdict, dataclass
+from typing import Literal
 
 import torch
 from einops import repeat
@@ -37,33 +38,50 @@ def to_dict(self, prefix: str = "") -> dict[str, float]:
  else {}
  )
  auroc_dict = {f"{prefix}auroc_{k}": v for k, v in asdict(self.roc_auc).items()}
- return {**acc_dict, **cal_acc_dict, **cal_dict, **auroc_dict}
+ return {**auroc_dict, **cal_acc_dict, **acc_dict, **cal_dict}
 
 
-def evaluate_preds(y_true: Tensor, y_logits: Tensor) -> EvalResult:
+def evaluate_preds(
+ y_true: Tensor,
+ y_logits: Tensor,
+ ensembling: Literal["none", "partial", "full"] = "none",
+) -> EvalResult:
  """
  Evaluate the performance of a classification model.
 
  Args:
  y_true: Ground truth tensor of shape (N,).
- y_pred: Predicted class tensor of shape (N, variants, n_classes).
+ y_logits: Predicted class tensor of shape (N, variants, n_classes).
 
  Returns:
  dict: A dictionary containing the accuracy, AUROC, and ECE.
  """
  (n, v, c) = y_logits.shape
  assert y_true.shape == (n,)
 
- # Clustered bootstrap confidence intervals for AUROC
- y_true = repeat(y_true, "n -> n v", v=v)
- auroc = roc_auc_ci(to_one_hot(y_true, c).long().flatten(1), y_logits.flatten(1))
- acc = accuracy_ci(y_true, y_logits.argmax(dim=-1))
-
+ if ensembling == "full":
+ y_logits = y_logits.mean(dim=1)
+ else:
+ y_true = repeat(y_true, "n -> n v", v=v)
+
+ y_pred = y_logits.argmax(dim=-1)
+ if ensembling == "none":
+ auroc = roc_auc_ci(to_one_hot(y_true, c).long().flatten(1), y_logits.flatten(1))
+ elif ensembling in ("partial", "full"):
+ # Pool together the negative and positive class logits
+ if c == 2:
+ auroc = roc_auc_ci(y_true, y_logits[..., 1] - y_logits[..., 0])
+ else:
+ auroc = roc_auc_ci(to_one_hot(y_true, c).long(), y_logits)
+ else:
+ raise ValueError(f"Unknown mode: {ensembling}")
+
+ acc = accuracy_ci(y_true, y_pred)
  cal_acc = None
  cal_err = None
 
  if c == 2:
- pos_probs = y_logits.softmax(-1)[..., 1]
+ pos_probs = torch.sigmoid(y_logits[..., 1] - y_logits[..., 0])
 
  # Calibrated accuracy
  cal_thresh = pos_probs.float().quantile(y_true.float().mean())

elk/run.py

@@ -173,7 +173,7 @@ def apply_to_layers(
  finally:
  # Make sure the CSVs are written even if we crash or get interrupted
  for name, dfs in df_buffers.items():
- df = pd.concat(dfs).sort_values(by="layer")
+ df = pd.concat(dfs).sort_values(by=["layer", "ensembling"])
  df.round(4).to_csv(self.out_dir / f"{name}.csv", index=False)
  if self.debug:
  save_debug_log(self.datasets, self.out_dir)

elk/training/train.py

@@ -24,7 +24,7 @@
  """Full specification of a reporter training run."""
 
  net: ReporterConfig = subgroups(
  {"ccs": CcsReporterConfig, "eigen": EigenReporterConfig}, default="eigen"
  )
  """Config for building the reporter network."""
 
@@ -124,24 +124,35 @@
  for ds_name, (val_h, val_gt, val_lm_preds) in val_dict.items():
  meta = {"dataset": ds_name, "layer": layer}
 
- val_result = evaluate_preds(val_gt, reporter(val_h))
- row_bufs["eval"].append(
- {
- **meta,
- "pseudo_auroc": pseudo_auroc,
- "train_loss": train_loss,
- **val_result.to_dict(),
- }
- )
-
- if val_lm_preds is not None:
- lm_result = evaluate_preds(val_gt, val_lm_preds)
- row_bufs["lm_eval"].append({**meta, **lm_result.to_dict()})
-
- for i, model in enumerate(lr_models):
- lr_result = evaluate_preds(val_gt, model(val_h))
- row_bufs["lr_eval"].append(
- {"inlp_iter": i, **meta, **lr_result.to_dict()}
+ val_credences = reporter(val_h)
+ for mode in ("none", "partial", "full"):
+ row_bufs["eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(val_gt, val_credences, mode).to_dict(),
+ "pseudo_auroc": pseudo_auroc,
+ "train_loss": train_loss,
+ }
  )
 
+ if val_lm_preds is not None:
+ row_bufs["lm_eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(val_gt, val_lm_preds, mode).to_dict(),
+ }
+ )
+
+ for i, model in enumerate(lr_models):
+ row_bufs["lr_eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ "inlp_iter": i,
+ **evaluate_preds(val_gt, model(val_h), mode).to_dict(),
+ }
+ )
+
  return {k: pd.DataFrame(v) for k, v in row_bufs.items()}