Add opt-in straight-through Gumbel-softmax combine weights to MoE router #26

Merged

lars merged 4 commits from feat/router-gumbel-softmax into master 2026-07-30 17:05:09 +02:00

Conversation 0 Commits 4 Files Changed 10

+550 -28

lars commented 2026-07-30 16:30:49 +02:00

Owner

Copy Link

Copy Source

Trains the routed trunk's forward combination as a hard one-hot sample
(matching eval-time top-1 dispatch exactly) while keeping a smooth gradient
on the backward pass, targeting the train/eval mismatch identified as a
likely contributor to experts overlapping instead of partitioning in the
first energy-router rollout benchmark. Off by default (model.router.gumbel);
existing routed configs/checkpoints are unaffected.

Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com

Trains the routed trunk's forward combination as a hard one-hot sample (matching eval-time top-1 dispatch exactly) while keeping a smooth gradient on the backward pass, targeting the train/eval mismatch identified as a likely contributor to experts overlapping instead of partitioning in the first energy-router rollout benchmark. Off by default (model.router.gumbel); existing routed configs/checkpoints are unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

lars added 1 commit 2026-07-30 16:30:49 +02:00

Add opt-in straight-through Gumbel-softmax combine weights to MoE router ... b51eafcfa5

Trains the routed trunk's forward combination as a hard one-hot sample
(matching eval-time top-1 dispatch exactly) while keeping a smooth gradient
on the backward pass, targeting the train/eval mismatch identified as a
likely contributor to experts overlapping instead of partitioning in the
first energy-router rollout benchmark. Off by default (model.router.gumbel);
existing routed configs/checkpoints are unaffected.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

lars added 1 commit 2026-07-30 16:40:09 +02:00

Add gumbel router configs sweeping learnable-knob combinations ... af2ee7c7ce

Extends the 10-expert EnergyRouter + physical-conditioning benchmark
config with the new opt-in gumbel combine weights, isolating the
learn_centers/learn_temperature axis: none, centers only, and
centers+temperature.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

lars added 1 commit 2026-07-30 16:43:01 +02:00

Add gumbel/learn_centers/learn_width/learn_temperature to out-dir naming ... 5eec4c250a

Extends default_out_dir_name's non-default-field convention to the
router's new gumbel combine-weight flag and its learnable-knob toggles,
so gumbel sweep configs (learn_centers on/off, learn_width,
learn_temperature) resolve to distinguishable checkpoint directory
names instead of colliding on the same r-<type><n> token.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

lars added 1 commit 2026-07-30 16:46:24 +02:00

Scope wandb run config to only-active hyperparameters ... 8065df896e

Router-only knobs (lambda_balance/lambda_proc/lambda_entropy/
gumbel_tau_start/_end) and WGAN-only knobs (n_critic/gp_weight) were
being logged to wandb's top-level run config unconditionally, even for
runs where routing or WGAN mode is off, implying hyperparameters from
an inactive code path. Extract _wandb_run_config and only include each
group when its gate is actually true (router.enabled / mode=="wgan");
the full model_config (with its router sub-dict) is still always
logged in full, so no information is lost.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

lars merged commit 22fdca7697 into master 2026-07-30 17:05:09 +02:00

lars deleted branch feat/router-gumbel-softmax 2026-07-30 17:05:09 +02:00

Sign in to join this conversation.

Reviewers

No Reviewers

Labels

Clear labels

architecture

experiment

investigation

performance

phase2

validation

No Label

Milestone

No items

No Milestone

Projects

Clear projects

No project

Assignees

Clear assignees

lars (Lars Bogner)

No Assignees

1 Participants

Notifications

Subscribe

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: lars/giant#26