Phase 4: flow matching, min-SNR, EDM2 defaults, interval CFG, FID - #14
Closed
AshishKumar4 wants to merge 5 commits into
Closed
Phase 4: flow matching, min-SNR, EDM2 defaults, interval CFG, FID#14AshishKumar4 wants to merge 5 commits into
AshishKumar4 wants to merge 5 commits into
Conversation
Three changes to the training objective layer, sharing one contract. Flow matching slots into the existing scheduler/predictor seams rather than next to them: FlowMatchingScheduler is the linear path alpha = 1 - t, sigma = t with SD3's logit-normal timesteps and Flux's resolution shift, and FlowMatchPredictionTransform predicts the velocity u = eps - x0. Because the path is linear, the unmodified DDIM and Euler updates already are the flow Euler step; there is a test proving the identity exactly. min-SNR-gamma (Hang et al. 2023) is defined on the x_0 loss, so converting it into the space the trainer actually computes the loss in needs the parameterization. get_weights becomes a template over get_schedule_weights and asks the transform for that conversion, which also makes the schedule/transform pairing explicit instead of implicit. No second weighting system. EDMNoiseScheduler's P_mean/P_std were hardcoded to EDM1's -1.2/1.2; they are now constructor arguments defaulting to EDM2's -0.4/1.0. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Guidance costs diversity at high noise and buys nothing at low noise (Kynkaanniemi et al. 2024). guidance_start/guidance_stop bound the fraction of the trajectory over which it applies; outside the interval the scale drops to 1, which is exactly the plain conditional prediction. The default interval is the full range, so existing samplers are unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The jax-fid port in metrics/inception.py has been sitting unused. Wire it into an EvaluationMetric the same way metrics/images.py builds the CLIP ones, sharing a single cached copy of the weights. Per-batch FID is noisy and only meaningful as a trend across checkpoints, which the docstring says outright. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Latents are normalized as (z - latent_shift) * latent_scale and inverted on decode, the SD3 convention, so a caller can hand in a dataset's own mean and 1/std instead of relying on the VAE's single fixed constant. The SD VAE's scaling_factor now rides on that same seam rather than being applied inside its jitted frame functions, so there is one normalization path and the defaults reproduce the previous behavior exactly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
--noise_schedule gains flow/flow_matching, plus --min_snr_gamma, --flow_shift, and 'fid' as a validation metric. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 27, 2026
Owner
Author
|
Landed on |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fifth PR of the revamp series — stacked on #13.
Flow matching
FlowMatchingScheduler(subclassing the existingContinuousNoiseScheduler) +FlowMatchPredictionTransformslot into the existing scheduler/predictor seams — no new machinery. Linear pathx_t = (1−t)x₀ + tε, logit-normal timesteps (SD3'sm=0, s=1), SD3 resolution shiftt' = st/(1+(s−1)t)with Flux-stylecompute_resolution_shift, andc_in = 1.The existing samplers work unmodified — on a linear path the DDIM/Euler update is the FM Euler step, and
test_sampler_step_is_the_flow_euler_stepproves that algebraically to <1e-5 for both.A toy 2-D test proves it actually learns, not merely runs: 1500 CPU steps on a two-mode Gaussian mixture, then 2048 samples drawn through the real
EulerSampler, asserting mode balance and per-mode statistics (untrained σ≈0.65 with wrong means → trained means within 0.006, σ within 0.008). Runs in ~4s.min-SNR-γ — integrated, not bolted on
The subtle part:
min(SNR,γ)is defined on the x₀ loss, but the loss is computed in each transform's target space. Rather than add a parallel weighting system,get_weightsbecame a template overget_schedule_weights, and each transform gainedtarget_error_scale(snr)= ‖target error‖²/‖x₀ error‖². This reproduces the paper formulas exactly (eps →min/SNR, v →min/(SNR+1)) and falls out correctly for Karras and FM.I verified all four derivations by hand and numerically: FM's
(1+√SNR)²is exactly1/t²; Karras's1/σ_d² + SNRis exactly1/c_out². And a good independent signal that the abstraction is right:KarrasPredictionTransform.target_error_scaleturns out to equal EDM's existing weighting, i.e. EDM is x₀-space weight 1.EDM's audited-correct weighting is bit-identical after the refactor (max relative error
0.0) when min-SNR is off.Also
P_mean/P_stdexposed and defaulted to EDM2's−0.4/1.0(was EDM1's−1.2/1.2hardcoded).guidance_start/guidance_stopstep-fractions onDiffusionSampler; default is full range, so no behavior change.metrics/inception.pyhad been sitting completely unused — now wired as anEvaluationMetric(--val_metrics fid).AutoEncoderseam, with the VAE'sscaling_factorfolded into it rather than stacked on top. Defaults reproduce prior behavior; verified by the network-marked VAE roundtrip test.Verification
112 passed, 4 deselected (from 73), independently re-run by me, not just reported. Network-marked tests (FID weights, SD VAE roundtrip) pass too.
Known scope note: no helper to compute a dataset's latent mean/std — the seam takes them as input; say the word if you want the measurement pass.
🤖 Generated with Claude Code