diff --git a/configs/wgan_ar_stoptoken.toml b/configs/wgan_ar_stoptoken.toml new file mode 100644 index 0000000..1ecb432 --- /dev/null +++ b/configs/wgan_ar_stoptoken.toml @@ -0,0 +1,122 @@ +# GIANT WGAN-GP + AR stop-token variant of configs/baseline.toml. +# +# Two roadmap axes, combined into one run: WGAN-GP generators for both +# stages (unbenchmarked since the 2026-08-03 pre-v0.3.0 failure, which was +# secondary-species mode collapse — the failure v0.3.0's AR/categorical +# pivot exists to fix) and the AR stop-token multiplicity mode +# (stage2_model.n_sec.mode = "stop_token", never benchmarked at all). +# Everything else is byte-identical to baseline.toml so a rollout compared +# against baseline's analysis_341dfb14 is attributable to these two axes +# alone: conditioning (physical/physical), hidden_dim 512 / n_res_blocks 6 / +# dropout 0.0 per stage, k_max 15, history "markov", teacher_forcing +# "always", particle_type.target "onehot" (n_classes 32, other_policy +# "sample"), lr 3e-4, warmup_epochs 3, weight_decay 0.01, ema_decay 0.9999, +# val_fraction 0.1, num_workers 4, seed 0, validate_steps 10, W&B on. +# +# No [stage1_model.wgan] / [stage2_model.wgan] block: the dataclass defaults +# (noise_dim 64, n_critic 5, gp_weight 10.0, critic_lr 0.0 = inherit +# train.lr, critic_hidden_dim/critic_n_res_blocks 0 = inherit the stage's +# 512/6, stage 2's gumbel_tau_start/_end 1.0/0.1) are what the earlier WGAN +# runs used — writing them out would add keys that don't vary. +# +# particle_type.class_weighting stays "none" (the default): config.py's +# validate_config rejects any other value under stage2_model.generator = +# "wgan", since that path feeds the type slice to the critic via a +# straight-through Gumbel relaxation instead of a weighted cross-entropy. +# +# Prior WGAN writeup (pre-v0.3.0, describes the failure this run re-tests): +# /home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md + +[meta] +# REQUIRED. Without it config.migrate_config reads this file as v0.2 and +# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot", +# particle_type.target = "physical" and the v0.2 default sizes, while still +# passing validate_config. +config_version = 3 + +[conditioning] +# Physical-property MLPs rather than learned vocab embeddings: computable for +# any PDG code / material, which is what the held-out-species and +# held-out-material generalization comparisons need. +out_dim = 128 +share_stages = false + +# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was +# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder +# identical to baseline.toml. +[conditioning.particle] +type = "physical" +emb_dim = 16 +n_layers = 2 + +[conditioning.material] +type = "physical" +emb_dim = 16 +n_layers = 2 + +[stage1_model] +generator = "wgan" +hidden_dim = 512 +n_res_blocks = 6 +dropout = 0.0 + +[stage2_model] +# Autoregressive in descending-energy order, as baseline.toml — this variant +# only swaps the generator (flow -> wgan) and the multiplicity mode +# (head -> stop_token), not the decoder shape. +decoder = "autoregressive" +generator = "wgan" +hidden_dim = 512 +n_res_blocks = 6 +dropout = 0.0 +k_max = 15 + +[stage2_model.autoregressive] +history = "markov" +teacher_forcing = "always" + +[stage2_model.n_sec] +# EOS-style per-slot stop head on the AR secondary decoder, replacing the +# n_sec classifier entirely (mutually exclusive — see NSecConfig's +# docstring in giant/config.py). Requires decoder = "autoregressive" and +# owner = "stage2" (both already true above/by default); validate_config +# enforces this. +mode = "stop_token" + +[stage2_model.particle_type] +target = "onehot" +# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the +# "other" bucket keeps essentially all real secondary species out of "other" +# without making the head expensive. +n_classes = 32 +other_policy = "sample" + +[train] +epochs = 30 +# Halved from baseline's 36864. That figure came from a measured linear fit +# of the *flow-AR* training step (peak reserved MiB = 0.9736 * batch_size + +# 115); WGAN invalidates it twice over — each stage gains a critic that by +# default inherits the stage's own 512/6 body, and gradient_penalty +# (giant/model/wgan.py, forced fp32 internally) runs a double-backward every +# batch. 18432 is a conservative choice pending a real memory measurement on +# this exact config, not a re-derived fit. Throughput is already flat above +# bs~4096 on the 4070, so this costs occupancy on the L40S, not step +# efficiency. +batch_size = 18432 +lr = 3e-4 +warmup_epochs = 3 +weight_decay = 0.01 +ema_decay = 0.9999 +val_fraction = 0.1 +num_workers = 4 +seed = 0 +# Tightened from baseline's 10: WGANStageTrainer.supports_val_loss = False, +# and with both stages adversarial there is no per-epoch val loss at all, so +# validate_every's marginal-KL pass (giant/training/trainers.py's +# val_objective) is the only comparable-across-epochs best-checkpoint +# selection signal available. 5 gives 6 evaluations over 30 epochs instead +# of baseline's 3, at ~6x5000s of extra walltime. +validate_every = 5 +validate_steps = 10 +wandb = true +wandb_project = "giant"