feat: comprehensive wandb logging (main-outcomes, panel routing, diag norms)

[RyanKim17920]

Summary

• Panel routing: all wandb keys are now routed into logical panels instead of a single flat namespace. Headline losses (dino, ibot, kde, total, grad_norm, lr, and their val_* counterparts) land in main-train-results/; bookkeeping metrics (throughput, memory, ETA, etc.) land in train-tracking/; diagnostic keys keep their diag/ prefix; probe keys keep their probe/ prefix.
• main-outcomes/ panel: after each probe run, headline downstream scores (linear, kNN, 16-shot, segmentation, progression, mutation, survival, robustness, mean) are logged to a clean main-outcomes/ panel and per-dataset breakdowns to main-outcomes-datasets/. This surfaces the numbers that actually matter at a glance rather than hunting through flat probe/* keys.
• Diagnostic token/attention norms (diag/): when train.log_diagnostics_every is set in the config (default 0 = off), the student backbone runs a diagnostic forward pass that captures per-layer q/k/v norms, qk similarity, and cls/reg/patch token norms every N log steps. Four evenly-spaced layers are sampled plus a per-metric mean across all layers.

What was NOT changed

No architecture, dataset, batch size, probe logic, or config defaults were modified. The diagnostic forward path is opt-in and unreachable unless log_diagnostics_every > 0. The normal (non-diagnostic) forward path in both model.py and train.py is identical to before this PR.

Test plan

• Smoke-test passes (WANDB_MODE=disabled python train.py ...) — existing training loop unchanged for default config
• With log_diagnostics_every: 10 in config, confirm diag/ keys appear in wandb at the correct steps
• After a probe run, confirm main-outcomes/mean and main-outcomes-datasets/ keys appear in wandb
• Val loop still works (unpacks 4 values from compute_losses, discards diag)

🤖 Generated with Claude Code

[PaulScotti]

Thanks for the PR! I appreciate trying to clean up our logging/plotting to be more informative

I have several concerns though, you made lots of changes and I'm not sure all are good/intended:

1. model.py (line 168): diagnostic mode bypasses activation checkpointing. With checkpoint=True, normal forward called checkpoint 4 times
2. train.py (line 571): if log_diagnostics_every: 1, diagnostics run on step 1 while FlopCounterMode measures measured_flops_per_step, so the permanent FLOP estimate includes diagnostic-only work
3. train.py (line 477): log_main_outcomes() runs every time log_probe_results() runs, not only when a new probe result is collected. Testing it out, two calls logged twice at the same W&B step
4. train.py (line 508): not sure why you need to rename these probe outputs, isnt "probe/bracs_score" fine?
5. train.py (line 488): usage of broad try/except conflicts with our preferred fail-loudly rule and MedARC coding standards
6. train.py (line 297): renames existing W&B keys from stable train/* and val/* namespaces into new namespaces? Why do this
7. train.py (line 536): train.log_diagnostics_every is an undocumented .get(..., 0) fallback, again not desirably coding style

[RyanKim17920]

Thanks for the feedback! I've updated the code accordingly and addressed all your major points:

1. The diagnostic forward now runs through the activation-checkpoint path too (matches the normal forward), verified on a GPU smoke test
2. Diagnostics no longer run on the FLOP-measurement step
3. Double logging is fixed now, main-outcomes logs on probe results instead of train log steps
4. We use the same names now (keeping _score for consistency)
5. Removed the try/excepts to match with code standards
6. train/* and val/* are unchanged from beforehand, but I mirror the main losses into a main-train-results/ panel for visibility since train/* had lots of diagnostic/throughput metrics.
7. log_diagnostics_every is defaulted to 0 in the yaml and uses this instead of a get fallback