Add KL divergence to marginal validation and hook it into the training loop

validate_marginals now estimates a per-dimension KL(real || generated) via
a shared histogram, alongside the existing mean/std comparison, so
distribution-shape drift shows up even when the first two moments match.

Wire it into giant/train.py: every validate_every epochs (default 10, 0
disables), the training loop runs validate_marginals against val_loader and
prints the table. validate_every flows through DEFAULT_CONFIG/config.toml
and is exposed as --validate-every on both giant train and scripts/train.py.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-06-18 13:45:12 +02:00
co-authored by Claude Sonnet 4.6
parent 43634ef77a
commit 893d91e749
6 changed files with 47 additions and 6 deletions
+2
View File
@@ -60,6 +60,7 @@ def train(
emb_dim: Annotated[Optional[int], typer.Option()] = None,
val_fraction: Annotated[Optional[float], typer.Option()] = None,
seed: Annotated[Optional[int], typer.Option(help="Random seed for reproducibility")] = None,
validate_every: Annotated[Optional[int], typer.Option(help="Run marginal+KL validation every N epochs (0 disables)")] = None,
shuffle_buffer: Annotated[int, typer.Option(help="Rows held in RAM per worker for shuffling")] = 65536,
out: Annotated[Optional[Path], typer.Option(help="Checkpoint dir (default: auto from hyperparams)")] = None,
device: Annotated[Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")] = None,
@@ -71,6 +72,7 @@ def train(
"mode": mode.value if mode is not None else None,
"epochs": epochs, "batch_size": batch_size, "lr": lr,
"val_fraction": val_fraction, "num_workers": num_workers, "seed": seed,
"validate_every": validate_every,
}.items() if v is not None}
cli_model = {k: v for k, v in {
"hidden_dim": hidden_dim, "n_blocks": n_blocks, "emb_dim": emb_dim,