Repository: huggingface/diffusers

Issue:
[Diffusion DPO] Loss function and reproducing results in the paper
Hi, 

I got a question on the loss function used in both training scripts for diffusion DPO. 

In train_diffusion_dpo.py line 863 and train_diffusion_dpo_sdxl.py line 978
` loss = -1 * F.logsigmoid(inside_term.mean())`
From my understanding, this mean is the average taken across the batches as the empirical mean for formula (14) in the paper.
<img width="300" alt="Screenshot 2024-01-24 at 3 54 39 PM" src="https://github.com/huggingface/diffusers/assets/119135690/bddea053-5d43-4496-a057-d120d09fcdaf">
So I think the loss should instead be 
` loss = -1 * F.logsigmoid(inside_term).mean()`?

Also, may I know if there's anything else I need to modify other than changing the training configuration to reproduce the results in the paper, say 28.16 HPSv2 reward for the DPO-SDXL HPSv2 generations? Thank you very much!







Discussion:
I tried to evaluate both the StableDiffusion 1.5 and StableDiffusion XL checkpoints using 120 unique sampled prompts from pick-a-pic v1 validation set and computed mean hpsv2/pick score, here's an example of my workflow of hpsv2
        
        model_id = "runwayml/stable-diffusion-v1-5"
        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
        # load finetuned model
        unet_id = "mhdang/dpo-sd1.5-text2image-v1"
        unet = UNet2DConditionModel.from_pretrained(unet_id, subfolder="unet", torch_dtype=torch.float16)
        pipe.unet = unet
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        pipe = pipe.to(accelerator.device)
        pipe.set_progress_bar_config(disable=True)
        
        generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
        all_images =[]
        all_images_ref = []
        all_scores = []
        all_scores_ref = []
        context = contextlib.nullcontext() 
        
        
        # load hpsv2 models and processors
        model_name = "ViT-H-14"
        model, preprocess_train, preprocess_val = create_model_and_transforms(
            model_name,
            'laion2B-s32B-b79K',
            precision=weight_dtype,
            device=accelerator.device,
            jit=False,
            force_quick_gelu=False,
            force_custom_text=False,
            force_patch_dropout=False,
            force_image_size=None,
            pretrained_image=False,
            image_mean=None,
            image_std=None,
            light_augmentation=True,
            aug_cfg={},
            output_dict=True,
            with_score_predictor=False,
            with_region_predictor=False
        )    
              
        tokenizer = get_tokenizer(model_name)
        
        checkpoint_path = checkpoint_path = f"/home/ygu/.cache/hpsv2/HPS_v2_compressed.pt"
        # force download of model via score
        hpsv2.score([], "")
        
        checkpoint = torch.load(checkpoint_path, map_location=accelerator.device)
        model.load_state_dict(checkpoint['state_dict'])
        tokenizer = get_tokenizer(model_name)
        model = model.to(accelerator.device, dtype=weight_dtype)
        model.eval()
        
        target_size =  224
        normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                                    std=[0.26862954, 0.26130258, 0.27577711])
        PIL_to_tensor = transforms.ToTensor()
        
        for index in tqdm(list(range(args.sampling_iters)),position=0,disable=not accelerator.is_local_main_process):
          BATCHED_PROMPTS = VALIDATION_PROMPTS[args.sampling_batch_size*index: args.sampling_batch_size*(index+1)]
          images = []
          images_ref = []
          for prompt in BATCHED_PROMPTS:
              with context:
                  image = pipe(prompt, num_inference_steps=25, guidance_scale=7.5).images[0].resize((512,512))
                  image_ref = pipe_ref(prompt, num_inference_steps=25, guidance_scale=7.5).images[0].resize((512,512))
                  
                  all_images.append(image)
                  all_images_ref.append(image_ref)
          
              
                  image = PIL_to_tensor(image).to(accelerator.device,weight_dtype)
                  image_ref = PIL_to_tensor(image_ref).to(accelerator.device,weight_dtype)


It turns out the score doesn't differ very much. 
<img width="400" alt="Screenshot 2024-01-24 at 11 15 44 PM" src="https://github.com/huggingface/diffusers/assets/119135690/d70250c9-7134-4cc3-b8d8-45b9abcd3d0f">
  
I know the score is only part of the metrics but I'm still wondering if I got something wrong in my code or evaluation setup, thank you!
Hi, I also have a question about loss

In the presented implementation:

```
model_diff = model_losses_w - model_losses_l
ref_diff = ref_losses_w - ref_losses_l
inside_term = scale_term * (model_diff - ref_diff)
```

but in paper code:

<img width="605" alt="Снимок экрана 2024-01-28 в 13 26 22" src="https://github.com/huggingface/diffusers/assets/69761539/5b0271df-7bfc-4f19-abc6-c4ef6ac2e73f">

maybe loss should look like this?

```
model_diff = model_losses_w- ref_losses_w  
ref_diff = model_losses_l -  ref_losses_l 
inside_term = scale_term * (model_diff - ref_diff)
```

> Hi, I also have a question about loss
> 
> In the presented implementation:
> 
> ```
> model_diff = model_losses_w - model_losses_l
> ref_diff = ref_losses_w - ref_losses_l
> inside_term = scale_term * (model_diff - ref_diff)
> ```
> 
> but in paper code:
> 
> <img alt="Снимок экрана 2024-01-28 в 13 26 22" width="605" src="https://private-user-images.githubusercontent.com/69761539/300254914-5b0271df-7bfc-4f19-abc6-c4ef6ac2e73f.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MDY1MDA4NDIsIm5iZiI6MTcwNjUwMDU0MiwicGF0aCI6Ii82OTc2MTUzOS8zMDAyNTQ5MTQtNWIwMjcxZGYtN2JmYy00ZjE5LWFiYzYtYzRlZjZhYzJlNzNmLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDAxMjklMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwMTI5VDAzNTU0MlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWViOTlkMjQxNGY4MTAwYzFlMDAwNWIyNzhmYTcxNDFiNDBlMzI0MTFkYTZiNzU0YzZhNTAzNDUzNjljMjViNDcmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.rvq7KCVkfm9cj5c6lCOfocAmXDD_HH16wS2Mqq0gw9E">
> maybe loss should look like this?
> 
> ```
> model_diff = model_losses_w- ref_losses_w  
> ref_diff = model_losses_l -  ref_losses_l 
> inside_term = scale_term * (model_diff - ref_diff)
> ```

Both of them are the same as: 
`inside_term = scale_term * (model_losses_w - ref_losses_w - model_losses_l + ref_losses_l ) `
@sayakpaul confirming that the loss .mean() call here got out of order. It wouldn't have an effect at batch size=1 but above that would change the gradient signal substantially. Checked our internal implementations and it's correct there so must have gotten mis-translated somewhere along the way. Sorry for not catching it during review!
@bram-w thanks! Would you be willing to open a PR? 

That's be of great help!

Resolving patch:
diff --git a/examples/research_projects/diffusion_dpo/train_diffusion_dpo.py b/examples/research_projects/diffusion_dpo/train_diffusion_dpo.py
--- a/examples/research_projects/diffusion_dpo/train_diffusion_dpo.py
+++ b/examples/research_projects/diffusion_dpo/train_diffusion_dpo.py
@@ -860,7 +860,7 @@ def collate_fn(examples):
                 # Final loss.
                 scale_term = -0.5 * args.beta_dpo
                 inside_term = scale_term * (model_diff - ref_diff)
-                loss = -1 * F.logsigmoid(inside_term.mean())
+                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)
diff --git a/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py b/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py
--- a/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py
+++ b/examples/research_projects/diffusion_dpo/train_diffusion_dpo_sdxl.py
@@ -975,7 +975,7 @@ def compute_time_ids(original_size, crops_coords_top_left):
                 # Final loss.
                 scale_term = -0.5 * args.beta_dpo
                 inside_term = scale_term * (model_diff - ref_diff)
-                loss = -1 * F.logsigmoid(inside_term.mean())
+                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)


Test patch: