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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user