Clarification: training sampler (generate_samples_from_batch) vs inference sampler mismatch (TrigFlow vs RF) — is this intended?

[LeslieZhoa]

Hi, thank you for open-sourcing the code and the great paper. I’m trying to better understand the sampling formulation and align my own evaluation pipeline with the training setup.

While reading the repo, I noticed that the sampling logic used in training (student visualization / rollout) appears different from the sampling logic in the provided inference script. Both work well in practice, but I’m not sure whether the difference is intentional (e.g., an equivalent re-parameterization) or just an implementation choice.

What I observed

1) Training sampler (student) — generate_samples_from_batch()

In T2VDistillModel_rCM.generate_samples_from_batch() the student sampling uses TrigFlow angle timesteps (call them tau):

• timesteps (TrigFlow angle):

  t_steps = [atan(sigma_max)] + mid_t + [0]

• at each step:
  • call self.denoise(x, t_cur, ..., net_type="student").x0 to get x0
  • update state using a TrigFlow-style “re-noising” step:

    x = cos(t_next) * x0 + sin(t_next) * randn()

Inside denoise(), the network time input is c_noise = sin(tau)/(cos(tau)+sin(tau)) * 1000 via RectifiedFlow_TrigFlowWrapper, and the input is scaled by c_in = 1/(cos+sin).

2) Inference script sampler — infer_*.py

In the inference script, the code starts from the same angle-like timesteps but then converts them to a RectifiedFlow-style time:

t_steps = [atan(sigma_max), *mid_t, 0]
t_steps = sin(t_steps) / (cos(t_steps) + sin(t_steps))   # RF time, call it s(tau)
x = init_noise * t_steps[0]

Then each step uses an RF update formula:

v_pred = net(x, timesteps=t_cur*1000)
x = (1 - t_next) * (x - t_cur * v_pred) + t_next * randn()

So inference:

• advances in RF time s = sin(tau)/(cos+sin)
• initializes x as noise * s0
• does not apply the same (cos+sin)^{-1} input scaling that denoise() applies during training sampling
• uses a different state update equation than the training TrigFlow-style cos/sin re-noising step

My question

Is this discrepancy intended / mathematically equivalent under a certain change of variables, or is inference using a different (but empirically good) heuristic sampler?

More specifically:

1. Should “faithful” inference for the distilled student follow the same TrigFlow iteration as training (x -> x0 -> cos/sin re-noise)?
2. If the RF sampler is intended, is there a recommended way to map the training formulation (with c_in/c_skip/c_out) into an RF-style sampler so the two match more closely?
3. Is there a reason inference initializes with x = noise * t0 (RF boundary scaling) while training sampling initializes with x = noise?

Any clarification (and/or recommended sampler to use for reporting metrics) would be greatly appreciated. Thanks again for the excellent work!

[meng-ustc]

I have the same question: why the inference script uses a different sampler. Though I think both samplers should work, and I tried both in my experiments, they could have quite different performance, which will depend on the time steps you actually use in the sampling. But it's not fair to compare side by side, because the SNR can't be matched. This is related to my another question, in the paper it said, the TrigFlow wrapper is supposed to match the input SNR of the original RF model. However, assume we have RF timestep t, the SNR will be (1-t)/t, then if we make a TrigFlow timestep t_trig, which enables c_noise(t_trig)=t, the SNR will be 1/tan(t_trig). So my understanding is the wrapper doesn't guarantee the SNR, instead it makes the mapped xt still be a sum of s(t)*x0 and (1-s(t))*e (e is noise epsilon).

Btw, regarding the third question in the post, I remember the starting time step is picked for higher generation diversity, then I think the init latent is a scaling of the SNR to match the starting time step.

Please correct me if I make anything wrong. Appreciate it!

[zhengkw18]

Thank you all for your engagement with our work, and great observation! This is an equivalent re-parameterization for sampling (maybe you can derive it yourself). In fact, you can try our earlier version of inference code: https://github.com/NVlabs/rcm/blob/b8fa13b098497f3a9c5a440e24976a8832af2e50/rcm/inference/wan2pt1_t2v_rcm_infer.py. It produces almost the same sampling result to the current RF-based version.

We use the RF-based version as our current inference code because it looks simpler. For JVP-related training loss, TrigFlow-based and RF-based losses will be different; we found that the TrigFlow-based implementation (used in sCM) is more stable in training, compared to the RF-based version (e.g., MeanFlow). We are still doing more ablation studies.

For your question 3, different noise schedules in diffusion models can be aligned by:

1. Aligning the timesteps to ensure the same SNR. This is achieved by the following code:

# Convert TrigFlow timesteps to RectifiedFlow
t_steps = torch.sin(t_steps) / (torch.cos(t_steps) + torch.sin(t_steps))

2. Adjusting the scale of xt. For example, for $x_t=\cos t x_0+\sin t \epsilon$ in TrigFlow, the corresponding xt at the same SNR in RF is
   $$x_t^{RF}=\frac{\cos t}{\cos t+\sin t}x_0+\frac{\sin t}{\cos t+\sin t}\epsilon$$
   This is because RF should satisfy the coefficients before $x_0$ and $\epsilon$ sum to 1. So there is a scaling factor $\cos t+\sin t$. That is the answer for your question 3.

[meng-ustc]

Thanks for your explanation!
It's clear to me now, and I just realized in my last message, when I expand (1-t)/t with the c_noise(t) formula, I'll get (1-t)/t=1/tan(t_trig). Though I still don't understand if we wrap it to RF and they are convertible, how does it impact the training stability. Looking forward to your findings in your ongoing ablations.

Btw, in your inference code, there is optional ode or sde. Did you compare the performance before? Do you think the sCM distilled model has a more straight ode trajectory than the before-distillation model (if this ode sampling still works well, though at the same time sCM loss function has no constrains to the trajectory itself)?

[zhengkw18]

ODE will produce sharper textures. However, SDE is generally more robust. For the ODE straightness, I do not prefer to explain distillation in this way.