[training scripts] --logit_mean biases timesteps in the opposite direction to the SD3 paper (positive m oversamples low noise)

Describe the bug

The SD3/FLUX-style training scripts sample timesteps with compute_density_for_timestep_sampling(weighting_scheme="logit_normal", logit_mean=m, logit_std=s), which returns u = sigmoid(N(m, s)). They then index the scheduler's timestep table with indices = (u * num_train_timesteps).long(). FlowMatchEulerDiscreteScheduler's table is ordered noise-first (timesteps[0] = 1000, sigmas[0] = 1), so the realised noise level is sigma ≈ shift(1 - u), i.e. sigma follows a (shifted) logit-normal(-m, s) distribution.

The SD3 paper (Esser et al. 2024, Sec. 3.1, Eq. 19), which the helper's docstring and README_sd3.md refer to, samples the noise level t of z_t = (1 - t) x_0 + t eps (the scripts' own comment: zt = (1 - texp) * x + texp * z1) from logit-normal(m, s), and says m biases training "towards either data p0 (negative m) or noise p1 (positive m)". In the scripts the sign of --logit_mean is therefore reversed: --logit_mean 1.0 --logit_std 0.6 (the paper's lognorm(1.00, 0.60)) concentrates training on low noise levels (FLUX: E[sigma] 0.284 instead of 0.717; P(sigma > 0.5) 0.048 instead of 0.952).

The defaults are not affected: logit_mean=0 and the mode / uniform schemes are symmetric. Only users who set --logit_mean get the opposite of what the paper describes.

Affected: the 28 scripts that use this indexing line (list below). Not affected:

train_dreambooth_lora_lumina2.py, which also reverses the interpolation ((1 - sigma) * noise + sigma * x0);
train_dreambooth_lora_ideogram4.py, which uses the sampled value directly as sigma, i.e. the paper's direction.

Reproduction

This uses the real compute_density_for_timestep_sampling and FlowMatchEulerDiscreteScheduler with the FLUX.1 and SD3 scheduler configs, and the scripts' indexing:

# The SD3/FLUX training scripts sample timesteps like this (e.g. examples/dreambooth/train_dreambooth_sd3.py):
#   u = compute_density_for_timestep_sampling(weighting_scheme="logit_normal", logit_mean=m, logit_std=s, ...)
#   indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
#   timesteps = noise_scheduler_copy.timesteps[indices]          -> sigmas via get_sigmas(timesteps)
#   noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise
# SD3 paper (Esser et al. 2024, Sec. 3.1, Eq. 19): t ~ logit-normal(m, s) is the noise level of
# z_t = (1 - t) x_0 + t eps, and m > 0 biases training towards noise. Here sigma plays the role of t.
import torch
import diffusers
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.training_utils import compute_density_for_timestep_sampling

print(f"diffusers {diffusers.__version__}, torch {torch.__version__}")
m, s, n = 1.0, 0.6, 200_000
configs = {  # scheduler configs of the FLUX.1 and SD3 checkpoints (scheduler/scheduler_config.json)
    "FLUX.1": dict(shift=3.0, use_dynamic_shifting=True, base_shift=0.5, max_shift=1.15,
                   base_image_seq_len=256, max_image_seq_len=4096),
    "SD3": dict(shift=3.0),
}
for name, cfg in configs.items():
    sched = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, **cfg)
    torch.manual_seed(0)
    u = compute_density_for_timestep_sampling(weighting_scheme="logit_normal", batch_size=n,
                                              logit_mean=m, logit_std=s, mode_scale=1.29)
    indices = (u * sched.config.num_train_timesteps).long()  # as in the scripts
    timesteps = sched.timesteps[indices]
    # get_sigmas() looks up each timestep's position in the same table, i.e. `indices` (checked on a subset)
    assert all((sched.timesteps == t).nonzero().item() == i for t, i in zip(timesteps[:500], indices[:500]))
    sigma = sched.sigmas[indices]
    # Paper direction: the noise level itself is logit-normal(m, s) (then the same table/shift)
    sigma_paper = sched.sigmas[((1 - u) * 1000).long().clamp(max=999)]
    print(f"{name:6s} logit_mean={m:+.1f}, logit_std={s}: E[sigma] scripts {sigma.mean():.3f}, paper {sigma_paper.mean():.3f}"
          f" | P(sigma > 0.5) scripts {(sigma > 0.5).float().mean():.3f}, paper {(sigma_paper > 0.5).float().mean():.3f}")
    assert sigma.mean() < 0.55 < sigma_paper.mean()  # m > 0 should bias towards noise, the scripts bias towards data
print("BUG REPRODUCED: the sign of --logit_mean is reversed relative to the SD3 paper")

Logs

diffusers 0.42.1, torch 2.15.0+cpu
FLUX.1 logit_mean=+1.0, logit_std=0.6: E[sigma] scripts 0.284, paper 0.717 | P(sigma > 0.5) scripts 0.048, paper 0.952
SD3    logit_mean=+1.0, logit_std=0.6: E[sigma] scripts 0.524, paper 0.877 | P(sigma > 0.5) scripts 0.567, paper 1.000
BUG REPRODUCED: the sign of --logit_mean is reversed relative to the SD3 paper

Expected: with logit_mean=+1.0 the noise level should be biased towards noise, as in the paper (E[sigma] 0.717 for FLUX.1, 0.877 for SD3).

Where: training_utils.py#L366-L390 (v0.42.1, 2a14339) returns u, which correctly plays the role of the paper's t. The table is noise-first (scheduling_flow_match_euler_discrete.py#L127-L136). The scripts index it with u, e.g. train_dreambooth_lora_sd3.py#L1784-L1797:

                u = compute_density_for_timestep_sampling(
                    weighting_scheme=args.weighting_scheme,
                    batch_size=bsz,
                    logit_mean=args.logit_mean,
                    logit_std=args.logit_std,
                    mode_scale=args.mode_scale,
                )
                indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
                timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)

                # Add noise according to flow matching.
                # zt = (1 - texp) * x + texp * z1
                sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
                noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise

The indexing line is identical, at the same line number, in all 28 scripts at v0.42.1 and on main (86ef420).

The 28 affected scripts (indexing line)

Proposed fix. The same one-line change in each of the 28 scripts. The shared helper cannot be flipped, because the ideogram4 script uses its output directly as sigma:

-                indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
+                indices = ((1 - u) * noise_scheduler_copy.config.num_train_timesteps).long().clamp(
+                    max=noise_scheduler_copy.config.num_train_timesteps - 1
+                )

The mode and uniform schemes, and logit_normal with logit_mean=0, are symmetric under u -> 1 - u, so their distributions do not change.

Verifying the fix. fix_check.py applies the fixed indexing with the same real helper and schedulers. It compares the realised sigmas with the paper's distribution sigma = shift(sigmoid(x)), x ~ N(m, s) (quantiles at 5/25/50/75/95%), and checks that the default settings keep the same distribution:

released = sched.sigmas[(u * N).long()]
fixed = sched.sigmas[((1 - u) * N).long().clamp(max=N - 1)]
t = torch.sigmoid(torch.randn(n) * s + m)            # paper: the noise level t ~ logit-normal(m, s)
paper = shift * t / (1 + (shift - 1) * t)            # shift = 1 for FLUX.1 (dynamic shifting), 3 for SD3
assert (fixed.quantile(qs) - paper.quantile(qs)).abs().max() < 5e-3

diffusers 0.42.1, torch 2.15.0+cpu
FLUX.1 logit_normal m=+1.0 s=0.6: E[sigma] released 0.284, fixed 0.717, paper 0.716 | max quantile gap fixed vs paper 0.0012
FLUX.1 logit_normal m=-0.5 s=1.0: E[sigma] released 0.603, fixed 0.398, paper 0.398 | max quantile gap fixed vs paper 0.0007
FLUX.1 logit_normal m=+0.0 s=1.0: E[sigma] released 0.501, fixed 0.500, paper 0.500 | max quantile gap fixed vs paper 0.0007 | max quantile gap fixed vs released 0.0020
FLUX.1 mode         m=+0.0 s=1.0: E[sigma] released 0.501, fixed 0.500 | max quantile gap fixed vs released 0.0020
FLUX.1 none         m=+0.0 s=1.0: E[sigma] released 0.500, fixed 0.501 | max quantile gap fixed vs released 0.0030
SD3    logit_normal m=+1.0 s=0.6: E[sigma] released 0.524, fixed 0.877, paper 0.877 | max quantile gap fixed vs paper 0.0009
SD3    logit_normal m=-0.5 s=1.0: E[sigma] released 0.794, fixed 0.621, paper 0.621 | max quantile gap fixed vs paper 0.0010
SD3    logit_normal m=+0.0 s=1.0: E[sigma] released 0.715, fixed 0.714, paper 0.714 | max quantile gap fixed vs paper 0.0011 | max quantile gap fixed vs released 0.0021
SD3    mode         m=+0.0 s=1.0: E[sigma] released 0.710, fixed 0.710 | max quantile gap fixed vs released 0.0020
SD3    none         m=+0.0 s=1.0: E[sigma] released 0.676, fixed 0.677 | max quantile gap fixed vs released 0.0027
FIX OK: the realised noise levels follow the SD3 paper's logit-normal(m, s); defaults unchanged

System Info

🤗 Diffusers version: 0.42.1
Platform: Linux x86_64
Python version: 3.13.13
PyTorch version: 2.15.0+cpu (CPU only)
Transformers version: 5.22.1
Accelerate version: 1.15.0

Who can help?

No response

I would be happy to open a PR with this fix for the 28 scripts.

@Andreas-Hum

I reproduced this on current main (diffusers 0.42.0.dev0, torch 2.15.0) with the script above. I got the same numbers: FLUX.1 E[sigma] 0.284 vs 0.717, SD3 0.524 vs 0.877. The indexing line now appears in 29 scripts under examples/, one more than the list.

I'd like to fix this in one PR across all of them, if that's welcome and @Andreas-Hum isn't already on it. One question first, since it changes what --logit_mean does for anyone who sets it:

Fix the indexing (indices = ((1 - u) * num_train_timesteps).long().clamp(max=num_train_timesteps - 1)), so --logit_mean matches the paper and the helper docstring, plus a short note in the affected READMEs about the sign change. Or
Keep current behaviour and only fix the docs/help text so they describe the reversed sign.

I'd lean towards 1. logit_mean=0, mode and uniform are symmetric, so default runs are unaffected. I'll use AI assistance (Claude Code) and run the self-review skill. I'll wait for your call before opening anything.

Thanks for reproducing it. I already have the fix, and it's option 1. It's on a local branch that changes the indexing line in all affected scripts:

indices = ((1 - u) * num_train_timesteps).long().clamp(max=num_train_timesteps - 1)

I'll open the PR once a maintainer acknowledges the issue, as the contribution guide asks.

One correction on the count: the 29th match on main is train_dreambooth_lora_lumina2.py, and that script is already correct, so it must not be changed. It uses the same indexing, but it also swaps the interpolation (noisy = (1 - sigma) * noise + sigma * x0). That means its sigma is the data fraction, and its noise level is 1 - sigma ≈ u, which already follows the paper. The two reversals cancel out, which is why the issue leaves it off the list.

Why the sign flips

compute_density_for_timestep_sampling returns u = sigmoid(ξ) with ξ ~ N(m, s²).
The FlowMatch table runs from noise to data. Before the shift, entry i has sigma_i = 1 - i/N. So indices = floor(u * N) gives sigma ≈ 1 - u (and then the shift shift * sigma / (1 + (shift - 1) * sigma) is applied).
1 - sigmoid(ξ) = sigmoid(-ξ), and -ξ ~ N(-m, s²). So the scripts' noise level is logit-normal(-m, s).
In the SD3 paper (Esser et al. 2024, Sec. 3.1, Eq. 19), t ~ logit-normal(m, s) is the noise level in z_t = (1 - t) x_0 + t ε, and a positive m biases training towards noise. The scripts use the same interpolation ((1 - sigma) * x0 + sigma * noise), so sigma should play the role of t. With the current indexing, --logit_mean therefore has the opposite sign.

The fix

Index with 1 - u, so that sigma ≈ u ~ logit-normal(m, s), as in the paper. The clamp is needed because for very small u, 1 - u rounds to exactly 1.0 in float32, which would give index N, one past the end of the table. Compared with the exact mirror N - 1 - floor(u * N), the result differs by at most one table step.

Why defaults don't change

logit_mean = 0: ξ and -ξ have the same distribution, so u and 1 - u do too.
mode: with v uniform, the scheme's f(v) = 1 - v - c(cos²(πv/2) - 1 + v) satisfies f(1 - v) = 1 - f(v), and 1 - v is also uniform.
uniform: symmetric for the same reason.

What I checked (real compute_density_for_timestep_sampling and FlowMatchEulerDiscreteScheduler, FLUX.1 and SD3 scheduler configs, 200k samples):

SD3    logit_normal m=+1.0 s=0.6: E[sigma] released 0.524, fixed 0.877, paper 0.877 | max quantile gap fixed vs paper 0.0009
SD3    logit_normal m=-0.5 s=1.0: E[sigma] released 0.794, fixed 0.621, paper 0.621 | max quantile gap fixed vs paper 0.0010
SD3    logit_normal m=+0.0 s=1.0: E[sigma] released 0.715, fixed 0.714 | max quantile gap fixed vs released 0.0021
SD3    mode         m=+0.0 s=1.0: E[sigma] released 0.710, fixed 0.710 | max quantile gap fixed vs released 0.0020
FLUX.1 mode         m=+0.0 s=1.0: E[sigma] released 0.501, fixed 0.500 | max quantile gap fixed vs released 0.0020

With the fix, the realised noise levels match the paper's distribution (and FLUX.1 goes from E[sigma] 0.284 to 0.717 at m = +1.0, s = 0.6). Defaults stay within sampling noise. I agree with your README point: a short note about the sign change for anyone who sets --logit_mean makes sense, and I'll add it to the PR.

You're right about lumina2, thanks. Line 1395 interpolates (1.0 - sigmas) * noise + sigmas * model_input, so the two reversals cancel out and it shouldn't change. My 29 counted it by mistake, so 28 is correct.

Since you already have the fix, it's yours. I'll step back and won't open a PR.

Navigation Menu

Repository navigation

[training scripts] `--logit_mean` biases timesteps in the opposite direction to the SD3 paper (positive m oversamples low noise) #15166

Describe the bug

Reproduction

Logs

System Info

Who can help?

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[training scripts] --logit_mean biases timesteps in the opposite direction to the SD3 paper (positive m oversamples low noise) #15166

Description

Describe the bug

Reproduction

Logs

System Info

Who can help?

Activity

mlzerotoheroyoutube commented on Nov 3, 2026

Andreas-Hum commented on Nov 3, 2026

mlzerotoheroyoutube commented on Nov 3, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

[training scripts] `--logit_mean` biases timesteps in the opposite direction to the SD3 paper (positive m oversamples low noise) #15166