Question about discrete sCM loss vs standard CM loss performance

[SpursLipu]

Hi, thanks for your great work!

I have a question regarding training with the consistency loss.

Currently, my model heavily relies on custom operators, which makes it difficult to directly compute
𝑑𝐹/𝑑𝑡 using JVP. Because of this, I tried to reproduce the discrete sCM loss (with fd_type == 2) as an alternative.

At the same time, I also implemented the standard consistency model (CM) loss, i.e., directly computing the L2 loss between:

the predicted 𝑥0 from 𝑥𝑡​, and
the predicted 𝑥0 from 𝑥𝑡+1​

However, in my experiments, I found that this simpler CM loss actually performs better than the discrete sCM formulation.

So I have a few questions:

Have you observed similar behavior where discrete sCM underperforms compared to the standard CM loss?
Do you have any suggestions on how to improve the performance of discrete sCM in practice?
Are there any recommended tricks (e.g., better finite difference schemes, weighting strategies, or stabilization techniques) when JVP is not available?

From my understanding, sCM can suffer from issues like error accumulation and instability in derivative estimation, especially when approximating time derivatives numerically , so I wonder if this might be related to what I am seeing.

Thanks a lot for your help!

[zhengkw18]

I also discovered that using finite differences to approximate JVP is not as good as directly using discrete-time CM. When JVP is not available, I recommend directly using discrete-time CM, or combining discrete-time CM with DMD following rCM.

From my understanding, sCM can suffer from issues like error accumulation and instability in derivative estimation, especially when approximating time derivatives numerically

If approximating time derivatives, the method is no longer sCM. sCM requires exact JVP computation. The finite-difference approximations are not presented in the sCM paper, but rCM, and they are not used on hard tasks like video generation because we find severe degradation.

[SpursLipu]

Thank you for your detailed reply! I’ll be sure to share any new findings with you as I continue experimenting.

[zhengkw18]

Great! Looking forward to your further exploration. I will also add a discrete-time CM implementation to this repo.

[zhengkw18]

I have added a discrete-time CM implementation https://github.com/NVlabs/rcm/blob/main/rcm/models/t2v_model_distill_rcm.py#L591, for your reference.

[SpursLipu]

Thanks a lot for sharing your implementation — I really appreciate it!

After going through your code, I noticed a slight difference compared to my approach. In my implementation, after computing ( x_{t_2} ) using the Euler step, I use teacher-denoised ( x_0 ) and student-denoised ( x_0 ) to compute the loss. In contrast, it seems your implementation relies on the student for denoising in both steps.

So the two implementations are a bit different. From my side, this setup seems to work reasonably well in practice, although it does introduce one extra CFG-based forward pass (from the teacher).

For reference, here is a simplified version of my implementation (cleaned up a bit for clarity):

def training_CM_loss(self, models, real_score_function, inputs_shared, inputs_posi, inputs_nega):
    max_timestep_boundary = int(inputs_shared.get("max_timestep_boundary", 1) * self.score_scheduler.num_train_timesteps)
    min_timestep_boundary = int(inputs_shared.get("min_timestep_boundary", 0) * self.score_scheduler.num_train_timesteps)

    input_latents = inputs_shared["input_latents"]
    noise_template = inputs_shared["noise"]
    cfg_scale = inputs_shared["cfg_scale"]

    bs, _, T, H, W = input_latents.shape

    self.score_scheduler.set_timesteps(
        self.score_scheduler.num_train_timesteps,
        dynamic_shift_len=(H * W) // 4,
        training=True,
        num_frames=T
    )

    timestep_id = torch.randint(min_timestep_boundary, max_timestep_boundary - 2, (1,))
    timestep = self.score_scheduler.timesteps[timestep_id].to(dtype=torch.float32, device=self.device)
    sigma = self.score_scheduler.sigmas[timestep_id].to(dtype=torch.float32, device=self.device)

    noise = torch.randn_like(noise_template).float()
    latents_t = ((1 - sigma) * input_latents.float() + sigma * noise).to(dtype=self.torch_dtype)

    with torch.inference_mode():
        # teacher vt at t
        teacher_posi_v = self.model_fn(**real_score_function, **inputs_posi, latents=latents_t, timestep=timestep).float()
        teacher_nega_v = self.model_fn(**real_score_function, **inputs_nega, latents=latents_t, timestep=timestep).float()
        teacher_v = teacher_nega_v + cfg_scale * (teacher_posi_v - teacher_nega_v)

    # student vt at t
    student_v = self.model_fn(**models, **inputs_posi, latents=latents_t, timestep=timestep).float()
    x0_t_student = latents_t.detach() - student_v * sigma

    # find next valid timestep
    gap_step = 1
    while True:
        timestep_2 = self.score_scheduler.timesteps[timestep_id + gap_step].to(dtype=torch.float32, device=self.device)
        sigma_2 = self.score_scheduler.sigmas[timestep_id + gap_step].to(dtype=torch.float32, device=self.device)
        if timestep_2 != timestep and sigma_2 != sigma:
            break
        gap_step += 1

    # Euler step
    xt_2 = latents_t.float() + (sigma_2 - sigma) * teacher_v

    with torch.inference_mode():
        # teacher vt at t2
        teacher_posi_v_2 = self.model_fn(**real_score_function, **inputs_posi, latents=xt_2, timestep=timestep_2).float()
        teacher_nega_v_2 = self.model_fn(**real_score_function, **inputs_nega, latents=xt_2, timestep=timestep_2).float()
        teacher_v_2 = teacher_nega_v_2 + cfg_scale * (teacher_posi_v_2 - teacher_nega_v_2)

        x0_t2_teacher = (xt_2 - teacher_v_2 * sigma_2).detach()

    weight = 1 / (sigma ** 2).to(dtype=torch.float64, device=self.device)
    loss = ((x0_t_student.double() - x0_t2_teacher.double()) ** 2).mean() * weight

    return loss

Curious to hear your thoughts on this difference — especially whether using teacher ( x_0 ) at the second step vs fully student-based denoising has any impact in your experiments.

Thanks again for sharing your work!

[zhengkw18]

Consistency models should ensure that the student prediction at the current step matches the student prediction at the previous step. So both terms should be student prediction ($\theta$ in the image). The teacher $\phi$ is only used to move the datapoint backward along the teacher ODE.

[SpursLipu]

Consistency models should ensure that the student prediction at the current step matches the student prediction at the previous step. So both terms should be student prediction ( θ in the image). The teacher ϕ is only used to move the datapoint backward along the teacher ODE.

Thanks for the discussion!

Based on my understanding of the original paper, the target model here is supposed to be the EMA version of the model. In our setup, we didn’t introduce EMA, so I previously tried directly using the student model to compute ( x_{t_2} ). However, the results were not very good—the training was somewhat unstable. That said, it’s also possible that my implementation at the time wasn’t fully correct, and there may have been other issues affecting the outcome.

After that, I switched to using the teacher model for this part. While this is not strictly aligned with the consistency loss formulation from the original method, it empirically works quite well in practice.

I think this variant could also be interpreted more like a regression loss. My intuition is that it may help stabilize training, although it might come with a different performance ceiling. I’ll continue experimenting with the original formulation to better understand the trade-offs.

Thanks again for sharing your implementation—this has been very helpful!

[zhengkw18]

Thanks for the discussion!

Based on my understanding of the original paper, the target model here is supposed to be the EMA version of the model. In our setup, we didn’t introduce EMA, so I previously tried directly using the student model to compute ( x_{t_2} ). However, the results were not very good—the training was somewhat unstable. That said, it’s also possible that my implementation at the time wasn’t fully correct, and there may have been other issues affecting the outcome.

After that, I switched to using the teacher model for this part. While this is not strictly aligned with the consistency loss formulation from the original method, it empirically works quite well in practice.

I think this variant could also be interpreted more like a regression loss. My intuition is that it may help stabilize training, although it might come with a different performance ceiling. I’ll continue experimenting with the original formulation to better understand the trade-offs.

Thanks again for sharing your implementation—this has been very helpful!

Yes, the initial CM paper uses EMA, but in Song Yang's following paper (IMPROVED TECHNIQUES FOR TRAINING CONSISTENCY MODELS), EMA is removed. sCM also uses no EMA. I think the time scheduler matters for CM stability. My version uses a uniform-shift 5 schedule, and in my previous training, it performed fine. Your version possibly does not converge to a real consistency model, but it may be more stable and also suitable to act as the initialization stage of DMD. Looking forward to your further exploration!