| import itertools |
| from dataclasses import dataclass |
|
|
| import hydra.utils |
| import lightning as L |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import transformers |
|
|
| import dataloader |
| import metrics |
| import models |
| import utils |
|
|
|
|
| @dataclass |
| class Loss: |
| loss: torch.FloatTensor |
| nlls: torch.FloatTensor |
| prior_loss: torch.FloatTensor |
| num_tokens: torch.FloatTensor |
|
|
|
|
| class LogLinear(torch.nn.Module): |
| def __init__(self, eps): |
| super().__init__() |
| self.eps = eps |
|
|
| def forward(self, t): |
| t = (1 - self.eps) * t |
| alpha_t = 1 - t |
| dalpha_t = - torch.ones_like(alpha_t) * (1 - self.eps) |
| return dalpha_t, alpha_t |
|
|
| def get_t_for_alpha(self, alpha_t): |
| return 1 - alpha_t |
|
|
|
|
| class Cosine(torch.nn.Module): |
| def __init__(self, eps): |
| super().__init__() |
| self.eps = eps |
| self.half_pi = torch.pi / 2 |
|
|
| def forward(self, t): |
| t = (1 - self.eps) * t |
| alpha_t = 1 - torch.cos(self.half_pi * (1 - t)) |
| dalpha_t = - torch.sin(self.half_pi * (1 - t)) * self.half_pi |
| return dalpha_t, alpha_t |
|
|
| def get_t_for_alpha(self, alpha_t): |
| is_tensor = torch.is_tensor(alpha_t) |
| if not is_tensor: |
| alpha_t = torch.tensor([alpha_t]) |
| t = 1 - 2 / torch.pi * torch.acos(1 - alpha_t) |
| if not is_tensor: |
| t = t.cpu().item() |
| return t |
|
|
|
|
| def sample_categorical(categorical_probs): |
| gumbel_norm = ( |
| 1e-10 |
| - (torch.rand_like(categorical_probs) + 1e-10).log()) |
| return (categorical_probs / gumbel_norm).argmax(dim=-1) |
|
|
|
|
| def _unsqueeze(x, reference): |
| return x.view( |
| * x.shape, |
| * ((1,) * (len(reference.shape) - len(x.shape)))) |
|
|
|
|
| class TrainerBase(L.LightningModule): |
| def __init__( |
| self, |
| config, |
| tokenizer: transformers.PreTrainedTokenizer, |
| vocab_size=None): |
| super().__init__() |
| self.save_hyperparameters() |
| self.config = config |
| if hasattr(self.config.algo, 'ignore_bos'): |
| self.ignore_bos = config.algo.ignore_bos |
| else: |
| self.ignore_bos = False |
| if hasattr(self.config.algo, 'loss_type'): |
| self.loss_type = config.algo.loss_type |
| self.tokenizer = tokenizer |
| if vocab_size is None: |
| self.vocab_size = len(self.tokenizer) |
| else: |
| self.vocab_size = vocab_size |
| self.sampler = self.config.sampling.predictor |
| self.antithetic_sampling = self.config.training.antithetic_sampling |
| self.parameterization = self.config.algo.parameterization |
| if self.config.algo.backbone == 'dit': |
| self.backbone = models.dit.DIT( |
| self.config, vocab_size=self.vocab_size) |
| elif self.config.algo.backbone == 'unet': |
| self.backbone = models.unet.UNet(self.config, self.vocab_size) |
| elif self.config.algo.backbone == 'dimamba': |
| self.backbone = models.dimamba.DiMamba( |
| self.config, |
| vocab_size=self.vocab_size, |
| pad_token_id=self.tokenizer.pad_token_id) |
| elif self.config.algo.backbone == 'hf_dit': |
| self.backbone = transformers.AutoModelForMaskedLM.from_pretrained( |
| config.eval.checkpoint_path, trust_remote_code=True) |
|
|
| self.T = self.config.algo.T |
| self.num_tokens = self.config.model.length |
| self.p_nucleus = self.config.sampling.p_nucleus |
| |
| if config.noise.type == 'log-linear': |
| self.noise = LogLinear(config.noise.eps) |
| elif config.noise.type == 'cosine': |
| self.noise = Cosine(config.noise.eps) |
| else: |
| raise ValueError(config.noise.type) |
| |
| self.num_classes = config.data.get('num_classes', None) |
| self.class_conditional = self.num_classes is not None |
| self.class_cond_dropout = config.training.class_dropout_p |
|
|
| self.metrics = metrics.Metrics( |
| gen_ppl_eval_model_name_or_path=\ |
| self.config.eval.gen_ppl_eval_model_name_or_path, |
| eval_ppl_batch_size=\ |
| self.config.eval.perplexity_batch_size) |
|
|
| if self.config.training.ema > 0: |
| self.ema = models.ema.ExponentialMovingAverage( |
| self._get_parameters(), |
| decay=self.config.training.ema) |
| else: |
| self.ema = None |
| |
| self.lr = self.config.optim.lr |
| self.sampling_eps = self.config.training.sampling_eps |
| self.time_conditioning = self.config.algo.time_conditioning |
| self.neg_infinity = -1000000.0 |
| self.fast_forward_epochs = None |
| self.fast_forward_batches = None |
|
|
| def _validate_configuration(self): |
| assert self.config.algo.backbone in {'dit', 'hf_dit', |
| 'unet'} |
| if self.config.algo.parameterization == 'ar': |
| assert not self.config.algo.time_conditioning |
| assert self.config.prior.type == 'none' |
| if self.time_conditioning: |
| assert self.config.sampling.predictor != 'ancestral_cache' |
|
|
| if self.parameterization in {'score', 'mean'}: |
| assert self.time_conditioning |
| if self.T > 0: |
| assert self.parameterization != 'score' |
|
|
| def to(self, *args, **kwargs): |
| self = super().to(*args, **kwargs) |
| self.metrics.to(*args, **kwargs) |
| return self |
|
|
| def q_xt(self, x, alpha_t): |
| raise NotImplementedError |
|
|
| def _get_parameters(self): |
| return itertools.chain(self.backbone.parameters(), |
| self.noise.parameters()) |
|
|
| def _eval_mode(self): |
| if self.ema: |
| self.ema.store(self._get_parameters()) |
| self.ema.copy_to(self._get_parameters()) |
| self.backbone.eval() |
| self.noise.eval() |
|
|
| def _train_mode(self): |
| if self.ema: |
| self.ema.restore(self._get_parameters()) |
| self.backbone.train() |
| self.noise.train() |
|
|
| def on_load_checkpoint(self, checkpoint): |
| if self.ema: |
| self.ema.load_state_dict(checkpoint['ema']) |
| |
| |
| self.fast_forward_epochs = checkpoint['loops'][ |
| 'fit_loop']['epoch_progress']['current']['completed'] |
| self.fast_forward_batches = checkpoint['loops'][ |
| 'fit_loop']['epoch_loop.batch_progress'][ |
| 'current']['completed'] |
|
|
| def on_save_checkpoint(self, checkpoint): |
| if self.ema: |
| checkpoint['ema'] = self.ema.state_dict() |
| |
| |
| |
| |
| checkpoint['loops']['fit_loop'][ |
| 'epoch_loop.batch_progress']['total'][ |
| 'completed'] = checkpoint['loops']['fit_loop'][ |
| 'epoch_loop.automatic_optimization.optim_progress'][ |
| 'optimizer']['step']['total'][ |
| 'completed'] * self.trainer.accumulate_grad_batches |
| checkpoint['loops']['fit_loop'][ |
| 'epoch_loop.batch_progress']['current'][ |
| 'completed'] = checkpoint['loops']['fit_loop'][ |
| 'epoch_loop.automatic_optimization.optim_progress'][ |
| 'optimizer']['step']['current'][ |
| 'completed'] * self.trainer.accumulate_grad_batches |
| |
| |
| |
| checkpoint['loops']['fit_loop'][ |
| 'epoch_loop.state_dict'][ |
| '_batches_that_stepped'] = checkpoint['loops']['fit_loop'][ |
| 'epoch_loop.automatic_optimization.optim_progress'][ |
| 'optimizer']['step']['total']['completed'] |
| if 'sampler' not in checkpoint.keys(): |
| checkpoint['sampler'] = {} |
| if hasattr(self.trainer.train_dataloader.sampler, |
| 'state_dict'): |
| sampler_state_dict = self.trainer.\ |
| train_dataloader.sampler.state_dict() |
| checkpoint['sampler'][ |
| 'random_state'] = sampler_state_dict.get( |
| 'random_state', None) |
| else: |
| checkpoint['sampler']['random_state'] = None |
|
|
| def on_train_start(self): |
| if self.ema: |
| self.ema.move_shadow_params_to_device(self.device) |
| |
| |
| distributed = ( |
| self.trainer._accelerator_connector.use_distributed_sampler |
| and self.trainer._accelerator_connector.is_distributed) |
| if distributed: |
| sampler_cls = dataloader.FaultTolerantDistributedSampler |
| else: |
| sampler_cls = dataloader.RandomFaultTolerantSampler |
| updated_dls = [] |
| for dl in self.trainer.fit_loop._combined_loader.flattened: |
| if hasattr(dl.sampler, 'shuffle'): |
| dl_sampler = sampler_cls( |
| dl.dataset, shuffle=dl.sampler.shuffle) |
| else: |
| dl_sampler = sampler_cls(dl.dataset) |
| if (distributed |
| and self.fast_forward_epochs is not None |
| and self.fast_forward_batches is not None): |
| dl_sampler.load_state_dict({ |
| 'epoch': self.fast_forward_epochs, |
| 'counter': (self.fast_forward_batches |
| * self.config.loader.batch_size)}) |
| updated_dls.append( |
| torch.utils.data.DataLoader( |
| dl.dataset, |
| batch_size=self.config.loader.batch_size, |
| num_workers=self.config.loader.num_workers, |
| pin_memory=self.config.loader.pin_memory, |
| sampler=dl_sampler, |
| shuffle=False, |
| persistent_workers=True)) |
| self.trainer.fit_loop._combined_loader.flattened = updated_dls |
|
|
| def optimizer_step(self, *args, **kwargs): |
| super().optimizer_step(*args, **kwargs) |
| if self.ema: |
| self.ema.update(self._get_parameters()) |
|
|
| def _process_sigma(self, sigma): |
| raise NotImplementedError |
|
|
| def _process_model_output(self, model_output, xt, sigma): |
| raise NotImplementedError |
|
|
| def forward(self, xt, sigma, labels=None, weights=None, |
| nn_input_idxs=None): |
| if nn_input_idxs is None: |
| nn_input_idxs = xt |
|
|
| sigma = self._process_sigma(sigma) |
| with torch.amp.autocast('cuda', dtype=torch.float32): |
| model_output = self.backbone( |
| x=nn_input_idxs, sigma=sigma, class_cond=labels, |
| weights=weights) |
| return self._process_model_output( |
| model_output=model_output, xt=xt, sigma=sigma) |
|
|
| def on_train_epoch_start(self): |
| self.metrics.reset() |
| assert self.metrics.train_nlls.nll.mean_value == 0 |
| assert self.metrics.train_nlls.nll.weight == 0 |
|
|
| def training_step(self, batch, batch_idx): |
| current_accumulation_step = ( |
| batch_idx % self.trainer.accumulate_grad_batches) |
| losses = self._loss(batch['input_ids'], |
| batch.get('labels', None), |
| batch['attention_mask'], |
| current_accumulation_step, |
| train_mode=True) |
| self.metrics.update_train(losses.nlls, losses.prior_loss, |
| losses.num_tokens) |
| self.log(name='trainer/loss', |
| value=losses.loss.item(), |
| on_step=True, |
| on_epoch=False, |
| sync_dist=True) |
| return losses.loss |
|
|
| def on_train_epoch_end(self): |
| for k, v in self.metrics.valid_nlls.items(): |
| self.log(name=k, value=v.compute(), on_step=False, |
| on_epoch=True, sync_dist=True) |
|
|
| def on_validation_epoch_start(self): |
| self.metrics.reset() |
| self._eval_mode() |
| assert self.metrics.valid_nlls.nll.mean_value == 0 |
| assert self.metrics.valid_nlls.nll.weight == 0 |
|
|
| def validation_step(self, batch, batch_idx): |
| del batch_idx |
| losses = self._loss(batch['input_ids'], |
| batch.get('labels', None), |
| batch['attention_mask']) |
| self.metrics.update_valid(losses.nlls, losses.prior_loss, |
| losses.num_tokens) |
| return losses.loss |
|
|
| def on_validation_epoch_end(self): |
| for k, v in self.metrics.valid_nlls.items(): |
| self.log(name=k, value=v.compute(), on_step=False, |
| on_epoch=True, sync_dist=True) |
| if ((self.config.eval.compute_perplexity_on_sanity |
| or not self.trainer.sanity_checking) |
| and self.config.eval.generate_samples): |
| samples, text_samples = None, None |
| for _ in range( |
| self.config.sampling.num_sample_batches): |
| samples = self.generate_samples( |
| num_samples=self.config.loader.eval_batch_size) |
| |
| self.metrics.record_entropy(samples) |
| |
| text_samples = self.tokenizer.batch_decode(samples) |
| if self.config.eval.compute_generative_perplexity: |
| self.metrics.record_generative_perplexity( |
| text_samples, self.num_tokens, self.device) |
| if text_samples is not None: |
| if self.trainer.global_rank == 0 and hasattr( |
| self.trainer.logger, 'log_table'): |
| |
| text_samples = text_samples[ |
| : self.config.sampling.num_sample_log] |
| self.trainer.logger.log_table( |
| key=f'samples@global_step{self.global_step}', |
| columns=['Generated Samples'], |
| data=[[s] for s in text_samples]) |
| if self.config.eval.compute_generative_perplexity: |
| self.log('val/gen_ppl', |
| self.metrics.gen_ppl.compute(), |
| on_epoch=True, |
| on_step=False, |
| sync_dist=True) |
| self.log('val/sample_entropy', |
| self.metrics.sample_entropy.compute(), |
| on_epoch=True, |
| on_step=False, |
| sync_dist=True) |
| self._train_mode() |
|
|
| def configure_optimizers(self): |
| optimizer = torch.optim.AdamW( |
| self._get_parameters(), |
| lr=self.config.optim.lr, |
| betas=(self.config.optim.beta1, |
| self.config.optim.beta2), |
| eps=self.config.optim.eps, |
| weight_decay=self.config.optim.weight_decay) |
|
|
| scheduler = hydra.utils.instantiate( |
| self.config.lr_scheduler, optimizer=optimizer) |
| scheduler_dict = {'scheduler': scheduler, |
| 'interval': 'step', |
| 'monitor': 'val/loss', |
| 'name': 'trainer/lr'} |
| return [optimizer], [scheduler_dict] |
|
|
| def generate_samples(self, num_samples, num_steps, eps): |
| raise NotImplementedError |
|
|
| def restore_model_and_sample(self, num_steps, eps=1e-5): |
| """Generate samples from the model.""" |
| |
| self._eval_mode() |
| samples = self.generate_samples( |
| num_samples=self.config.loader.eval_batch_size, |
| num_steps=num_steps, |
| eps=eps) |
| self._train_mode() |
| return samples |
|
|
| def _process_model_input(self, x0, valid_tokens): |
| raise NotImplementedError |
|
|
| def nll(self, input_tokens, labels, output_tokens, |
| current_accumulation_step=None, train_mode=False): |
| raise NotImplementedError |
|
|
| def _loss(self, x0, labels, valid_tokens, |
| current_accumulation_step=None, |
| train_mode=False): |
| (input_tokens, output_tokens, |
| valid_tokens) = self._process_model_input( |
| x0, valid_tokens) |
| loss = self.nll(input_tokens, labels, output_tokens, |
| current_accumulation_step, train_mode) |
| assert loss.ndim == 2 |
| if self.ignore_bos: |
| loss[:, 1:] = loss[:, 1:] |
| valid_tokens[:, 1:] = valid_tokens[:, 1:] |
|
|
| nlls = (loss * valid_tokens).sum() |
| num_tokens = valid_tokens.sum() |
| token_nll = nlls / num_tokens |
|
|
| return Loss(loss=token_nll, |
| nlls=nlls, |
| prior_loss=0.0, |
| num_tokens=num_tokens) |
|
|
|
|
| class Diffusion(TrainerBase): |
| def _validate_configuration(self): |
| super()._validate_configuration() |
| assert self.config.sampling.noise_removal in { |
| 'none', 'ancestral', 'greedy'} |
| assert self.loss_type in {'elbo', 'low_var'} |
| if self.config.sampling.noise_removal == 'greedy': |
| assert self.sampler != 'analytic' |
| assert self.parameterization in {'mean', 'subs'} |
|
|
| def _process_model_input(self, x0, valid_tokens): |
| return x0, None, valid_tokens |
|
|
| def _process_sigma(self, sigma): |
| assert sigma.ndim == 2 |
| sigma = sigma.mean(-1).squeeze() |
| if sigma.ndim == 0: |
| sigma = sigma.unsqueeze(0) |
| if not self.time_conditioning: |
| sigma = torch.zeros_like(sigma) |
| assert sigma.ndim == 1, sigma.shape |
| return sigma |
|
|
| def _sample_t(self, n, accum_step): |
| if accum_step is not None: |
| |
| batch_dim = n |
| n = self.config.loader.global_batch_size |
| _eps_t = torch.rand(n, device=self.device) |
| if self.antithetic_sampling: |
| offset = torch.arange(n, device=self.device) / n |
| _eps_t = (_eps_t / n + offset) % 1 |
| t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps |
| if accum_step is not None: |
| t = t.chunk(self.trainer.num_nodes)[self.trainer.node_rank] |
| t = t.chunk(self.trainer.num_devices)[self.trainer.local_rank] |
| t = t.chunk(self.trainer.accumulate_grad_batches)[ |
| accum_step] |
| |
| t = t[:batch_dim] |
| return t |
|
|
| def _sigma_from_alphat(self, alpha_t): |
| return -torch.log(alpha_t) |
|
|
| def _reconstruction_loss(self, x0): |
| t0 = torch.zeros(1, x0.shape[0], dtype=self.dtype, |
| device=self.device) |
| sigma_t0 = self._sigma_from_alphat(self.noise(t0)[1]) |
| model_output_t0 = self.forward(x0, sigma_t0) |
| return - torch.gather(input=model_output_t0, |
| dim=-1, |
| index=x0[:, :, None]).squeeze(-1) |
|
|
| def nll_per_token(self, model_output, xt, x0, alpha_t, |
| dalpha_t, low_var): |
| raise NotImplementedError |
|
|
| def nll(self, x0, labels, output_tokens, |
| current_accumulation_step=None, train_mode=False): |
| del output_tokens |
| t = self._sample_t(x0.shape[0], |
| current_accumulation_step) |
| assert t.shape[0] == x0.shape[0] |
| if self.T > 0: |
| t = (t * self.T).to(torch.int) |
| t = t / self.T |
| |
| t += (1 / self.T) |
| |
| dalpha_t, alpha_t = self.noise(t) |
| alpha_t = alpha_t.unsqueeze(-1) |
| dalpha_t = dalpha_t.unsqueeze(-1) |
| assert alpha_t.ndim == 2 |
| assert dalpha_t.ndim == 2 |
| sigma = self._sigma_from_alphat(alpha_t) |
|
|
| xt = self.q_xt(x0, alpha_t) |
| |
| if self.class_conditional: |
| assert labels is not None |
| rand = torch.rand(size=labels.shape, dtype=torch.float32, |
| device=self.device) |
| |
| labels = torch.where(rand < self.class_cond_dropout, |
| self.num_classes, labels) |
| else: |
| assert labels is None |
| log_x_theta = self.forward(xt, sigma=sigma, labels=labels) |
| utils.print_nans(log_x_theta, 'model_output') |
| return self.nll_per_token( |
| log_x_theta=log_x_theta, |
| xt=xt, |
| x0=x0, |
| alpha_t=alpha_t, |
| dalpha_t=dalpha_t, |
| low_var=train_mode and self.loss_type == 'low_var') |
|
|
| def _get_score(self, **kwargs): |
| del kwargs |
| raise NotImplementedError |
|
|
| def _denoiser_update(self, x, t): |
| raise NotImplementedError |
|
|
| def _analytic_update(self, x, t, dt): |
| raise NotImplementedError |
|
|
| def _posterior_from_x0(self, x0, xt, alpha_s, alpha_t): |
| """From the clean x0, or denoiser predictions, implement |
| the mathematical expression for q(z_s | z_t, x)""" |
| raise NotImplementedError |
|
|
| def _forward_process(self, q_x0, alpha_s): |
| """Apply forward noising: q(x_s | x_0), where x_0 is a |
| K-dimensional vector""" |
| raise NotImplementedError |
|
|
| def _get_ancestral_posterior(self, xt, sigma, labels, |
| alpha_s, alpha_t, p_x0): |
| """Top-level call from generate_samples. Compute the |
| standard posterior sampling distribution from the |
| noisy sequence xt""" |
| gamma = self.config.sampling.guid_weight |
| |
| |
| if not self.class_conditional or gamma in (None, 0.0): |
| if labels is not None: |
| labels = torch.full_like(labels, self.num_classes) |
| return self._get_posterior_from_xt(xt, sigma, labels, |
| alpha_s, alpha_t, p_x0) |
| elif gamma == 1.0: |
| |
| assert labels is not None |
| return self._get_posterior_from_xt(xt, sigma, labels, |
| alpha_s, alpha_t, p_x0) |
| else: |
| |
| |
| return self._get_guided_posterior_from_xt(self, xt, |
| sigma, labels, gamma, alpha_s, alpha_t, p_x0) |
|
|
| def _get_posterior_from_xt(self, xt, sigma, labels, alpha_s, |
| alpha_t, p_x0=None): |
| """From xt, compute a single denoiser predictions, cast |
| to correct dtype, and call _posterior_from_x0.""" |
| if p_x0 is None: |
| log_x0_pred = self.forward(xt, sigma, labels) |
|
|
| if self.config.sampling.use_float64: |
| log_x0_pred = log_x0_pred.to(torch.float64) |
|
|
| if self.p_nucleus < 1: |
| log_x0_pred = utils.top_k_top_p_filtering( |
| log_x0_pred, top_p=self.p_nucleus) |
| p_x0 = log_x0_pred.exp() |
|
|
| return p_x0, self._posterior_from_x0(x0=p_x0,xt=xt, |
| alpha_s=alpha_s, alpha_t=alpha_t) |
|
|
| def _get_guided_posterior_from_xt(self, xt, sigma, labels, |
| gamma, alpha_s, alpha_t, p_x0=None): |
| """From xt, combine the class cond / uncond predictions |
| of the denoiser. Call _get_posterior_from_xt.""" |
| |
| if p_x0 is None: |
| p_x0_cond = p_x0_uncond = None |
| else: |
| p_x0_cond, p_x0_uncond = p_x0 |
| |
| p_x0_cond, cond_posterior = self._get_posterior_from_xt( |
| xt, sigma, labels, alpha_s, alpha_t, p_x0_cond) |
| log_cond_posterior = cond_posterior.log() |
| |
| |
| p_x0_uncond, uncond_posterior = self._get_posterior_from_xt( |
| xt, sigma, torch.full_like(labels, self.num_classes), |
| alpha_s, alpha_t, p_x0_uncond) |
| log_uncond_posterior = uncond_posterior.log() |
|
|
| un_normalized_posterior = gamma * log_cond_posterior \ |
| + (1 - gamma) * log_uncond_posterior |
| |
| |
| is_inf_mask = torch.logical_or( |
| log_cond_posterior.isinf(), log_uncond_posterior.isinf()) |
| un_normalized_posterior[is_inf_mask] = self.neg_infinity |
| |
| return ((p_x0_cond, p_x0_uncond), |
| un_normalized_posterior.softmax(-1)) |
|
|
| def _ancestral_update(self, x, t, labels, dt, p_x0=None, |
| noise_removal_step=False): |
| _, alpha_t = self.noise(t) |
| if noise_removal_step: |
| alpha_s = torch.ones_like(alpha_t) |
| else: |
| _, alpha_s = self.noise(t - dt) |
| assert alpha_t.ndim == 2 |
|
|
| sigma = self._sigma_from_alphat(alpha_t) |
| assert alpha_t.ndim == 2, f'{alpha_t.ndim=}' |
|
|
| p_x0, q_xs = self._get_ancestral_posterior(x, sigma, |
| labels, alpha_s, alpha_t, p_x0) |
|
|
| return p_x0, sample_categorical(q_xs) |
|
|
| def _psi_update(self, x, t, labels, dt, kappa, p_x0=None, |
| noise_removal_step=False): |
| _, alpha_t = self.noise(t) |
| if noise_removal_step: |
| alpha_s = torch.ones_like(alpha_t) |
| else: |
| _, alpha_s = self.noise(t - dt) |
| alpha_0 = torch.ones_like(alpha_t) |
| sigma = self._sigma_from_alphat(alpha_t) |
|
|
| |
| p_x0, q_xs = self._get_ancestral_posterior( |
| x, sigma, labels, alpha_s, alpha_t, p_x0) |
|
|
| |
| _, q_x0 = self._get_ancestral_posterior( |
| x, sigma, labels, alpha_0, alpha_t, p_x0) |
|
|
| |
| pc_q_xs = self._forward_process(q_x0, alpha_s) |
|
|
| q_sample = kappa * q_xs + (1 - kappa) * pc_q_xs |
| return p_x0, sample_categorical(q_sample) |
|
|
| def _get_sampling_time_profile(self, eps, num_steps): |
| profile = self.config.sampling.psi.time_profile |
| num_steps += 1 |
| if profile == 'linear' \ |
| or self.config.sampling.predictor != 'psi': |
| |
| return torch.linspace(1, eps, num_steps) |
| if not profile.startswith('linear-constant-linear'): |
| raise ValueError(profile) |
| c = float(profile.split('-')[3]) |
| if 'inv' in profile: |
| c = self.noise.get_t_for_alpha(c) |
| psi_cfg = self.config.sampling.psi |
| n_hi = round(psi_cfg.high_frac * num_steps) |
| n_mid = round(psi_cfg.middle_frac * num_steps) |
| return torch.cat([ |
| torch.linspace(1, c, n_hi), |
| torch.full((n_mid,), c), |
| torch.linspace(c, eps, num_steps - n_hi - n_mid)]) |
|
|
| def _mode_to_psi_kappas(self, mode, timesteps): |
| n = len(timesteps) |
| if mode == 'pure-posterior': |
| return torch.ones(n) |
| if mode == 'pure-pc': |
| return torch.zeros(n) |
|
|
| eta = float(mode.split('-')[-1]) |
| if (mode.startswith('constant-') |
| and not mode.startswith('constant-remdm')): |
| return torch.full((n,), eta) |
|
|
| |
| _, all_alphas = self.noise(timesteps) |
| alpha_t, alpha_s = all_alphas[:-1], all_alphas[1:] |
| eta_t = torch.tensor([eta]) |
| if mode.startswith('max-capped-'): |
| sigma = torch.minimum( |
| eta_t.expand_as(alpha_t), (1 - alpha_s) / alpha_t) |
| sigma = torch.where(alpha_t == 0, eta, sigma) |
| elif mode.startswith('max-rescale-'): |
| sigma_max = torch.minimum( |
| eta_t.expand_as(alpha_t), (1 - alpha_s) / alpha_t) |
| sigma = torch.where(alpha_t > 0, sigma_max, 1) * eta |
| elif mode.startswith('constant-remdm'): |
| sigma = eta_t |
| else: |
| raise ValueError(mode) |
| kappas = torch.clip(1 - sigma / (1 - alpha_s), 0, 1) |
| if len(kappas) > 0: |
| kappas = torch.cat([kappas, torch.ones(1)]) |
| return kappas |
|
|
| def _get_kappas(self, timesteps): |
| cfg = self.config.sampling.psi |
| n = len(timesteps) |
| n_hi = round(cfg.high_frac * n) |
| n_mid = round(cfg.middle_frac * n) |
| kappas = torch.cat([ |
| self._mode_to_psi_kappas(cfg.high_mode, timesteps[:n_hi]), |
| self._mode_to_psi_kappas(cfg.middle_mode, |
| timesteps[n_hi:n_hi + n_mid]), |
| self._mode_to_psi_kappas(cfg.low_mode, |
| timesteps[n_hi + n_mid:])]) |
| assert (kappas >= 0).all() and (kappas <= 1).all() |
| assert len(kappas) == n |
| return kappas |
|
|
| @torch.no_grad() |
| def generate_samples(self, num_samples, labels=None, |
| num_steps=None, eps=1e-5): |
| """Generate samples from the model.""" |
| |
| if num_steps is None: |
| num_steps = self.config.sampling.steps |
| x = self.prior_sample(num_samples, self.num_tokens) |
| use_psi_sampler = self.config.sampling.predictor == 'psi' |
| timesteps = self._get_sampling_time_profile(eps, |
| num_steps) |
| if use_psi_sampler: |
| kappas = self._get_kappas(timesteps).to(self.device) |
|
|
| p_x0_cache = None |
| if labels is not None: |
| labels = labels.to(self.device) |
|
|
| for i in range(num_steps): |
| t = timesteps[i] * torch.ones( |
| x.shape[0], 1, device=self.device) |
| dt = timesteps[i] - timesteps[i+1] |
| if self.sampler == 'ancestral': |
| _, x = self._ancestral_update( |
| x=x, t=t, labels=labels, dt=dt, p_x0=None) |
| elif self.sampler == 'ancestral_cache': |
| p_x0_cache, x_next = self._ancestral_update( |
| x=x, t=t, labels=labels, dt=dt, p_x0=p_x0_cache) |
| if (not torch.allclose(x_next, x) |
| or self.time_conditioning): |
| |
| p_x0_cache = None |
| x = x_next |
| elif self.sampler == 'psi': |
| _, x = self._psi_update(x=x, t=t, kappa=kappas[i], |
| labels=labels, dt=dt, |
| p_x0=None) |
| elif self.sampler == 'analytic': |
| assert labels is None, 'class-conditional sampling ' \ |
| 'is not implemented with the analytic sampler' |
| x = self._analytic_update(x=x,t=t, dt=dt) |
| else: |
| raise ValueError(self.sampler) |
|
|
| t0 = timesteps[-1] * torch.ones(x.shape[0], 1, |
| device=self.device) |
| if self.config.sampling.noise_removal == 'ancestral': |
| if self.sampler == 'analytic': |
| x = self._denoiser_update(x=x, t=t0) |
| else: |
| _, x = self._ancestral_update(x=x, t=t0, labels=labels, |
| dt=None, p_x0=p_x0_cache, noise_removal_step=True) |
| elif self.config.sampling.noise_removal == 'greedy': |
| sigma = self._sigma_from_alphat(self.noise(t0)[1]) |
| x = self.forward(xt=x, sigma=sigma).argmax(dim=-1) |
| return x |
|
|
| @torch.no_grad |
| def _semi_ar_sampler( |
| self, n_samples, stride_length, num_strides, dt=0.001): |
| |
| ones = torch.ones(n_samples, dtype=self.dtype, |
| device=self.device) |
|
|
| num_steps = int(1 / dt) |
| sampling_steps = 0 |
| intermediate_tokens = [] |
| target = None |
| for _ in range(num_strides + 1): |
| p_x0_cache = None |
| x = self.prior_sample(n_samples, self.num_tokens) |
| if target is not None: |
| x[:, : -stride_length] = target |
| for i in range(num_steps + 1): |
| p_x0_cache, x_next = self._ancestral_update( |
| x=x, t=(1 - i * dt) * ones, dt=dt, p_x0=p_x0_cache) |
| if (not torch.allclose(x_next, x) |
| or self.time_conditioning): |
| p_x0_cache = None |
| sampling_steps += 1 |
| x = x_next |
| x = self.forward(x, 0 * ones).argmax(dim=-1) |
| intermediate_tokens.append( |
| x[:, :stride_length].cpu().numpy()) |
| target = x[:, stride_length:] |
| |
| intermediate_tokens.append(target.cpu().numpy()) |
| intermediate_text_samples = [] |
| sequence_lengths = (( |
| np.concatenate(intermediate_tokens, axis=1)[:, 1:] |
| == self.tokenizer.eos_token_id).cumsum(-1) == 0).sum(-1) |
| for i in range(2, len(intermediate_tokens) + 1): |
| intermediate_text_samples.append( |
| self.tokenizer.batch_decode( |
| np.concatenate(intermediate_tokens[:i], axis=1))) |
| return (sampling_steps, intermediate_text_samples, |
| sequence_lengths) |
|
|
| def restore_model_and_semi_ar_sample( |
| self, stride_length, num_strides, dt=0.001): |
| """Generate samples from the model.""" |
| |
| |
| self._eval_mode() |
| (sampling_steps, samples, |
| sequence_lengths) = self._semi_ar_sampler( |
| n_samples=self.config.loader.eval_batch_size, |
| stride_length=stride_length, |
| num_strides=num_strides, |
| dt=dt) |
| self._train_mode() |
| return sampling_steps, samples, sequence_lengths |
|
|
|
|
| class AbsorbingState(Diffusion): |
| def __init__(self, config, tokenizer): |
| |
| |
| |
| |
| |
| vocab_size = tokenizer.vocab_size |
| if (not hasattr(tokenizer, 'mask_token') |
| or tokenizer.mask_token is None): |
| self.mask_index = vocab_size |
| vocab_size += 1 |
| else: |
| self.mask_index = tokenizer.mask_token_id |
| self.subs_masking = config.algo.subs_masking |
| super().__init__(config, tokenizer, |
| vocab_size=vocab_size) |
| self.save_hyperparameters() |
|
|
| def _validate_configuration(self): |
| super()._validate_configuration() |
| if self.parameterization in {'score', 'mean'}: |
| assert self.time_conditioning |
| assert not (self.parameterization == 'mean' |
| and self.T == 0) |
| if self.T > 0: |
| assert self.parameterization in {'mean', 'subs'} |
| if self.subs_masking: |
| assert self.parameterization == 'mean' |
|
|
| def q_xt(self, x, alpha_t): |
| """Computes the noisy sample xt. |
| |
| Args: |
| x: int torch.Tensor with shape (batch_size, |
| diffusion_model_input_length), input. |
| alpha_t: float torch.Tensor with shape (batch_size, 1). |
| """ |
| move_indices = torch.rand( |
| * x.shape, device=x.device) < 1 - alpha_t |
| xt = torch.where(move_indices, self.mask_index, x) |
| if self.ignore_bos: |
| xt[:, 0] = x[:, 0] |
| return xt |
|
|
| def prior_sample(self, *batch_dims): |
| return self.mask_index * torch.ones( |
| * batch_dims, dtype=torch.int64, device=self.device) |
|
|
| def _posterior_from_x0(self, x0, xt, alpha_s, alpha_t): |
| """From the clean x0, or denoiser predictions, implement |
| the mathematical expression for q(z_s | z_t, x)""" |
| assert x0.dtype == torch.float64, 'Requires float64 prec.' |
| |
| orig_mask = xt[:, :, None] != self.mask_index |
| orig_mask = orig_mask.expand(-1, -1, x0.shape[-1]) |
| orig_output_on_clean = x0[orig_mask] |
|
|
| q_xs = ((alpha_s - alpha_t) / (1 - alpha_t))[..., None] * x0 |
| q_xs[..., self.mask_index] = (1 - alpha_s) / (1 - alpha_t) |
| q_xs[orig_mask] = orig_output_on_clean |
| return q_xs |
|
|
| def _forward_process(self, x0, alpha_s): |
| out = alpha_s[..., None] * x0 |
| out[..., self.mask_index] = 1 - alpha_s |
| return out |
|
|
| def _staggered_score(self, score, dsigma): |
| score = score.clone() |
| extra_const = (1 - dsigma.exp()) * score.sum(dim=-1) |
| score *= dsigma.exp()[:, None] |
| score[..., self.mask_index] += extra_const |
| return score |
|
|
| def _analytic_update(self, x, t, dt): |
| sigma_t = self._sigma_from_alphat(self.noise(t)[1]) |
| sigma_s = self._sigma_from_alphat(self.noise(t - dt)[1]) |
| dsigma = sigma_t - sigma_s |
| score = self._get_score(x, sigma_t) |
| if self.config.sampling.use_float64: |
| score = score.to(torch.float64) |
| stag_score = self._staggered_score(score, dsigma) |
| probs = stag_score * self._transp_transition(x, dsigma) |
| return sample_categorical(probs) |
|
|
| def _denoiser_update(self, x, t): |
| sigma = self._sigma_from_alphat(self.noise(t)[1]) |
| score = self._get_score(x, sigma) |
| if self.config.sampling.use_float64: |
| score = score.to(torch.float64) |
| stag_score = self._staggered_score(score, sigma) |
| probs = stag_score * self._transp_transition(x, sigma) |
| probs[..., self.mask_index] = 0 |
| samples = sample_categorical(probs) |
| return samples |
|
|
| def _transp_transition(self, i, sigma): |
| sigma = _unsqueeze(sigma, reference=i[..., None]) |
| edge = torch.exp(-sigma) * F.one_hot( |
| i, num_classes=self.vocab_size) |
| edge += torch.where(i == self.mask_index, |
| 1 - torch.exp(-sigma).squeeze(-1), |
| 0)[..., None] |
| return edge |
|
|
|
|
| class UniformState(Diffusion): |
| def _validate_configuration(self): |
| super()._validate_configuration() |
| assert self.time_conditioning |
| assert self.parameterization == 'mean' |
| if self.config.algo.name != 'distillation': |
| assert self.T == 0 |
|
|
| def _forward_process(self, x0, alpha_s): |
| return (alpha_s[..., None] * x0 |
| + (1 - alpha_s[..., None]) / self.vocab_size) |
|
|
| def q_xt(self, x, alpha_t): |
| """Computes the noisy sample xt. |
| |
| Args: |
| x: int torch.Tensor with shape (batch_size, |
| diffusion_model_input_length), input. |
| move_chance: float torch.Tensor with shape |
| (batch_size, 1). |
| """ |
| move_indices = torch.rand( |
| *x.shape, device=x.device) < 1 - alpha_t |
| uniform_tensor = torch.randint( |
| 0, self.vocab_size, x.shape, device=x.device) |
| xt = torch.where(move_indices, uniform_tensor, x) |
| if self.ignore_bos: |
| xt[:, 0] = x[:, 0] |
| return xt |
|
|
| def prior_sample(self, *batch_dims): |
| return torch.randint( |
| 0, self.vocab_size, batch_dims, dtype=torch.int64, |
| device=self.device) |
