feat: weighted loss, AUC, ECE (#17)

* feat(probe): add weighted cross-entropy loss, sweep-to-final pipeline, and vectorized binning

- Add --loss {cross_entropy,weighted_cross_entropy} with per-bin adaptive
  class weights w_pos = λ · n_neg/n_pos to counter class imbalance
- Add --pos-weight λ (log-uniform swept in [0.1, 10.0]) as tunable multiplier
- Add sweep --then-final to automatically fetch best W&B config and train
  all layers in one command
- Add final --from-sweep to load best HPs from an existing sweep via W&B API
- Add fetch_best_sweep_config() to query sweep.best_run() programmatically
- Vectorize bin_mask computation: pre-compute bin_ids tensor instead of
  per-element .item() Python loop (O(N × n_bins) → O(N))

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* docs(readme): document weighted CE loss, sweep→final pipeline, and --from-sweep

Add sections for --loss {cross_entropy,weighted_cross_entropy}, --pos-weight,
--then-final (one-command sweep→final), and --from-sweep (auto-load best HPs).
Update SLURM examples to show new flags.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(probe): add fixed_params to sweep, ECE metric, and extend dashboard exports

- create_sweep() now accepts optional fixed_params dict to pin specific HPs
  while sweeping the remaining ones (e.g., sweep only pos_weight)
- sweep subcommand gains --fixed JSON arg for CLI-driven fixed params
- Added Expected Calibration Error (ECE) to ProbeResult and training loop
- Dashboard exports now include test_auc, val_auc, test_ece, val_ece
- Bumped SLURM sweep time limit from 2h to 4h for sweep+final chains

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat(probe): save model weights after final training, add ECE to tests

- train_probe_layer() now returns (results, weights) tuple — weights dict
  contains per-bin state_dict and training mean for inference
- run_final() saves weights.pt alongside results.pt
- run_sweep() discards weights (sweeps don't need them)
- Fixed tests for new ProbeResult fields (val_ece, test_ece) and new
  return type of train_probe_layer()

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* fix(probe): require explicit n_bins=1 for pooled training with n_eval_bins

Replace implicit train_n_bins override (1 if n_eval_bins set) with explicit
validation: n_eval_bins must equal n_bins or n_bins must be 1. Update CLI
help text to reflect the new constraint.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Apply suggestion from @andre15silva

---------

Co-authored-by: Tux <tux.tu@outlook.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

Author: 3 people

README.md

@@ -22,8 +22,9 @@ run_build_cache.py       concatenate .pt files into per-(layer, probe) tensors
                          → cache/<run_id>/<probe>_layer<n>.pt
         │
         ▼
-run_probe.py sweep       W&B hyperparameter sweep → pick best lr/wd/batch-size
+run_probe.py sweep       W&B hyperparameter sweep → pick best lr/wd/batch-size/loss
 run_probe.py final       train linear probe on best hparams → W&B run with metrics
+                          (or pass --from-sweep to auto-load best HPs from a sweep)
         │
         ▼
 run_figures.py           accuracy heatmaps per layer × relative-position bin
@@ -192,29 +193,69 @@ uv run python run_build_cache.py \
   --probe will_be_correct
 ```
 
-### 3. Probe sweep (W&B)
+### 3. Probe training
+
+Probe training has three modes: **manual**, **sweep-only**, and **sweep→final** (one-shot).
+
+#### Loss function
+
+The `--loss` flag selects the objective:
+
+| `--loss` | Behavior |
+|---|---|
+| `cross_entropy` (default) | Standard `CrossEntropyLoss` — no class weighting |
+| `weighted_cross_entropy` | `CrossEntropyLoss(weight=[1.0, λ · n_neg/n_pos])` — adaptive per-bin class balancing |
+
+When using `weighted_cross_entropy`, the `--pos-weight` λ parameter (default 1.0) multiplies the inverse class ratio:
+- λ = 1.0 → exactly balanced classes per bin
+- λ > 1.0 → penalize false negatives more (up-weight positives)
+- λ < 1.0 → penalize false positives more (down-weight positives)
+
+#### Mode A: Manual (specify all HPs)
 
 ```bash
-uv run python run_probe.py sweep \
-  --run-id my_run \
-  --probe will_be_correct \
-  --model-config configs/models/qwen3_8b.yaml
+uv run python run_probe.py \
+  --run-id my_run --probe will_be_correct \
+  --model-config configs/models/qwen3_8b.yaml \
+  final \
+  --lr 1e-3 --weight-decay 1e-4 --batch-size 512 --patience 10 \
+  --loss weighted_cross_entropy --pos-weight 1.0
 ```
 
-### 4. Probe final run
+#### Mode B: Sweep only
 
 ```bash
-uv run python run_probe.py final \
-  --run-id my_run \
-  --probe will_be_correct \
+uv run python run_probe.py \
+  --run-id my_run --probe will_be_correct \
   --model-config configs/models/qwen3_8b.yaml \
-  --lr 1e-3 \
-  --weight-decay 1e-4 \
-  --batch-size 512 \
-  --patience 10
+  sweep --count 50
 ```
 
-### 5. Figures
+The sweep explores `lr`, `weight_decay`, `batch_size`, `patience`, `loss` (cross_entropy vs weighted_cross_entropy), and `pos_weight` (log-uniform in [0.1, 10.0]) via Bayesian optimisation. Copy the best HPs from the W&B dashboard.
+
+#### Mode C: Sweep → final (fully automatic)
+
+```bash
+uv run python run_probe.py \
+  --run-id my_run --probe will_be_correct \
+  --model-config configs/models/qwen3_8b.yaml \
+  sweep --count 50 --then-final
+```
+
+Runs the sweep, waits for completion, fetches the best config from W&B via `fetch_best_sweep_config()`, and trains on all probe layers — no manual copy-paste.
+
+#### Mode D: Final from an existing sweep
+
+```bash
+uv run python run_probe.py \
+  --run-id my_run --probe will_be_correct \
+  --model-config configs/models/qwen3_8b.yaml \
+  final --from-sweep abc123xyz
+```
+
+Queries `sweep.best_run()` via the W&B API to auto-populate `lr`, `weight_decay`, `batch_size`, `patience`, `loss`, and `pos_weight`.
+
+### 4. Figures
 
 ```bash
 uv run python run_figures.py \
@@ -223,7 +264,7 @@ uv run python run_figures.py \
   --model-config configs/models/qwen3_8b.yaml
 ```
 
-### 6. Export dashboard
+### 5. Export dashboard
 
 ```bash
 uv run python run_export_dashboard.py \
@@ -257,10 +298,18 @@ sbatch --array=0-29 slurm/probe_sweep.sh \
   --run-id my_run --probe will_be_correct \
   --model-config configs/models/qwen3_8b.yaml
 
+# Final with manual HPs:
+sbatch slurm/probe_final.sh \
+  --run-id my_run --probe will_be_correct \
+  --model-config configs/models/qwen3_8b.yaml \
+  --lr 1e-3 --weight-decay 1e-4 --batch-size 512 --patience 10 \
+  --loss weighted_cross_entropy --pos-weight 1.0
+
+# Final auto-loaded from sweep:
 sbatch slurm/probe_final.sh \
   --run-id my_run --probe will_be_correct \
   --model-config configs/models/qwen3_8b.yaml \
-  --lr 1e-3 --weight-decay 1e-4 --batch-size 512 --patience 10
+  --from-sweep abc123xyz
 
 sbatch slurm/figures.sh --run-id my_run --probe will_be_correct \
   --model-config configs/models/qwen3_8b.yaml

run_probe.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 from src.configs import ModelConfig, load_config
 from src.probe import run_sweep, run_final
 
@@ -28,12 +29,22 @@ def main():
     sweep_p = subparsers.add_parser("sweep")
     sweep_p.add_argument("--sweep-id", default=None, help="Join an existing W&B sweep instead of creating a new one")
     sweep_p.add_argument("--count", type=int, default=None, help="Max number of runs this agent will execute")
+    sweep_p.add_argument("--then-final", action="store_true",
+                        help="After sweep, automatically train on all layers with best HPs")
+    sweep_p.add_argument("--fixed", type=json.loads, default=None,
+                        help='JSON dict of fixed HP values, e.g. \'{"lr": 0.001, "loss": "weighted_cross_entropy"}\'')
 
     final_p = subparsers.add_parser("final")
-    final_p.add_argument("--lr", type=float, required=True)
-    final_p.add_argument("--weight-decay", type=float, required=True)
-    final_p.add_argument("--batch-size", type=int, required=True)
-    final_p.add_argument("--patience", type=int, required=True)
+    final_p.add_argument("--lr", type=float, default=None)
+    final_p.add_argument("--weight-decay", type=float, default=None)
+    final_p.add_argument("--batch-size", type=int, default=None)
+    final_p.add_argument("--patience", type=int, default=None)
+    final_p.add_argument("--from-sweep", default=None,
+                        help="Sweep ID to load best hyperparameters from")
+    final_p.add_argument("--loss", choices=["cross_entropy", "weighted_cross_entropy"],
+                        default="cross_entropy")
+    final_p.add_argument("--pos-weight", type=float, default=1.0,
+                        help="Positive class weight multiplier (only used with --loss weighted_cross_entropy)")
 
     args = parser.parse_args()
     if args.n_eval_bins is not None and args.n_eval_bins != args.n_bins and args.n_bins != 1:
@@ -55,6 +66,9 @@ def main():
             probe_arch=args.probe_arch,
             n_eval_bins=args.n_eval_bins,
             eval_bin_axis=args.eval_bin_axis,
+            then_final=args.then_final,
+            results_dir=args.results_dir,
+            fixed_params=args.fixed,
             shuffle_labels=args.shuffle_labels,
         )
     else:
@@ -67,11 +81,14 @@ def main():
             batch_size=args.batch_size,
             patience=args.patience,
             seed=args.seed,
+            sweep_id=args.from_sweep,
             cache_dir=args.cache_dir,
             cache_run_id=cache_run_id,
             results_dir=args.results_dir,
             n_bins=args.n_bins,
             probe_arch=args.probe_arch,
+            loss=args.loss,
+            pos_weight=args.pos_weight,
             n_eval_bins=args.n_eval_bins,
             eval_bin_axis=args.eval_bin_axis,
             shuffle_labels=args.shuffle_labels,

slurm/probe_sweep.sh

@@ -3,7 +3,7 @@
 #SBATCH -p berzelius-cpu
 #SBATCH -n 8
 #SBATCH --mem=150G
-#SBATCH -t 02:00:00
+#SBATCH -t 04:00:00
 #SBATCH -o logs/probe_sweep_%j.out
 #SBATCH -e logs/probe_sweep_%j.err
 

src/export_dashboard.py

@@ -146,6 +146,10 @@ def export_dashboard(
                 probe_results[probe_name][str(layer_idx)][str(bin_idx)] = {
                     "test_acc": r.test_acc if hasattr(r, "test_acc") else r["test_acc"],
                     "val_acc": r.val_acc if hasattr(r, "val_acc") else r["val_acc"],
+                    "test_auc": r.test_auc if hasattr(r, "test_auc") else r.get("test_auc", None),
+                    "val_auc": r.val_auc if hasattr(r, "val_auc") else r.get("val_auc", None),
+                    "test_ece": r.test_ece if hasattr(r, "test_ece") else r.get("test_ece", None),
+                    "val_ece": r.val_ece if hasattr(r, "val_ece") else r.get("val_ece", None),
                     "n_train": r.n_train if hasattr(r, "n_train") else r["n_train"],
                     "n_val": r.n_val if hasattr(r, "n_val") else r["n_val"],
                     "n_test": r.n_test if hasattr(r, "n_test") else r["n_test"],

src/export_swebench_dashboard.py

@@ -290,6 +290,8 @@ def _load_labels_for(f: Path) -> dict:
                     "val_acc": r.val_acc if hasattr(r, "val_acc") else r["val_acc"],
                     "test_auc": r.test_auc if hasattr(r, "test_auc") else r.get("test_auc"),
                     "val_auc": r.val_auc if hasattr(r, "val_auc") else r.get("val_auc"),
+                    "test_ece": r.test_ece if hasattr(r, "test_ece") else r.get("test_ece"),
+                    "val_ece": r.val_ece if hasattr(r, "val_ece") else r.get("val_ece"),
                     "n_train": r.n_train if hasattr(r, "n_train") else r["n_train"],
                     "n_val": r.n_val if hasattr(r, "n_val") else r["n_val"],
                     "n_test": r.n_test if hasattr(r, "n_test") else r["n_test"],