Method qt_reverse

diffusion_condition.py:434–464

Get q(x_{t+1} | x_t), the one-step posterior efficiently. Args: qt_plus_1: an array of floats specifying a distribution over p(x_0). t: t in q(x_{t+1} | x_t). return_logits: if True, return the output logits make_one_hot: if True, will convert q0 to fl

(self,
                   qt_plus_1,
                   t,
                   return_logits=False,
                   make_one_hot=False,
                   epsilon=1e-20
                   )

        return (1 - beta) * torch.eye(self.dim) + beta * self._get_mask()

    def qt_reverse(self,
                   qt_plus_1,
                   t,
                   return_logits=False,
                   make_one_hot=False,
                   epsilon=1e-20
                   ):
        """Get q(x_{t+1} | x_t), the one-step posterior efficiently.
        Args:
          qt_plus_1: an array of floats specifying a distribution over p(x_0).
          t: t in q(x_{t+1} | x_t).
          return_logits: if True, return the output logits
          make_one_hot: if True, will convert q0 to floats if needed.
          epsilon: a small number to normalize logits conversion with, if needed.
        Returns:
          q(x_{t+1} | x_t).
        """
        if make_one_hot:
            assert qt_plus_1.dtype == torch.int64
            qt_plus_1 = torch.nn.functional.one_hot(qt_plus_1, num_classes=self.dim)

        beta = self.schedule(t)
        qtpls1_at_mask = qt_plus_1[Ellipsis, self.tokenizer.mask_token_id: self.tokenizer.mask_token_id + 1]
        non_mask_prob0 = (1 - beta) * qt_plus_1[Ellipsis, :self.tokenizer.mask_token_id] + beta * qtpls1_at_mask
        non_mask_prob1 = (1 - beta) * qt_plus_1[Ellipsis, self.tokenizer.mask_token_id + 1:] + beta * qtpls1_at_mask
        prob_at_time_t = torch.cat((non_mask_prob0, qtpls1_at_mask, non_mask_prob1), dim=-1)

        if return_logits:
            return torch.log(prob_at_time_t + epsilon)
        else:
            return prob_at_time_t

    def get_qt_given_q0(self,
                        q0,