""" Train the Stage 3B frequency classifier on cached freq-v1 features. Mirrors `scripts/train_head.py` in idiom (torch, seeded, early stop, JSON report) but operates on the 16-dim frequency features from `scripts/precompute_frequency_features.py` instead of CLIP embeddings. Model ----- A standardised linear logistic-regression head (the "classical first model" required by S3B-07): feature standardisation (mean/std from train) baked into the saved checkpoint, then `Linear(16 -> 2)`. Interpretable, tiny, no new deps (torch only). A small MLP is available via --hidden if the linear model underperforms, but linear is the default per the tracker's guidance. Checkpoint format (torch.save dict) ----------------------------------- - feature_schema: "freq-v1" (must match the features at inference) - feature_mean: float32 (16,) - feature_std: float32 (16,) - state_dict: the classifier head weights - hidden: int (0 = linear logreg) - classes: ["authentic", "ai_generated"] Usage ----- python scripts/train_frequency_head.py \\ --feat-dir data/freq_features \\ --out data/checkpoints/freq_v1.pt \\ --report-out data/reports/freq_v1_train.json """ from __future__ import annotations import argparse import json import time from pathlib import Path import numpy as np FEATURE_SCHEMA_VERSION = "freq-v1" _CLASSES = ["authentic", "ai_generated"] def _load_feat(feat_dir: Path, split: str): path = feat_dir / f"{split}_freq.npz" with np.load(path, allow_pickle=False) as z: schema = str(z["feature_schema"]) if schema != FEATURE_SCHEMA_VERSION: raise ValueError( f"{path}: feature_schema {schema!r} != {FEATURE_SCHEMA_VERSION!r}" ) return { "features": z["features"].astype(np.float32), "labels": z["labels"].astype(np.int64), "generators": z["generators"], "sources": z["sources"], } def _rates(pred: np.ndarray, y: np.ndarray) -> dict: tp = int(((pred == 1) & (y == 1)).sum()) tn = int(((pred == 0) & (y == 0)).sum()) fp = int(((pred == 1) & (y == 0)).sum()) fn = int(((pred == 0) & (y == 1)).sum()) return { "n": int(len(y)), "accuracy": float((pred == y).mean()) if len(y) else 0.0, "confusion": {"tp": tp, "tn": tn, "fp": fp, "fn": fn}, "false_positive_rate": float(fp / (fp + tn)) if (fp + tn) else 0.0, "false_negative_rate": float(fn / (fn + tp)) if (fn + tp) else 0.0, } def main() -> None: parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--feat-dir", type=Path, required=True) parser.add_argument("--out", type=Path, required=True) parser.add_argument("--report-out", type=Path, default=None) parser.add_argument("--epochs", type=int, default=200) parser.add_argument("--batch-size", type=int, default=512) parser.add_argument("--lr", type=float, default=1e-2) parser.add_argument("--weight-decay", type=float, default=1e-4) parser.add_argument("--patience", type=int, default=15) parser.add_argument("--seed", type=int, default=0) parser.add_argument( "--hidden", type=int, default=0, help="Hidden units (0 = linear logistic regression).", ) parser.add_argument( "--eval-split", action="append", default=[], help="Extra split(s) to evaluate, e.g. --eval-split heldout.", ) args = parser.parse_args() import torch import torch.nn as nn torch.manual_seed(args.seed) np.random.seed(args.seed) train = _load_feat(args.feat_dir, "train") val = _load_feat(args.feat_dir, "val") test = _load_feat(args.feat_dir, "test") extra = {name: _load_feat(args.feat_dir, name) for name in args.eval_split} # Standardisation from train only (baked into the checkpoint). mean = train["features"].mean(axis=0) std = train["features"].std(axis=0) + 1e-8 def prep(split): x = (split["features"] - mean) / std return torch.from_numpy(x.astype(np.float32)), torch.from_numpy( split["labels"] ) x_tr, y_tr = prep(train) x_val, y_val = prep(val) n_feat = x_tr.shape[1] if args.hidden > 0: head = nn.Sequential( nn.Linear(n_feat, args.hidden), nn.ReLU(), nn.Linear(args.hidden, 2), ) else: head = nn.Linear(n_feat, 2) # Class weights to counter the authentic/ai imbalance in the splits. counts = np.bincount(train["labels"], minlength=2).astype(np.float64) w = counts.sum() / (2.0 * np.maximum(counts, 1.0)) class_weight = torch.tensor(w, dtype=torch.float32) loss_fn = nn.CrossEntropyLoss(weight=class_weight) opt = torch.optim.Adam( head.parameters(), lr=args.lr, weight_decay=args.weight_decay ) print( f"train={tuple(x_tr.shape)} val={tuple(x_val.shape)} " f"class_counts={counts.tolist()} hidden={args.hidden}" ) best_val = float("inf") best_state = None no_improve = 0 for epoch in range(1, args.epochs + 1): head.train() perm = torch.randperm(len(x_tr)) losses = [] for i in range(0, len(x_tr), args.batch_size): idx = perm[i : i + args.batch_size] opt.zero_grad() loss = loss_fn(head(x_tr[idx]), y_tr[idx]) loss.backward() opt.step() losses.append(loss.item()) head.eval() with torch.no_grad(): vloss = loss_fn(head(x_val), y_val).item() vacc = float((head(x_val).argmax(-1) == y_val).float().mean()) if epoch % 20 == 0 or epoch == 1: print( f" epoch {epoch:3d} train_loss={np.mean(losses):.4f} " f"val_loss={vloss:.4f} val_acc={vacc:.4f}" ) if vloss < best_val - 1e-4: best_val = vloss best_state = {k: v.clone() for k, v in head.state_dict().items()} no_improve = 0 else: no_improve += 1 if no_improve >= args.patience: print(f" early stop at epoch {epoch}") break if best_state is not None: head.load_state_dict(best_state) head.eval() def evaluate(split) -> dict: x = (split["features"] - mean) / std with torch.no_grad(): logits = head(torch.from_numpy(x.astype(np.float32))) pred = logits.argmax(-1).numpy() y = split["labels"] rep = _rates(pred, y) # Per-generator breakdown over AI rows. gens = np.asarray([str(g) for g in split["generators"]]) by_gen = {} for g in sorted(set(gens) - {""}): m = gens == g by_gen[g] = _rates(pred[m], y[m]) if by_gen: rep["by_generator"] = by_gen return rep report = { "feature_schema": FEATURE_SCHEMA_VERSION, "hidden": args.hidden, "splits": { "val": evaluate(val), "test": evaluate(test), **{name: evaluate(s) for name, s in extra.items()}, }, } for split_name in ("test", *args.eval_split): r = report["splits"][split_name] print( f" [{split_name}] acc={r['accuracy']:.4f} " f"fpr={r['false_positive_rate']:.4f} fnr={r['false_negative_rate']:.4f}" ) args.out.parent.mkdir(parents=True, exist_ok=True) torch.save( { "feature_schema": FEATURE_SCHEMA_VERSION, "feature_mean": mean.astype(np.float32), "feature_std": std.astype(np.float32), "state_dict": head.state_dict(), "hidden": args.hidden, "classes": _CLASSES, }, args.out, ) print(f"\nsaved frequency head to {args.out}") if args.report_out is not None: args.report_out.parent.mkdir(parents=True, exist_ok=True) with args.report_out.open("w", encoding="utf-8") as fh: json.dump(report, fh, indent=2, sort_keys=True) print(f"saved report to {args.report_out}") if __name__ == "__main__": main()