Masked-diffusion KL reduces to weighted cross-entropy

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.