Question about the three schedulers from generator and real/fake teacher

[lxxie298]

This is truly a remarkable piece of work! Many thanks to the author for open-sourcing it.

I have a question about the critic loss computation.

From my understanding, the student predicts $x_0$ using the generator's scheduler (denoted as scheduler_stu). Then scheduler_stu is used to add noise, and the resulting noisy sample is passed to the fake/real teachers.

However, the teachers use a scheduler with a different shift value from the student's scheduler to predict $x_0$ from pred_flow. After that, the code converts $x_0$ back to pred_flow again using the generator's scheduler.

I am wondering why the critic loss is not computed directly from the teacher's original output (pred_flow).

Instead, pred_flow is first transformed to $x_0$ and then converted back using a different shift setting. Is this intentional, or could it be a bug?

file path model/dmd.py

  _, pred_fake_image = self.fake_score(
      noisy_image_or_video=noisy_generated_image, # !!!!!!!!!!!!!!!!!
      conditional_dict=conditional_dict,
      timestep=critic_timestep
  ) # Using fake teacher's scheduler, default shift = 8

flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
    scheduler=self.scheduler,
    x0_pred=pred_fake_image.flatten(0, 1),  # !!!!!!!!!!!!!!!!!
    xt=noisy_generated_image.flatten(0, 1),
    timestep=critic_timestep.flatten(0, 1)
) # Using generator's scheduler, shift was set to 5.

[zhuhz22]

Hi @lxxie298 , thank you for your attention to our work!
Actually, the whole asymmetric DMD part is copied directly from Self Forcing, and the code is the same to theirs.
As for the issue you raise, I personally think it doesn't matter: critic_timestep has already been shifted before, and in neither self.fake_score() nor WanDiffusionWrapper._convert_x0_to_flow_pred() would the timeshift be called. So although the models have different shift, it doesn't actually matter.

As far as I understand, what deserves your attention is the following implementation:

    @staticmethod
    def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
        """
        Convert x0 prediction to flow matching's prediction.
        x0_pred: the x0 prediction with shape [B, C, H, W]
        xt: the input noisy data with shape [B, C, H, W]
        timestep: the timestep with shape [B]

        pred = (x_t - x_0) / sigma_t
        """
        # use higher precision for calculations
        original_dtype = x0_pred.dtype
        x0_pred, xt, sigmas, timesteps = map(
            lambda x: x.float().to(x0_pred.device), [x0_pred, xt,
                                                      scheduler.sigmas,
                                                      scheduler.timesteps]
        )
        timestep_id = torch.argmin(
            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
        flow_pred = (xt - x0_pred) / sigma_t
        return flow_pred.to(original_dtype)

, where timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) guarantees that the input timestep is exactly what the function uses, avoiding the influence of timeshift.

[lxxie298]

Thanks for the clarification!

I understand that critic_timestep is already shifted before being passed into self.fake_score(), and that neither self.fake_score() nor _convert_x0_to_flow_pred() applies the time-shift function again.

However, my concern is slightly different. Even though the same critic_timestep value is passed to both functions, the sigma used in the two conversions may still be different if the schedulers have different shift values.

In self.fake_score(), the model first predicts the original flow output, say $v_\theta$, and then internally converts it to $x_0$ using the fake teacher's scheduler:

$$ \hat{x}_0 = x_t - \sigma_{\text{fake}} v_\theta $$

Later, in the critic loss, this predicted $x_0$ is converted back to flow using the generator's scheduler:

$$ \hat{v} = \frac{x_t - \hat{x}_0}{\sigma_{\text{gen}}} = \frac{\sigma_{\text{fake}}}{\sigma_{\text{gen}}} v_\theta $$

Therefore, the recovered flow is exactly equal to the fake teacher's original pred_flow only when:

$$ \sigma_{\text{fake}} = \sigma_{\text{gen}} $$

My understanding is that timestep_id = argmin(...) only finds the nearest timestep value inside each scheduler's own scheduler.timesteps. But since scheduler.timesteps is computed from scheduler.sigmas, and scheduler.sigmas depends on the scheduler's shift, two schedulers with different shift values can map the same input critic_timestep to different sigma values.

So I agree that the shift function is not called again, but I am not fully convinced that this completely removes the influence of different shifts. It seems that the current implementation may introduce a scaling factor of approximately:

$$ \frac{\sigma_{\text{fake}}}{\sigma_{\text{gen}}} $$

unless both schedulers use the same shift or their sigma grids happen to match closely at that timestep.

That is why I was wondering whether, for denoising_loss_type == "flow", it would be more direct and mathematically cleaner to compute the critic loss using the fake teacher's original pred_flow, i.e. the first output of self.fake_score(), instead of converting pred_flow -> x0 -> pred_flow through potentially different schedulers.

Maybe this scaling difference is negligible in practice because the timestep grid is dense, or maybe the intention is to express the fake teacher prediction under the generator scheduler's noise parameterization. But if the generator scheduler uses shift=5 while the fake teacher scheduler uses shift=8, I think the two approaches are not strictly equivalent.

Please let me know if I am missing something here. My main question is whether this x0 round-trip through different schedulers is intentional, or whether the critic flow loss should directly use the fake teacher's original flow prediction.

[lxxie298]

I found a self.scheduler.set_timesteps(1000, training=True) call in WanDiffusionWrapper, and the code is:

def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
    sigma_start = self.sigma_min + \
        (self.sigma_max - self.sigma_min) * denoising_strength
    if self.extra_one_step:
        self.sigmas = torch.linspace(
            sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
    else:
        self.sigmas = torch.linspace(
            sigma_start, self.sigma_min, num_inference_steps)
    if self.inverse_timesteps:
        self.sigmas = torch.flip(self.sigmas, dims=[0])
    self.sigmas = self.shift * self.sigmas / \
        (1 + (self.shift - 1) * self.sigmas)
    if self.reverse_sigmas:
        self.sigmas = 1 - self.sigmas
    self.timesteps = self.sigmas * self.num_train_timesteps

So I now understand your point better.

Since scheduler.timesteps is computed from the shifted scheduler.sigmas, the input critic_timestep is already in the shifted timestep / sigma coordinate. Therefore, when _convert_flow_pred_to_x0() or _convert_x0_to_flow_pred() uses

timestep_id = torch.argmin(
    (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
)

it is not necessarily using the same raw index between the generator scheduler and the fake-teacher scheduler. Instead, it finds the nearest timestep value in each scheduler's own shifted timestep grid.

In this sense, even if the generator uses shift=5 and the fake teacher uses shift=8, the sigma used by the fake teacher should be close to the sigma used by the generator for the same critic_timestep, because both schedulers satisfy approximately:

$$ \text{timestep} = 1000 \cdot \sigma $$

So the ratio

$$ \frac{\sigma_{\text{fake}}}{\sigma_{\text{gen}}} $$

should usually be close to 1, except for small discretization / nearest-neighbor errors due to the two shifted timestep grids not being exactly identical.

Therefore, I agree that this should not introduce a large systematic mismatch in practice. My previous concern was based on thinking that both schedulers would use the same raw index, but actually the argmin lookup aligns them by the shifted timestep value.

Thank you very much for the explanation!

[zhuhz22]

Hi @lxxie298, thank you for your enthusiasm and dedication!

After carefully reading your analysis, I now tend to agree with you that this is indeed a typo. As you later analyzed, the error here is very small, but this design may indeed be rather redundant and somewhat awkward.

BTW, as mentioned earlier, the asymmetric DMD code here is copied from Self Forcing. Perhaps you could consider opening an issue in their official repo and ask why they deliberately designed this mechanism of repeated mapping rather than direct calculation?

[lxxie298]

Ok, thanks for your suggestion ! 😊🌹