In [1]:

---
execute:
  eval: false
  echo: true
  warning: false
format:
  html:
    code-fold: false
---

Importing the packages.

In [1]:

import torch

Creating a small state space.

In [14]:

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.

In [15]:

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.

In [17]:

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_