research: Token-level training with span-max loss from Obeso et al. §3

[peaktwilight]

What

Move from snippet-level logistic regression to the proper token-level + span-max formulation described in Obeso et al. §3 (https://arxiv.org/abs/2509.03531).

The current baseline trains sklearn LogisticRegression on snippet-level labels (one label per snippet, taken at the last-token activation) and reaches AUC 0.846. The paper's formulation trains a small PyTorch head with:

L = (1-ω) · Σ w_i · BCE(y_i, p_i) + ω · Σ BCE(y_s, max p_i for i in s)

where ω is annealed 0 → 1 over training (token-level supervision early, span-level later), and entity tokens are up-weighted by α = 10.

Why

Token-level supervision should give us calibrated per-token scores (which the UI already pretends to render) instead of broadcasting a single snippet label across every token. The span-max term then re-aligns the loss with the snippet-level objective we actually care about at inference.

Status / context

• src/train_probe_spanmax.py — partial implementation of the loss + training loop
• src/extract_token_activations.py — per-token + span activation extractor (done)
• Integration gap: wire the activation extraction pipeline end-to-end and report AUC

Definition of done

• Span-max probe trained on the merged CyberSecEval + SVEN dataset
• AUC ≥ 0.86 token-level, ≥ 0.88 example-level
• Replaces data/probe.npz

[peaktwilight]

Token-level extraction attempted on the merged 2,758-row dataset. The script hung at 950/2758 (~34%) — process state went stuck with 36 GB RSS, suggesting MPS unified-memory pressure from holding all positional hidden states across all 4 captured layers in memory simultaneously.

Fix needed for the next run:

• Stream activations to disk per-example instead of accumulating in RAM
• Or shrink to a single layer at a time (recompute is cheap once weights are cached)
• Or use a server with discrete GPU + more memory headroom

The single-position extraction (the one that produced the 0.963 probe) was fine; the per-token version is what hits the memory ceiling.

Pinning so we don't lose the diagnosis.