feat: add WGAN + AR stop-token config variant
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m23s
CI / Format (ruff format) (push) Successful in 1m23s
CI / Type check (ty) (push) Successful in 1m24s
CI / Tests (push) Successful in 3m32s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 8s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m23s
CI / Format (ruff format) (push) Successful in 1m23s
CI / Type check (ty) (push) Successful in 1m24s
CI / Tests (push) Successful in 3m32s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 8s
Branches off configs/baseline.toml with both stages on WGAN-GP and stage-2 n_sec.mode = stop_token, 30 epochs — combines two unbenchmarked roadmap axes (post-v0.3.0 WGAN, and the AR stop-token multiplicity mode) into one variant that stays a single edit away from baseline for attribution. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KzPghrmFcAJYrWUvHApY9N
This commit is contained in:
@@ -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"
|
||||
Reference in New Issue
Block a user