Question about discrete score understanding

[tulvgengenr]

Hello! I am very interested in your paper “Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution”, it is very excellent paper!

In my project, I would like to use the theory of this paper to optimize, but there is a question I would like to ask you. I would be very grateful for your answer.

The work that my project is based on uses discrete diffusion, who uses $Q^{absorb}$ for additive noise scheduling. For a token, it has a $1-e^{-sigma(t)}$ probability of becoming “mask_id” and a $e^{sigma(t)}$ probability of remaining the same. But the model outputs a probability distribution for each token in the sequence (same as the LLM), and the output result tensor shape is (B, L, N). How can I convert the model results to scores? Thus being able to calculate the LOSS directly using your code.

To my knowledge, I think I need to calculate the ratio of the probability that each position is predicted to be “mask_id” to the probability that it is predicted to be itself. For example the following code:

def p2score(pred).
    '''
    pred: [B, L, N]
    '''
    # Probability for mask (vocab_size-1) / itself
    pred_mask = pred[:, :, -1].repeat(1, 1, pred.shape[-1])
    score = pred_mask / pred
    return score

Then you can use your code to calculate the loss, here is your code:

    def score_entropy(self, score, sigma, x, x0):
        rel_ind = x == self.dim - 1
        esigm1 = torch.where(
            sigma < 0.5,
            torch.expm1(sigma),
            torch.exp(sigma) - 1
        )
        ratio = 1 / esigm1.expand_as(x)[rel_ind]
        other_ind = x0[rel_ind]
        # negative_term
        neg_term = ratio * torch.gather(score[rel_ind], -1, other_ind[... , None]).squeeze(-1)
        #positive term
        pos_term = score[rel_ind][:, :-1].exp().sum(dim=-1)
        # constant term
        const = ratio * (ratio.log() - 1)
        entropy = torch.zeros(*x.shape, device=x.device)
        entropy[rel_ind] += pos_term - neg_term + const
        return entropy

Is my understanding correct and I would like to get your reply!

[TediousBoredom]

def get_loss_fn(noise, graph, train, sampling_eps=1e-3, lv=False):

def loss_fn(model, batch, cond=None, t=None, perturbed_batch=None):
    """
    Batch shape: [B, L] int. D given from graph
    """

    if t is None:
        if lv:
            raise NotImplementedError("Yeah I gotta do this later")
        else:
            t = (1 - sampling_eps) * torch.rand(batch.shape[0], device=batch.device) + sampling_eps
        
    sigma, dsigma = noise(t)
    
    if perturbed_batch is None:
        perturbed_batch = graph.sample_transition(batch, sigma[:, None])

    log_score_fn = mutils.get_score_fn(model, train=train, sampling=False)
    log_score = log_score_fn(perturbed_batch, sigma)
    loss = graph.score_entropy(log_score, sigma[:, None], perturbed_batch, batch)

    loss = (dsigma[:, None] * loss).sum(dim=-1)

    return loss

return loss_fn

[yusowa0716]

Hello,

I have another question.

My understanding is that the goal of this function is to model the probability ratio $s_\theta(x_t, t)_y \approx \frac{p_t(y)}{p_t(x_t)}$. I've been trying to derive the final expression from the underlying principles of the absorbing state diffusion process.

My line of reasoning is as follows:

1. The probability of a token $x_0$ not being masked at a given noise level $\sigma(t)$ is the survival probability, which can be expressed as $p(x_t=x_0 | x_0) = e^{-\sigma(t)}$.
2. The probability of the original token $x_0$ transitioning to another specific token $y$ is proportional to the probability of being masked, $(1 - e^{-\sigma(t)})$, distributed according to the model's prediction of the original data, which we can denote as $p_\theta(x_0)$. This gives a transition probability of $p(x_t=y | x_0) = (1 - e^{-\sigma(t)}) \cdot \langle p_\theta(x_0), y \rangle$.
3. Taking the logarithm of the ratio of these probabilities, $\log \left( \frac{p(x_t=y | x_0)}{p(x_t=x_0 | x_0)} \right)$, seems to correspond to the first part of the expression in the code, related to the model's logits and the esigm1_log term.

However, I am having trouble understanding the origin of the final term in this line of code:

# From the _sedd_parameterization function
x = x - esigm1_log - np.log(x.shape[-1] - 1)

My specific question is:

What is the theoretical purpose of subtracting np.log(x.shape[-1] - 1)? This term does not appear in my derivation above. ?

Any insight or clarification on the mathematical justification for this term would be greatly appreciated.

Thank you!

[dingfangyu]

-np.log(x.shape[-1] - 1) is -log(vocabulary_size), approximately -10, it is designed for numerical issues.

Without it the score network will output log_score values about 0, while the first item of DSE loss is computed by log_score.exp().sum(along seq dim and vocab dim), which will be extremely large, the loss will be approximately seqlen * vocabsize, and the model cannot be optimized.

In contrast, if you add -10 to all of the log_score values, the DSE loss will be a safe value. You can regard it as a safe network initialization, where the initial output score values are all 1/V.