Hi, I wanted to add a quick comment with some further justification for this change. The PR addresses a subtle but critical issue where changing num_train_timesteps can cause the asymptotic variance of the forward process to deviate from unity, breaking the assumption that $x_T$ matches the standard Gaussian prior that the reverse process starts from.

Problem

Using the standard notation

the variance of the final state $x_T$ is given by $Var(x_T) = 1 - ᾱ_T$. For the sampling process to match a standard Gaussian prior N(0, 1), we require $Var(x_T) \approx 1$. The value $ᾱ_T$ is controlled by the sum of betas since $-\log ᾱ_T \approx \sum_{i=1}^T \beta_i$. If this sum is too small, $ᾱ_T$ will not be close to zero and the variance will be incorrect.

How the Current Implementation Fails

The current implementation leads to an inconsistent $\sum_i \beta_i$ when $T$ is changed, which breaks the unit variance assumption.

• Default T = 1000: $\sum_i \beta_i \approx 10.05$ and $Var(x_T) \approx 1$.
• Naive change T = 200: $\sum_i \beta_i \approx 2.01$ which is far too small. This results in $Var(x_T) \approx 0.87$ and the prior is incorrect, which can cause significant issues during the reverse sampling process.

How the Proposed Fix Works

The proposed fix of scaling beta_end ensures the sum of betas remains approximately constant, thereby preserving the unit variance of the final state.

• With T = 200, beta_end is scaled to 0.1 and $\sum_i \beta_i \approx 10.01$.
• This ensures $\bar{\alpha}_T \approx 0$ and in turn $Var(x_T) \approx 1$, preserving the mathematical integrity of the diffusion process.
  This change ensures the scheduler produces a theoretically sound noise schedule by default, preventing users from having to manually correct for variance issues when experimenting with different numbers of training timesteps.

Hi @yiyixuxu ,

I wanted to gently follow up on this PR and provide some new empirical evidence for the practical impact of this change. I ran a simple experiment training a diffusion model on the butterfly dataset with num_train_timesteps=200, where the only difference between the two runs was the scale_betas_for_timesteps flag.

The images below show a clear improvement in sample quality, directly confirming the theoretical motivation for this PR.

Image 1: Before Fix         Image 2: After Fix (This PR)