Skip to content

Flow-matching fidelity: loss functions, SNR weighting, channel balancing, EMA, telemetry - #68

Open
scragnog wants to merge 15 commits into
koda-dernet:mainfrom
scragnog:pr/fidelity
Open

Flow-matching fidelity: loss functions, SNR weighting, channel balancing, EMA, telemetry#68
scragnog wants to merge 15 commits into
koda-dernet:mainfrom
scragnog:pr/fidelity

Conversation

@scragnog

@scragnog scragnog commented Jul 29, 2026

Copy link
Copy Markdown

Flow-matching fidelity: loss functions, SNR weighting, channel balancing, EMA, telemetry

Part 2 of the training-core series — stacked on #67; until that merges this diff includes its commits.

A body of work on training fidelity for the rectified-flow objective, developed and A/B-listened over several weeks of LoRA runs:

Loss & weighting

  • loss_fn: mse (default, unchanged), huber (+huber_delta), pseudo_huber, and x0 prediction-space loss with dual-metric telemetry (both spaces logged whichever trains)
  • loss_weighting: adds flow_snr (SNR weighting derived for rectified flow) alongside the existing min_snr; snr_gamma applies to both
  • t_bias, latent_noise, per-channel balancing with optional VAE channel priors, and a dynamic_channel_balance mode that rebalances from running statistics
  • legacy_loss escape hatch reproduces the old exact objective bit-for-bit

EMA overhaul

  • ema_start_step replaces the linear warmup ramp: EMA simply starts late (cleaner semantics, no shadow-drift during warmup)
  • apply() guarded as a no-op before start step — prevents a stale zero-initialized shadow overwriting live weights on early checkpoint/preview paths
  • restore-without-apply downgraded to debug (harmless before start)

Telemetry

  • fidelity metrics, EMA effective-decay, timestep-distribution histograms, SNR fix
  • GUI: loss-function dropdown with help text and chart tags

Passthrough fixes — wizard/preset/GUI-CLI now correctly propagate all of the above (several fields previously fell back to defaults silently on non-CLI paths, including --loss-weighting none which trained with flow_snr).

Defaults are unchanged: a config that doesn't opt into anything trains identically to before (mse + existing behavior).

🤖 Generated with Claude Code

Rob Work PC and others added 15 commits July 29, 2026 11:34
Rob Work PC @koda-dernet @claude
feat: dynamic genre conditioning — encode both variants, pick per step
dbee15d 
Preprocessing now encodes BOTH the caption and the genre-tag text
variant into each .pt (backward-compatible optional keys); training
selects caption vs genre per sample per step according to genre_ratio
(auto-detected from preprocess_meta.json when unset) instead of the old
static per-sample choice at preprocess time. Old tensors keep working.

Co-authored-by: Gabriel <gadna166@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Rob Work PC @koda-dernet @claude
feat: conditioning monitor — per-batch caption/genre telemetry in the…
85ef97a 
… GUI

conditioning_type flows from the dataset through collate into
ProgressWriter events; the training monitor shows live caption-vs-genre
usage. Also ports target-loss scale/EMA telemetry to the Fabric trainer
path (previously basic-loop only).

Co-authored-by: Gabriel <gadna166@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Rob Work PC @claude
feat: flow-matching fidelity math — Huber loss, flow_snr weighting, c…
f9afd4d 
…hannel balancing, latent noise, t_bias

(from 5e400f6; scripts/ debug artifacts excluded)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@koda-dernet Rob Work PC
fix: split fidelity controls into 3 form-rows for proper layout
4e00dcf
@koda-dernet Rob Work PC
fix: default loss_fn back to mse (correct for flow matching) + EMA wa…
547dd87 
…rmup ramp

- MSE guarantees convergence to conditional expectation (required by flow matching)
- Huber converges to median, under-penalizes extreme flow vectors -> clipped highend
- Huber still available as option, just not default
- EMA warmup: effective_decay = min(target, (1+step)/(10+step)) so shadow catches up fast early
@koda-dernet Rob Work PC
feat: configurable EMA warmup steps (linear ramp from 0 to target decay)
00cb7d1 
- ema_warmup_steps (default 2000): prevents EMA shadow from lagging early training
- At step 0 decay=0 (shadow=current weights), linearly ramps to target decay
- Wired through: defaults, configs, config_factory, CLI, wizard UI, frontend
- GUI: warmup field auto-hides when EMA decay=0
@koda-dernet Rob Work PC
refactor: EMA delayed start instead of linear ramp (ema_start_step)
240bc2b 
- No tracking until start_step, then snapshot current weights as shadow
- Mathematically clean: shadow is always a proper exponential average at full target decay
- Linear ramp was creating a Frankenstein blend of different averaging rates
- Renamed ema_warmup_steps -> ema_start_step across all 16 touchpoints
- Default 2000 steps: model learns basic structure before EMA locks in
@koda-dernet Rob Work PC
fix: guard EMA apply() as no-op before start_step to prevent stale sh…
d989c30 
…adow overwrite
@koda-dernet Rob Work PC
feat: full training telemetry for fidelity metrics, EMA, and timestep…
f53d39b 
… distribution

- training_step now stashes auxiliary metrics: raw/weighted loss, SNR weight stats,
  timestep mean, per-channel loss breakdown (max/min/ratio)
- drain_step_metrics() returns and clears the dict at log_every intervals
- Both Fabric and basic training loops log to TensorBoard: fidelity/*, ema/*
- Fixed timestep histogram to fire for flow_snr too (was min_snr only)
- Pipe fidelity dict + ema_active through progress writer for real-time GUI
- Frontend: unpack fidelity dict from progress messages into _tbScalars
- Added curated labels/colors for all new tags in mini-chart registry
@koda-dernet Rob Work PC
feat: dynamic channel rebalancing + SNR telemetry fix + EMA decay metric
771ff2a 
- dynamic_channel_balance (opt-in toggle): blends 50/50 static VAE prior
  with running EMA of per-channel loss, so struggling channels get more gradient
- Wired through all config touchpoints: defaults, configs, factory, CLI, GUI
- GUI: checkbox auto-hides when channel_balance is off or legacy mode
- Fix: SNR weight charts were always 1.0 (logged post-normalization);
  now logs pre-normalization values + spread ratio
- New ema/effective_decay metric (0.0 when inactive, target decay when active)
- Added curated mini-chart labels for snr_weight_spread, ch_dynamic_blend_range,
  ema/effective_decay
@koda-dernet Rob Work PC
feat: add x0-prediction loss and pseudo-huber loss modes with dual-me…
8daf308 
…tric telemetry

- Add 3 new loss_fn options: pseudo_huber, x0_mse, x0_pseudo_huber
- x0_mse implements paper's intended formulation (implicit t² weighting)
- pseudo_huber: δ²(√(1+(e/δ)²) - 1) smooth Huber approximation
- Log fidelity/arrangement_loss (t>0.5) and fidelity/detail_loss (t≤0.5)
- Update config validation to accept all 5 loss_fn values
@koda-dernet Rob Work PC
feat(cli/wizard): surface new loss_fn options in CLI args and wizard …
27330b0 
…prompts

- Expand --loss-fn choices to include pseudo_huber, x0_mse, x0_pseudo_huber
- Wizard shows delta prompt for all delta-using modes (huber, pseudo_huber, x0_pseudo_huber)
- Update review summary label to Huber/Pseudo-Huber delta
@koda-dernet Rob Work PC
feat(frontend): add loss function dropdown options, help text, and ch…
2eb4a51 
…art tags

- Add pseudo_huber, x0_mse, x0_pseudo_huber to loss function dropdown
- Update help panel with descriptions for all 5 modes
- Show huber delta input for all delta-using loss modes
- Add Arrangement Loss and Detail Loss to chart tag metadata
Rob Work PC
Downgrade EMA restore-without-apply to debug (harmless before start_s…
83a21a5 
…tep)
Rob Work PC @claude
fix: complete wizard/preset/GUI-CLI passthrough for fidelity fields
15fa0f7 
- build_train_namespace passes dynamic_channel_balance through (was
  silently reset to default for wizard and preset runs)
- presets recognise dynamic_channel_balance
- GUI CLI builder emits --loss-weighting unconditionally (backend
  default is flow_snr, so omitting 'none' trained with flow_snr),
  plus --loss-fn/--huber-delta/--t-bias/--latent-noise/
  --channel-balance/--dynamic-channel-balance/--vae-channel-prior/
  --legacy-loss/--ema-start-step, and snr-gamma for flow_snr too
- preset loader maps dynamic_channel_balance into the form

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

coderabbitai Bot commented Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Warning

Review limit reached

@scragnog, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 59 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 32d510fb-9b2f-47ff-aae4-ff4cb207fdf6

📥 Commits

Reviewing files that changed from the base of the PR and between ecd13bd and 15fa0f7.

📒 Files selected for processing (25)
  • frontend/index.html
  • frontend/js/api-cli.js
  • frontend/js/defaults.json
  • frontend/js/training.js
  • frontend/js/workspace-behaviors.js
  • frontend/js/workspace-config.js
  • frontend/js/workspace-setup.js
  • sidestep_engine/cli/args.py
  • sidestep_engine/core/adaptive_timestep.py
  • sidestep_engine/core/config_factory.py
  • sidestep_engine/core/configs.py
  • sidestep_engine/core/ema.py
  • sidestep_engine/core/lora_module.py
  • sidestep_engine/core/timestep_sampling.py
  • sidestep_engine/core/trainer.py
  • sidestep_engine/core/trainer_loop.py
  • sidestep_engine/data/preprocess.py
  • sidestep_engine/training_defaults.py
  • sidestep_engine/ui/config_panel.py
  • sidestep_engine/ui/flows/common.py
  • sidestep_engine/ui/flows/review_summary.py
  • sidestep_engine/ui/flows/train_steps_levers.py
  • sidestep_engine/ui/flows/train_steps_training.py
  • sidestep_engine/ui/presets.py
  • sidestep_engine/vendor/data_module.py

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

This was referenced Jul 29, 2026
Open
Open
@scragnog
scragnog marked this pull request as ready for review July 29, 2026 10:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@scragnog @koda-dernet