SDPO example for rectified flow

[cryptexis]

Hi,

thanks for your great work!

Would it be possible to release a small script that will indicate how to adapt your script for rectified flow?

[heekhero]

Hi, thanks a lot for your interest in our work!

To clarify how winner-preserving SDPO is integrated to rectified flow: the actual code change is quite small.
You only need to rescale the loser branch losses at the right place, right after you have
the per-sample winner / loser losses. Conceptually, you insert something like this:

# model_losses_w, model_losses_l: per-sample losses for winner / loser halves

if args.use_winner_preserving:
    lose_l_scale = get_adaptive_lose_l_scale(
        model_pred,
        target,
        args.winner_preserving_mu,
    )
else:
    lose_l_scale = 1.0

# replace the gradient on the loser branch only
model_losses_l_scaled = (
    model_losses_l.detach()
    + lose_l_scale * (model_losses_l - model_losses_l.detach())
)

# use the scaled loser loss when forming the preference gap
model_diff = model_losses_w - model_losses_l_scaled

Below is a slightly cleaned-up version of the core DPO block we use on FLUX, showing where this winner-preserving logic lives in context:

  # 1) Reference model prediction (no grad)
  with torch.inference_mode():
      ref_pred = self.ref_model(
          img=latents,
          img_ids=latent_pos_enc,
          txt=t5_encodings,
          txt_ids=text_pos_enc,
          y=clip_encodings,
          timesteps=timesteps,
      )

  # 2) Current model prediction
  model_pred = model(
      img=latents,
      img_ids=latent_pos_enc,
      txt=t5_encodings,
      txt_ids=text_pos_enc,
      y=clip_encodings,
      timesteps=timesteps,
  )

  # model_diff_w is the winner half of (model_pred - target)
  model_diff_w, _ = (model_pred - target).chunk(2)

  # per-sample MSE losses for winner / loser halves
  model_losses_w, model_losses_l = (
      self.loss_fn(model_pred, target)
      .mean(dim=(1, 2))   # average over spatial / channel dims
      .chunk(2)           # split batch into winner / loser
  )
  
  ############################################
  ############################################
 #These are the only lines of code you should change
  # 3) Winner-preserving scaling on the loser branch only
  if self.training.use_winner_preserving:
      lose_l_scale = get_adaptive_lose_l_scale(
          model_pred,
          target,
          self.training.winner_preserving_mu,
      )
  else:
      lose_l_scale = 1.0

  # replace gradient on model_losses_l with the scaled version
  model_losses_l_scaled = (
      model_losses_l.detach()
      + lose_l_scale * (model_losses_l - model_losses_l.detach())
  )

  # SDPO-adjusted preference gap for the student
  model_diff = model_losses_w - model_losses_l_scaled
  ############################################
  ############################################

  # 4) Reference model preference gap (no winner-preserving here)
  ref_losses_w, ref_losses_l = (
      self.loss_fn(ref_pred, target)
      .mean(dim=(1, 2))
      .chunk(2)
  )
  ref_diff = ref_losses_w - ref_losses_l

  # 5) DPO-style objective
  scale_term = -0.5 * self.job_config.training.beta_dpo
  inside_term = scale_term * (model_diff - ref_diff)

  # difference between student and reference predictions (winner half)
  pred2, _ = (model_pred - ref_pred).chunk(2)

  if self.job_config.training.dpo_method == "dspo":
      # DSPO-style quadratic form
      dpo_loss = (
          model_diff_w
          - self.job_config.training.beta_dpo
            * (1.0 - torch.sigmoid(inside_term))[:, None, None]
            * pred2
      ).pow(2).mean(dim=(1, 2)).mean()

  elif self.job_config.training.dpo_method == "diffusion-dpo":
      # vanilla Diffusion-DPO / DDPO-style objective
      dpo_loss = -torch.nn.functional.logsigmoid(inside_term).mean()

  else:
      raise NotImplementedError

  loss = dpo_loss

[cryptexis]

thank you! @heekhero - and the target stays as before correct ?

[heekhero]

@cryptexis Yes, the target stays exactly the same as naive flow-matching target like target = noise - image_latent. The change in this issue only affects how we form the loss, not how target is defined.

[cryptexis]

excellent, thank you very much!