[FetureRequest]Add DPO train

[sdbds]

source: https://github.com/huggingface/diffusers/blob/main/examples/research_projects/diffusion_dpo

[sdbds]

1、load lora

unet_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
    init_lora_weights="gaussian",
    target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)
unet.add_adapter(unet_lora_config)

2、# Final loss.

scale_term = -0.5 * args.beta_dpo
inside_term = scale_term * (model_diff - ref_diff)
loss = -1 * F.logsigmoid(inside_term.mean())

implicit_acc = (inside_term > 0).sum().float() / inside_term.size(0)
implicit_acc += 0.5 * (inside_term == 0).sum().float() / inside_term.size(0)

3、save lora

unet = accelerator.unwrap_model(unet)
unet = unet.to(torch.float32)
unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))

StableDiffusionXLLoraLoaderMixin.save_lora_weights(
    save_directory=args.output_dir,
    unet_lora_layers=unet_lora_state_dict,
    text_encoder_lora_layers=None,
    text_encoder_2_lora_layers=None,
)

[KohakuBlueleaf]

In Kohya I think we just need different dataset and loss term.

[artificialguybr]

Some users like me reporting some errors with the trainer... I dont have time to test more. Plan to back to test this weekend.

https://github.com/huggingface/diffusers/issues/6565

[feffy380]

In Kohya I think we just need different dataset and loss term.

The DPO loss requires predictions from the reference model, so we also need a way to toggle LoRA modules during training. I've been faking it with network_multiplier = 0 and adding a fast path that skips the rest of LoRAModule.forward.

[KohakuBlueleaf]

In Kohya I think we just need different dataset and loss term.

The DPO loss requires predictions from the reference model, so we also need a way to toggle LoRA modules during training. I've been faking it with network_multiplier = 0 and adding a fast path that skips the rest of LoRAModule.forward.

I think it is good to add some toggle api in network module Will add this API into kohya and LyCORIS

[rockerBOO]

Would like to discuss how this improvement would look.

Dataset:

Pairing images together with their preferences. Also possibly supported 4 image preference data. As an example this dataset kashif/pickascore has images in pairs with preference (+/-/=) or winner/loser/even or preferred/not-preferred/no preference.

1. Adding a pattern to dreambooth dataset could be an option where you have images paired with the same name and an indication of preference after.

image_w.png image_l.png image_e.png where w is preferred and l is not. e is even or no preference.

image_0.png image_2.png image_3.png image_1.png where they are listed in order

image_0_l.png image_1_l.png image_2_w.png image_3_l.png where the images are preferred w and the others are set as l

2. Preference json/yaml/toml

List out image pairs and preferences.

[{ image_0: "image.png", pref_0: -1, image_1: "image1.png", pref_1: 1}] 

Captions and other dataset features can be declared the same way as dreambooth dataset. This allows the dataset to remain the same but we can set the preferences in this separate file. Loading the dataset would required aligning these image pairs together so a complete preference file would be needed.

3. New dataset similar to fine-tuning or in conjunction

Complete a full metadata file including preference data. This would align well since this JSON/dataset file would need to be made and would put it all together. Would probably want to collate the images into pairs and have this be worked through separately in the fine-tuning dataset vs having each image separately listed and trying to pair their after the fact.

edit

4. Image pairing using A/B subsets using identical naming

For each dataset, having a A/B or W/L subset using identical naming.

subset/w/image.png subset/l/image.png. Captions should be set on both sides.

Alternatively

subset/image.png subset_l/image.png

Training

In training we have pairs or up to 4 preferences (if desired to go down the 4 preference path). We calculate the loss like normal and then separate the losses into the appropriate pairings.

                model_losses = F.mse_loss(model_pred.float(), target.float(), reduction="none")
                model_losses = model_losses.mean(dim=list(range(1, len(model_losses.shape))))
                model_losses_w, model_losses_l = model_losses.chunk(2)

https://github.com/huggingface/diffusers/blob/1d9a6a81b9a00b2dfab992ba75eb55bf8afb4eae/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py#L1010-L1012

Then we calculate the loss on the model without the LoRA.

                    ref_loss = F.mse_loss(ref_preds.float(), target.float(), reduction="none")
                    ref_loss = ref_loss.mean(dim=list(range(1, len(ref_loss.shape))))

                    ref_losses_w, ref_losses_l = ref_loss.chunk(2)
                    ref_diff = ref_losses_w - ref_losses_l

https://github.com/huggingface/diffusers/blob/1d9a6a81b9a00b2dfab992ba75eb55bf8afb4eae/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py#L1027-L1031

Then apply the loss

                # Final loss.
                scale_term = -0.5 * args.beta_dpo
                inside_term = scale_term * (model_diff - ref_diff)
                loss = -1 * F.logsigmoid(inside_term).mean()

https://github.com/huggingface/diffusers/blob/1d9a6a81b9a00b2dfab992ba75eb55bf8afb4eae/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py#L1037C1-L1040C61

So to get this working, being able to run on the original model without the network applied, or bypassing the network. This has been stated above about how to do that, as well as modifications to make this more streamlined.

---

If anyone has input about the following, or progress they have been making towards these goals. We can then collaborate and bring together the necessary features to have this be a possibility.

---

Separately would be interesting to discuss separate preference dataset storage formatting. This would allow extensions and such to create datasets for use in training, and possibly formatting them to work well together. Using the pickascore dataset as a possible option but willing to propose different options. The pickascore dataset stores the individual images as well as their preference, prompt and other metadata so it is all inclusive.

Personally not sure what is the best option, but I have explored what could be possible here. I am looking to add preference workflows into different tools so would like to consider what others would like to have as a dataset format or concerns with the pickascore like dataset.

---

Thanks for reading.

[sdbds]

Now SOTA will be SPO https://rockeycoss.github.io/spo.github.io/

[rockerBOO]

SPO looks interesting but they have a model evaluator that hasn't yet described the architecture edit Looks like they have described it, but would still need to create it and train it. Possibly they will release a pre-trained model of it.

We implement the step-aware preference model by modifying the PickScore Kirstain et al. [2024] model2 into our step-aware preference model as follows. First, we inject the time condition into the vision encoder of PickScore by applying adaptive layer normalization after each attention layer and FFN layer, following the adaptive layer norm (adaLN) block design of DiT Peebles and Xie [2023]. Then, we fine-tune the modified model on the human preference pairs dataset, Pick-a-Pic Kirstain et al. [2024].

Relevant papers related to preference optimization.

DDPO DP3O Diffusion-DPO ORPO SPO

I think having the preference dataset will allow us the flexibility at implementing these different optimization paths.

[rockerBOO]

New paper related to ORPO MaPO