拡散モデルの先読み勾配ガイダンス

def cond_score_v2(self, x, t, classes):
    r"""
    Gradient Guidance G_loss
    Following Definiton 1/ Lemma 1 in the original paper

    G_loss = - beta_t * \nabla_{x_t} (y - g^T * E[x_0 | x_t])^2
    """

    alpha_t = extract(self.sqrt_alphas_cumprod, t, (x.shape[0], 1))
    h_t = extract(self.sqrt_one_minus_alphas_cumprod, t, (x.shape[0], 1)) ** 2

    #### compute the gradient of g^T * E[x_0 | x_t] w.r.t. x_t
    with torch.enable_grad():
        x_in = x.detach().requires_grad_(True) # x_t, batched samples
        pred_noise = self.model(x_in, t, classes=None) # predict noise
        noise_svd = extract(
            self.sqrt_one_minus_alphas_cumprod, t, pred_noise.shape)
        uncond_score = -1 * pred_noise * 1 / noise_svd
        x_0_hat = (x_in + h_t * uncond_score) / alpha_t # get E[x_0 | x_t]
        value = torch.sum(x_0_hat * self.g, dim=1, keepdim=True)
        gradient = torch.autograd.grad(value.sum(), x_in)[0] # Get graidents for each sample in a batch parallelly

    ## 
    # beta_t 
    ## theoretical value: a variant of the beta_t shown in Lemma 1
    ## beta_t := alpha_t^2/2 / (alpha_t^2 * \sigma^2 + h(t) ||\nabla_{x_t} (g^T * E[x_0 | x_t])||^2)
    beta_t = (alpha_t ** 2 / 2) / (
        alpha_t ** 2 * self.sigma ** 2 + h_t * torch.norm(gradient, dim=1, keepdim=True) ** 2
        ) if self.beta is None else torch.full((x.shape[0], 1), self.beta, device=x.device)
    beta_t = beta_t * self.beta_coef 

    ## gradient of the loss (y - g^T * E[x_0 | x_t]) w.r.t. x_t
    gradient_2 = 2 * (value - classes.view(-1, 1)) * gradient 
    
    G = uncond_score - beta_t * gradient_2 # get the conditional score
    pred_noise = -1 * G * noise_svd # transform the conditional score to noise
    return pred_noise

def optimize(args, model, diffusion, generator, func, func_type, query_samples, interval, opt_logger, reward_K):
    '''
    Reward optimization with guided diffusion model
    '''
    print("====================Optimization Starts====================")
    if args.score_matching_finetune:
        optimizer_ft_sm = torch.optim.Adam(diffusion.model.parameters(), lr=args.sm_ft_lr, betas=(0.9, 0.99))
        sm_ft_round = args.sm_ft_round if args.sm_ft_round is not None and args.sm_ft_round > 0 else args.opt_rounds - args.sm_ft_start_round
        finetune_rounds = 0
        
    samples_list = []
    reward_list = []
    top_k_samples_list = []

    sample_time = 0
    opt_start_time = time.time()
    for i in range(args.opt_rounds):
        ## 1. get y_k (B_k, d_outer)
        # Query gradient g_k (get gradient for each sample z_(k,i))
        bs_gradient, bs_value = BatchGrad(query_samples, func)         
        diffusion.update_grad(bs_gradient) #, batch_size=bs_gradient.shape[0])
        # y_k = delta + g_k^T * z_k
        guidance = interval + torch.sum(bs_gradient * query_samples, dim=1)
        opt_logger.log(guidance=guidance[0].item(), mean_guidance=torch.mean(guidance).item())

        if args.score_matching_finetune:
            ## Alg.2: Fine-tuning Diffusion Model with Score Matching Loss
            if i >= args.sm_ft_start_round and finetune_rounds < sm_ft_round:
                model.train()
                steps = args.sm_ft_step_per_round # [Deprecated] default: 1
                batch_size = len(query_samples) // steps
                for j in range(steps):
                    loss_diffusion = diffusion(query_samples[j*batch_size: (j+1)*batch_size, :])
                    optimizer_ft_sm.zero_grad()
                    loss_diffusion.backward()
                    optimizer_ft_sm.step()
                    opt_logger.log(loss_diffusion=loss_diffusion.item())
                finetune_rounds += 1
                print("finetune the pretrained score model | round: ", finetune_rounds)
                model.eval()

                if finetune_rounds == sm_ft_round:
                    save_checkpoint(...)
        
        query_samples = diffusion.sample(
            num_samples=args.generate_bs,
            classes= guidance if not args.disable_guidance else None,
            )

        reward = func(query_samples)
        mean_reward = torch.mean(reward)
        ratio = generator.off_support_ratio(query_samples)
        reward_list.append(reward.cpu().numpy())

    return samples_list, reward_list, top_k_samples_list