1.1.7. Discrete Denoising Diffusion Probabilistic Models (D3PMs)

While Cold Diffusion (CD) showed that diffusion could be performed with more general corruption patterns, previous formulations focused exclusively on continuous variables, adding Gaussian noise to corrupt data. In this context, Discrete Denoising Diffusion Probabilistic Models (D3PMs) (Austin et al. 2021) introduced a diffusion paradigm for discrete data, enabling the application to text, graphs and more general categorical data.

Let a single data component at diffusion step \(t\) be represented by a one-hot vector \(\boldsymbol x_t \in \{0, 1\}^K\) with \(K \in \mathbb N\) discrete states. The forward diffusion process in D3PM was defined as: \[ q(\boldsymbol x_t \mid \boldsymbol x_{t - 1}) = \boldsymbol x_t^T \boldsymbol Q_t \boldsymbol x_{t - 1}, \tag{1}\] where \(\boldsymbol Q_t \in \mathbb R^{K\times K}\) is the Markov transition matrix. The components of this matrix define the probability of state \(j\) transitioning to state \(i\), that is, \(q(\boldsymbol x_t = i \mid \boldsymbol x_t = j)\). This definition enables sampling at any time step \(t\) given a clean data \(\boldsymbol x_0\): \[ q(\boldsymbol x_t \mid \boldsymbol x_0) = \boldsymbol x_t^T \boldsymbol{\overline{Q}}_t \boldsymbol x_0, \tag{2}\] where \(\boldsymbol{\overline{Q}}_t = \boldsymbol Q_1 \boldsymbol Q_2\dots\boldsymbol Q_t\). The transition matrices defined different corruption processes, and were chosen to match the structure of the input data. For example, non-ordinal data could be modelled using an uniform transition state, where each component can change to any other state. Another proposed design was the absorbing state transition, where the variables can transition to a \([\text{MASK}]\) state, but can not transition out of it. This process is applied independently to each component of the input data.

Given the diffusion process, a Neural Network \(\boldsymbol{\tilde{p}}_{\boldsymbol \theta} : \{0, 1\}^K \times \mathbb R \to [0, 1]^K\) parameterized by \(\boldsymbol \theta\) is trained to reverse the corruption process. The network is trained to output a categorical distribution, such that it is able to identify which state each component was initially. The reverse transition step is then computed by marginalizing over this prediction: \[ p_{\boldsymbol \theta} (\boldsymbol x_{t - 1} \mid \boldsymbol x_t) = \sum_{\boldsymbol{\tilde{x}}_0 \in \{0, 1\}^K} q(\boldsymbol x_{t-1} \mid \boldsymbol x_t, \boldsymbol{\tilde{x}}_0) \left( \boldsymbol{\tilde{p}}_{\boldsymbol \theta}(\boldsymbol x_t, t)^\top \boldsymbol{\tilde{x}}_0 \right), \tag{3}\] where the true forward posterior \(q(\boldsymbol x_{t-1} \mid \boldsymbol x_t, \boldsymbol x_0)\) is calculated using Bayes’ rule and the Markov transition matrices: \[ q(\boldsymbol x_{t-1} \mid \boldsymbol x_t, \boldsymbol x_0) = \frac{q(\boldsymbol x_t \mid \boldsymbol x_{t-1}) q(\boldsymbol x_{t-1} \mid \boldsymbol x_0)}{q(\boldsymbol x_t \mid \boldsymbol x_0)} = \frac{(\boldsymbol x_t^\top \boldsymbol Q_t \boldsymbol x_{t-1})(\boldsymbol x_{t-1}^\top \boldsymbol{\overline{Q}}_{t-1} \boldsymbol x_0)}{\boldsymbol x_t^\top \boldsymbol{\overline{Q}}_t \boldsymbol x_0}. \tag{4}\]

To optimize the network, D3PM minimized the Variational Lower Bound (VLB) adapted for discrete spaces. The training objective is to minimize the expected Kullback-Leibler (KL) divergence between the true discrete posterior and the parameterized reverse step: \[ \mathcal{L}(\boldsymbol \theta) = \mathbb{E}_{t, \boldsymbol{x}_0, \boldsymbol{x}_t} \left[ D_{\text{KL}} \left( q(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0) \parallel p_{\boldsymbol{\theta}}(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t) \right) - \lambda\log p_{\boldsymbol{\theta}}(\boldsymbol{x}_0 \mid \boldsymbol{x}_t) \right], \tag{5}\] where \(\lambda \in \mathbb R_{\ge 0}\) is a weighting parameter for the cross-entropy between the original data and the predicted output of the network.

Code

Importing the packages.

import torch

Creating a small state space.

chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ "

char2int = {c: i for i, c in enumerate(chars)}
char2int["[MASK]"] = len(chars)

int2char = {i: c for i, c in enumerate(chars)}
int2char[len(chars)] = "[MASK]"

mask_id = char2int["[MASK]"]

K = len(char2int)

Defining an absorbing and uniform transition matrices.

def get_absorbing_Q(beta, K, mask_id):
    Q = torch.eye(K) * (1 - beta)
    Q[:, mask_id] += beta

    Q[mask_id, :] = 0.0
    Q[mask_id, mask_id] = 1.0

    return Q

def get_uniform_Q(beta, K):
    Q = torch.ones(K, K) * (beta / (K - 1))
    Q.fill_diagonal_(1 - beta)

    return Q

Simulating a forward process.

text = "DISCRETE DIFFUSION"
x_0 = torch.tensor([char2int[c] for c in text])
T = 5
betas = torch.linspace(0.05, 0.5, T)

def simulate_corruption(x_start, matrix_type="absorbing"):
    x_t = x_start.clone()
    print(f"t=0 : {''.join([int2char[i.item()] for i in x_t])}")

    for t in range(T):
        beta = betas[t].item()

        if matrix_type == "absorbing":
            Q_t = get_absorbing_Q(beta, K, mask_id)
        else:
            Q_t = get_uniform_Q(beta, K)

        transition_probs = Q_t[x_t]

        x_t = torch.multinomial(transition_probs, num_samples=1).squeeze()

        current_str = "".join([int2char[i.item()] for i in x_t])
        current_str = current_str.replace("[MASK]", "_")

        print(f"t={t+1:<2}: {current_str}")

print("--- Absorbing Transitions ---")
simulate_corruption(x_0, matrix_type="absorbing")

print("\n--- Uniform Transitions ---")
simulate_corruption(x_0, matrix_type="uniform")

--- Absorbing Transitions ---
t=0 : DISCRETE DIFFUSION
t=1 : DISCRETE DIFFUSION
t=2 : DISCR__E DIFFUSION
t=3 : DI__R___ DIF_USION
t=4 : DI__R___ ______I_N
t=5 : _I______ _________

--- Uniform Transitions ---
t=0 : DISCRETE DIFFUSION
t=1 : DYQCRETEKGIFFUSION
t=2 : TYQCRETDKGIFFUSION
t=3 : TXQCRETSKNITFUGIOO
t=4 : TJSIK_FSKNIDFUGQC_
t=5 : T PIK_DQMNESFNAQW_

Source: d3pm.ipynb

References

Austin, Jacob, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. 2021. “Structured Denoising Diffusion Models in Discrete State-Spaces.” Advances in Neural Information Processing Systems 34: 17981–93.