How should `sde_steps` be chosen relative to the sigma schedule?

[Auraithm]

Thank you for your nice work!
Description:

Different models have very different sigma schedules due to different shift parameters. For example, with 10 inference steps:

Low shift (mu ≈ 0.5-1.15):

step:  0     1     2     3     4     5     6     7     8     9
sigma: 1.00  0.93  0.86  0.78  0.69  0.59  0.49  0.38  0.27  0.15

High shift (mu ≈ 2.0):

step:  0     1     2     3     4     5     6     7     8     9
sigma: 1.00  0.98  0.95  0.91  0.87  0.81  0.73  0.61  0.43  0.10

The same sde_steps: [1, 2, 3] maps to very different sigma ranges depending on the schedule. How should sde_steps be set for different sigma distributions?

[Jayce-Ping]

Here is the current rationale for sde_steps:

• No Rigorous Theory: There is currently no definitive study mapping specific $\sigma$ distributions to optimal SDE steps. My settings are primarily empirical.
• Heuristic: Based on MixGRPO and TempFlow-GRPO, applying SDE in the early denoising stages (regardless of the exact $\sigma$ value) promotes better structural diversity.
• Relative Position: I believe the relative step index (e.g., the first 30% of steps) is a more robust proxy than absolute $\sigma$ values. Even with different $\mu$, the early steps consistently handle global structure.
• Empirical Choice: While FlowGRPO uses num_timesteps // 2, I found [1, 2, 3] to be a "sweet spot" for stability in my tests.

In short: There is generally no need to adjust sde_steps for different $\mu$ values.

[Auraithm]

Thanks for the explanation! One concern about high-shift sigma schedules:

We're working on integrating a new model whose default sigma schedule is heavily front-loaded (as simga_shift is high). For example, with 10 steps:

step:  0     1     2     3     4     5     6     7     8     9
sigma: 1.00  0.98  0.95  0.91  0.87  0.81  0.73  0.61  0.43  0.10

At the early steps, sigma is close to 1.0, meaning the latent is dominated by noise with very weak model signal. If SDE is applied in this high-noise regime, is there a concern that the signal-to-noise ratio is too low for the model to produce meaningful updates?

We tried sde_steps: [1, 2, 3] with this schedule and observed that the reward fluctuated heavily and failed to increase consistently.

[Jayce-Ping]

If SDE is applied in this high-noise regime, is there a concern that the signal-to-noise ratio is too low for the model to produce meaningful updates?

That’s a great point. I’m not entirely sure of the exact answer either, but my intuition suggests this shouldn't be a major concern. Since GRPO formulates denoising as a Markov process, the updates should theoretically remain effective regardless of the SNR. However, I could be mistaken.

the reward fluctuated heavily and failed to increase consistently.

I’ve observed the same phenomenon with models like Wan, where the reward fluctuates heavily and takes a long time to show a consistent increase. I haven't pinned down the exact cause yet, but the time shift could definitely be a factor.

I suggest to set sde_steps to a larger range, such as [0,1,2,3,4,5,6,7,8,9], i.e., across all timesteps, with more training steps like sde_steps=3 and see what happens.

[Auraithm]

Thank you again, I’ll try your approach. Happy New Year!