| from __future__ import annotations |
|
|
| import torch |
|
|
|
|
| class MultiMaskDiffusion: |
| """Forward corruption from the MultiMDM paper. |
| |
| p_t(x_t|x_0) = alpha_t * delta(x_t=x_0) |
| + (1-alpha_t) * [beta_t * delta(x_t=m_x0) |
| + (1-beta_t) * Uniform(M masks)]. |
| """ |
|
|
| def __init__( |
| self, |
| vocab_size: int, |
| clean_token_start: int, |
| num_masks: int, |
| mask_token_start: int, |
| pad_id: int = 0, |
| ): |
| self.vocab_size = vocab_size |
| self.clean_token_start = clean_token_start |
| self.num_masks = num_masks |
| self.mask_token_start = mask_token_start |
| self.pad_id = pad_id |
|
|
| def designated_mask(self, clean_ids: torch.Tensor) -> torch.Tensor: |
| if torch.any(clean_ids < self.clean_token_start): |
| raise ValueError("designated_mask expects clean vocabulary token IDs") |
| return self.mask_token_start + ( |
| (clean_ids - self.clean_token_start) % self.num_masks |
| ) |
|
|
| @staticmethod |
| def alpha(t: torch.Tensor) -> torch.Tensor: |
| return (1.0 - t.clamp(0.0, 1.0)).clamp(0.0, 1.0) |
|
|
| @staticmethod |
| def beta(t: torch.Tensor) -> torch.Tensor: |
| return (1.0 - t.clamp(0.0, 1.0)).clamp(0.0, 1.0) |
|
|
| def q_sample( |
| self, |
| x0: torch.Tensor, |
| t: torch.Tensor, |
| attention_mask: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| if t.ndim == 0: |
| t = t.expand(x0.shape[0]) |
| if t.ndim != 1 or t.shape[0] != x0.shape[0]: |
| raise ValueError("t must be scalar or have shape [batch]") |
| view_shape = [x0.shape[0]] + [1] * (x0.ndim - 1) |
| alpha = self.alpha(t).view(*view_shape) |
| beta = self.beta(t).view(*view_shape) |
|
|
| valid = x0.ge(self.clean_token_start) |
| if attention_mask is not None: |
| valid = valid & attention_mask.bool() |
|
|
| keep_clean = torch.rand_like(x0, dtype=torch.float32) < alpha |
| use_designated = torch.rand_like(x0, dtype=torch.float32) < beta |
| uniform_masks = self.mask_token_start + torch.randint( |
| self.num_masks, x0.shape, device=x0.device |
| ) |
| safe_clean = torch.where(valid, x0, torch.full_like(x0, self.clean_token_start)) |
| designated = self.designated_mask(safe_clean) |
| sampled_mask = torch.where(use_designated, designated, uniform_masks) |
| corrupted = torch.where(keep_clean, x0, sampled_mask) |
| return torch.where(valid, corrupted, x0) |
|
|
| def is_mask(self, ids: torch.Tensor) -> torch.Tensor: |
| return ids.ge(self.mask_token_start) & ids.lt( |
| self.mask_token_start + self.num_masks |
| ) |
|
|
