Discrete masked diffusion and finite-state reverse generation

Fix a finite nonempty clean vocabulary $V$, a mask symbol $m\mathrm{\notin }V$, integers $L,N\ge 1$, a grid $0=t_0<t_1<\cdots <t_N=1$, a survival schedule $1={\alpha }_0>{\alpha }_1>\cdots >{\alpha }_N=0$, and a probability vector $\mu$ on $V^L$ (nonnegative entries summing to one). Put $S=(V\cup \{m\}{)}^L$ and $r_j={\alpha }_j\mathrm{/}{\alpha }_{j-1}$ for $1\le j\le N$. On one coordinate, the step-$j$ transition from $v\in V$ stays at $v$ with probability $r_j$ and moves to $m$ with probability $1-r_j$; from $m$ it stays at $m$ with probability one. All other one-coordinate transitions have probability zero. The sequence-state transition from $z$ to $z^{\mathrm{\prime }}$ in $S$ is the product of these one-coordinate probabilities over the $L$ coordinates. Starting with $Z_0=X\sim \mu$, the conditional path mass is $\mathrm{Pr}(Z_0=z_0,\dots ,Z_N=z_N\mid X=x)={\mathbf{1}}_{\{z_0=x\}}{\prod }_{j=1}^N{\prod }_{i=1}^L{a}_j(z_{j-1,i},z_{j,i}),$ where $a_j$ is the stated one-coordinate transition rule. Thus, conditional on the clean sequence, the coordinate trajectories are independent. At each step, the next-coordinate states are independent conditional on the preceding sequence state, with the transition probabilities above; successive states of a coordinate need not be independent. The mask is absorbing.

The following one-coordinate diagram records the three transition formulas at step $j$; each arrow label is its transition probability. The clean-state loop has probability $r_N=0$ at the terminal step $j=N$, so it is absent from the support there.

Measure kernel and probability kernel: a probability kernel assigns a probability measure to each state and has measurable evaluation functions; on the finite power-set space above this is exactly a row-stochastic transition matrix.

Let $S$ be a finite nonempty set and let $Q$ and $P$ be probability kernels on $S$, meaning that each row is a probability vector. For probability vectors $u,v$ on $S$, use $\mathrm{TV}(u,v)={\sup }_{A\subseteq S}\mathrm{\mid }u(A)-v(A)\mathrm{\mid }$, where $u(A)={\sum }_{x\in A}u(x)$. If ${\sup }_{x\in S}\mathrm{TV}(Q(x,\cdot ),P(x,\cdot ))\le \epsilon$ for $0\le \epsilon \le 1$, then every probability vector $\mu$ on $S$ satisfies $\mathrm{TV}(\mu Q,\mu P)\le \epsilon ,$ where $(\mu Q)(A)={\sum }_{x\in S}\mu (x)Q(x,A)$ and similarly for $\mu P$.

Total variation distance for probability laws: total variation is the supremum of absolute event-probability differences.

Measure kernel and probability kernel: each kernel row is a probability measure; on finite power-set spaces this is a row-stochastic matrix.

Let $r\ge 1$, let $S_1,\dots ,S_r$ be finite nonempty sets, and let $u_k,v_k$ be probability vectors on $S_k$. Put $U={⨂}_{k=1}^r{u}_k$ and $V={⨂}_{k=1}^r{v}_k$ on $S_1\times \cdots \times S_r$. With $D$ the finite directional KL divergence, zero source masses contributing zero and positive source mass against zero reference mass contributing $+\mathrm{\infty }$, one has $D(U\mathrm{\parallel }V)={\sum }_{k=1}^{r}D(u_k\mathrm{\parallel }{v}_k),$ with equality also when one or both sides are $+\mathrm{\infty }$.

Finite Kullback–Leibler divergence: $D(P\mathrm{\parallel }Q)={\sum }_{x:p_x>0}{p}_x\log (p_x\mathrm{/}{q}_x)$, with $+\mathrm{\infty }$ for positive source mass where the reference mass is zero.

Let $S$ be a finite nonempty state space and let $N\ge 1$. For each $j=1,\dots ,N$, let $Q_j$ and $P_j$ be probability kernels on $S$, used in reverse order from level $j$ to level $j-1$. Starting from probability vectors ${\mu }_N,{\nu }_N$, define ${\mu }_{j-1}={\mu }_j{Q}_j$ and ${\nu }_{j-1}={\nu }_j{P}_j$. If $\mathrm{TV}({\mu }_N,{\nu }_N)\le {\epsilon }_{\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}}$ and ${\sup }_{z\in S}\mathrm{TV}(Q_j(z,\cdot ),P_j(z,\cdot ))\le {\epsilon }_j(1\le j\le N),$ where ${\epsilon }_{\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}},{\epsilon }_1,\dots ,{\epsilon }_N\in [0,1]$, then $\mathrm{TV}({\mu }_0,{\nu }_0)\le \min \{1,{\epsilon }_{\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}}+{\sum }_{j=1}^N{\epsilon }_j\}\mathrm{.}$ Here $\mathrm{TV}(u,v)={\sup }_{A\subseteq S}\mathrm{\mid }u(A)-v(A)\mathrm{\mid }$.

Total variation distance for probability laws: for probability laws, total variation is the supremum of absolute event-probability differences.

Measure kernel and probability kernel: a probability kernel has probability-law rows; on $S$ its action is matrix multiplication.

For a finite absorbing-mask chain with finite clean vocabulary, mask symbol, sequence length, strict grid, and survival schedule as specified above, each sequence-state transition row is a probability vector, and the displayed joint path mass $\mu (x){\mathbf{1}}_{\{z_0=x\}}{\prod }_{j=1}^N{\prod }_{i=1}^L{a}_j(z_{j-1,i},z_{j,i})$ sums to one over $x\in V^L$ and $z_0,\dots ,z_N\in S$. Thus it defines a probability law on the finite clean-sequence and path space.

Finite masking diffusion chain: fixes $V,m,L,N,\alpha ,\mu$, the one-coordinate rows, product sequence kernels, and the path-mass formula.

For the finite masking chain with state space $S=(V\cup \{m\}{)}^L$ and survival schedule ${\alpha }_0,\dots ,{\alpha }_N$, take probability vectors $h_{i,j}(z)\in \mathrm{\Delta }(V)$ for every position $i\in \{1,\dots ,L\}$, reverse step $j\in \{1,\dots ,N\}$, and current state $z\in S$, where $\mathrm{\Delta }(V)=\{p:V\to [0,1]:{\sum }_{v\in V}p(v)=1\}$. Initialize $Z_N=m^L$. For $j=N,N-1,\dots ,1$, set $c_j=({\alpha }_{j-1}-{\alpha }_j)\mathrm{/}(1-{\alpha }_j)$ and update all positions from the same current state $z=Z_j$ as follows, independently across positions conditional on $z$: if $z_i\mathrm{\ne }m$, set $Z_{j-1,i}=z_i$; if $z_i=m$ and $c_j=1$, draw $Z_{j-1,i}$ from $h_{i,j}(z)$; if $z_i=m$ and $c_j<1$, set $Z_{j-1,i}=m$ when $U_{j,i}^{\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}}\ge c_j$, and otherwise draw the clean token from $h_{i,j}(z)$ using $U_{j,i}^{\mathrm{t}\mathrm{o}\mathrm{k}\mathrm{e}\mathrm{n}}$. The resulting one-coordinate row is $p_{i,j}(w\mid z)={\mathbf{1}}_{\{w=z_i\}}$ when $z_i\mathrm{\ne }m$, and, when $z_i=m$, has mass $1-c_j$ at $m$ and mass $c_j{h}_{i,j}(z)(v)$ at each $v\in V$; the sequence-state kernel is $P_j(z,z^{\mathrm{\prime }})={\prod }_{i=1}^L{p}_{i,j}(z_i^{\mathrm{\prime }}\mid z)$. Return $Z_0$ after the finite loop. The mathematical random source is a finite product of mutually independent uniform $[0,1]$ variables $U_{j,i}^{\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}},U_{j,i}^{\mathrm{t}\mathrm{o}\mathrm{k}\mathrm{e}\mathrm{n}}$, one pair for each step and position; an implementation seed is valid only insofar as its supplied random kernel has this law. For inverse-CDF token draws, fix any order $V=(v_1,\dots ,v_K)$, choose the first $v_k$ whose cumulative mass strictly exceeds $U_{j,i}^{\mathrm{t}\mathrm{o}\mathrm{k}\mathrm{e}\mathrm{n}}$, and return $v_K$ if that draw equals $1$. At $j=1$, $c_1=1$, so every remaining mask is unmasked.

Finite masking diffusion chain: supplies the finite vocabulary, mask, length, grid, and survival schedule.

Measure kernel and probability kernel: a probability kernel is a family of probability-law rows with measurable evaluation functions; the finite $P_j$ above is a product of such rows.

Let $V$ be finite and nonempty, $m\mathrm{\notin }V$, $L,N\ge 1$, $0=t_0<\cdots <t_N=1$, and $1={\alpha }_0>\cdots >{\alpha }_N=0$. For each $i,j,z$, let $h_{i,j}(z)$ be a probability vector on $V$, put $c_j=({\alpha }_{j-1}-{\alpha }_j)\mathrm{/}(1-{\alpha }_j)$, and define the reverse coordinate rows to copy an unmasked token and, from $m$, assign mass $1-c_j$ to $m$ and mass $c_j{h}_{i,j}(z)(v)$ to each $v\in V$. Their product rows are probability vectors. Composing these $N$ kernels from the point mass at $m^L$ therefore gives a normalized finite law, and its final state lies in $V^L$ because $c_1=1$ removes every remaining mask.

Finite-grid masked reverse sampler: defines $c_j$, the position-wise reverse rows, their product kernel, and the start state $m^L$.

Measure kernel and probability kernel: a probability kernel has probability-law rows.

Let $V$ be a finite nonempty vocabulary, let $m\mathrm{\notin }V$, let $L,N\ge 1$, and let $1={\alpha }_0>{\alpha }_1>\cdots >{\alpha }_N=0$ be the survival schedule of the finite masking chain. Fix a clean sequence $x\in V^L$ and grid indices $0\le s<t\le N$. For every state $z_t\in (V\cup \{m\}{)}^L$ with $\mathrm{Pr}(Z_t=z_t\mid X=x)>0$, the conditional law of $Z_s$ given $X=x,Z_t=z_t$ is a product over coordinates. At coordinate $i$, if $z_{t,i}\in V$, then $Z_{s,i}=z_{t,i}$ with probability one. If $z_{t,i}=m$, then, with $c_{s,t}=({\alpha }_s-{\alpha }_t)\mathrm{/}(1-{\alpha }_t)$, the only possible values are $x_i$ and $m$, with probabilities $\mathrm{Pr}(Z_{s,i}=x_i\mid X=x,Z_{t,i}=m)=c_{s,t},\mathrm{Pr}(Z_{s,i}=m\mid X=x,Z_{t,i}=m)=\frac{1-{\alpha }_s}{1-{\alpha }_t}\mathrm{.}$ The denominator is positive because $t\ge 1$ implies ${\alpha }_t<1$; at $s=0$ the mask posterior has zero mass.

Finite masking diffusion chain: conditional on $X=x$, coordinates and masking draws are independent, clean tokens never change, and $m$ is absorbing.

Fix a clean string $x\in V^L$ and let $q_x$ be the finite forward path law of the absorbing-mask chain. For every path $(z_0,\dots ,z_N)$ with $q_x(z_0,\dots ,z_N)>0$, the reversed path probability factors as $q_x(z_0,\dots ,z_N)=q_x(Z_N=z_N){\prod }_{j=N}^1{q}_x(Z_{j-1}=z_{j-1}\mid Z_j=z_j)\mathrm{.}$ Each factor on the right is the true one-step posterior conditional on the clean string $x$; the identity is asserted only on positive-probability paths.

Finite masking diffusion chain: the conditional forward path mass is the product of its one-step transition rows.

Posterior of a finite masking chain: gives the finite conditional law of each preceding state given $x$ and the current state.

Let $V$ be a finite nonempty vocabulary, let $m\mathrm{\notin }V$, let $L,N\ge 1$, and let $1={\alpha }_0>\cdots >{\alpha }_N=0$. Draw $X=(X_1,\dots ,X_L)$ from any probability law on $V^L$. Conditional on $X$, at step $j$ each clean coordinate survives with probability $r_j={\alpha }_j\mathrm{/}{\alpha }_{j-1}$ using fresh Bernoulli innovations independent across coordinates and steps, and otherwise becomes $m$; an already masked coordinate remains masked. Thus $Z_j$ is the state after step $j$. For $j=1,\dots ,N$, put $c_j=({\alpha }_{j-1}-{\alpha }_j)\mathrm{/}(1-{\alpha }_j),d_j=(1-{\alpha }_{j-1})\mathrm{/}(1-{\alpha }_j)\mathrm{.}$ Let $h_{i,j}:S\to \mathrm{\Delta }(V)$ be any predictors, where $S=(V\cup \{m\}{)}^L$ and $\mathrm{\Delta }(V)$ is the set of probability vectors on $V$. Define the learned reverse kernel $P_j$ to have product rows: an unmasked coordinate is copied, while a masked coordinate assigns mass $d_j$ to $m$ and mass $c_j{h}_{i,j}(z)(v)$ to each $v\in V$. For every positive-probability current state $z_j$, the true conditional law $q(Z_{j-1}\mid Z_j=z_j,X)$ also has product rows: it copies every unmasked coordinate and, at a masked coordinate $i$, assigns mass $c_j$ to $X_i$ and $d_j$ to $m$. For $1\le i\le L$, let $M_{j,i}=\{Z_{j,i}=m\}$, which has probability $1-{\alpha }_j>0$, and let $P_{j,i}^{\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}}$ be the conditional law of $(Z_j,X_i)$ given $M_{j,i}$. View $h_{i,j}$ as a predictor in the finite supervised problem with feature $Z_j$, label $X_i$, action space $\mathrm{\Delta }(V)$, its singleton predictor class, and loss $\mathrm{\ell }(p,v)=-\log p(v)$, with $-\log 0=+\mathrm{\infty }$. Its population risk under $P_{j,i}^{\mathrm{m}\mathrm{a}\mathrm{s}\mathrm{k}}$ is $R_{j,i}=\mathbb{E}[\mathrm{\ell }(h_{i,j}(Z_j),X_i)\mid M_{j,i}]\mathrm{.}$ For finite-alphabet probability vectors $P=(p_a)$ and $Q=(q_a)$, write $D_{\mathrm{K}\mathrm{L}}(P\mathrm{\parallel }Q)={\sum }_{a:p_a>0}{p}_a\log (p_a\mathrm{/}{q}_a)$, with zero source terms set to zero and positive source mass against zero reference mass set to $+\mathrm{\infty }$. Define the extended-nonnegative conditional reverse-KL contribution ${\mathcal{L}}_j={\mathbb{E}}_{X,Z_j}\left[D_{\mathrm{K}\mathrm{L}}(q(Z_{j-1}\mid Z_j,X)\parallel P_j(Z_j,\cdot ))\right]\mathrm{.}$ Then ${\sum }_{j=1}^N{\mathcal{L}}_j={\sum }_{j=1}^N{\sum }_{i=1}^L({\alpha }_{j-1}-{\alpha }_j)R_{j,i}$ as an equality of extended nonnegative quantities. If a predictor assigns zero mass to the true token on a positive-probability masked event, both sides are $+\mathrm{\infty }$; finite cross-entropy is not assumed.

Finite Kullback–Leibler divergence: for finite probability vectors, $D(P\mathrm{\parallel }Q)={\sum }_{x:p_x>0}{p}_x\log (p_x\mathrm{/}{q}_x)$, with $+\mathrm{\infty }$ when $p_x>0=q_x$.

Finite product KL is the sum of coordinate KLs: the finite KL of product probability vectors is the sum of their coordinate KLs, including the singular-support case.

Posterior of a finite masking chain: conditional on the clean string, the true reverse posterior is a product of copy rows and two-point mask/unmask rows.

Finite-grid masked reverse sampler: the learned reverse rows copy unmasked tokens and use the same mask/unmask probability with clean-token vector $h_{i,j}$.

Supervised prediction problems: a supervised problem specifies measurable feature, label, action spaces, a joint law, a predictor class and a nonnegative measurable loss.

Population and empirical risk: population risk is the extended nonnegative integral of the loss under the data law.

Population and empirical risks are measurable: the finite-state loss and conditional risk are measurable and defined.

For the finite masking state space $S=(V\cup \{m\}{)}^L$ and reverse kernels $P_j$ of the finite-grid sampler, put a point-mass prior at $Z_N=m^L$. Define a decoder $D$ from $S$ to $V^L$ by $D(x\mid z_0)={\mathbf{1}}_{\{x=z_0\}}$ when $z_0\in V^L$, and $D(x\mid z_0)=\mathrm{\mid }V{\mathrm{\mid }}^{-L}$ for every $x\in V^L$ when $z_0\mathrm{\notin }{V}^L$. The finite joint mass of a clean output $x\in V^L$ and latent path $(z_0,\dots ,z_N)\in S^{N+1}$ is $p(x,z_0,\dots ,z_N)={\mathbf{1}}_{\{z_N=m^L\}}({\prod }_{j=1}^N{P}_j(z_j,z_{j-1}))D(x\mid z_0)\mathrm{.}$ Thus reverse step $j$ maps the state at level $j$ to the state at level $j-1$; the decoder is deterministic on clean endpoint states and is explicitly normalized on all other states.

Finite masking diffusion chain: fixes the finite vocabulary, mask, sequence length, grid, survival schedule, and sequence state space.

Finite-grid masked reverse sampler: supplies each reverse kernel $P_j$ and starts at $m^L$.

Measure kernel and probability kernel: a probability kernel supplies normalized transition rows.

Let $V$ be finite and nonempty, $m\mathrm{\notin }V$, $L,N\ge 1$, and $S=(V\cup \{m\}{)}^L$. Let $P_j$ be probability kernels on $S$, put prior mass one at $z_N=m^L$, and let $D(x\mid z_0)={\mathbf{1}}_{\{x=z_0\}}$ when $z_0\in V^L$ and $D(x\mid z_0)=\mathrm{\mid }V{\mathrm{\mid }}^{-L}$ otherwise. Then the finite joint mass $p(x,z_0,\dots ,z_N)={\mathbf{1}}_{\{z_N=m^L\}}({\prod }_{j=1}^N{P}_j(z_j,z_{j-1}))D(x\mid z_0)$ is nonnegative and sums to one over $x\in V^L$ and $(z_0,\dots ,z_N)\in S^{N+1}$.