[Diffusion DPO] Loss function and reproducing results in the paper

[yigu1008]

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.

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!

[yigu1008]

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.

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!

[kodinkod]

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:

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)

[tangjiuqi097]

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:

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 )

[bram-w]

@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!

[sayakpaul]

@bram-w thanks! Would you be willing to open a PR?

That's be of great help!