Leave-One-Out Corrective Parallel Sampler

Integrate this at inference time in a masked or uniform discrete diffusion Transformer, replacing the current tau-leaping transition from state x at time t to t-Delta. The Transformer should output logits s_theta[i,b] for the clean-token posterior at every position. At each step: (1) construct the model input containing the current noisy sequence; (2) obtain q_i(b|x^{-i},t)=softmax_b(s_theta[i,b]) using a leave-one-out approximation, implemented either by replacing the token at position i with a learned MASK or REMASK embedding, or by training a dedicated leave-one-out prediction head; (3) compute the exact forward transition likelihood K_t(x_i|b) for every candidate vocabulary token; (4) multiply q_i by K_t and normalize over b according to the displayed formula; (5) sample all selected coordinates in parallel; and (6) continue from the new state. For efficiency, first implement the simultaneous-mask approximation, where every position's own current token is hidden from its prediction head, then compare it with an expensive exact leave-one-out baseline on short sequences. The paper supplies the posterior-times-transition correction; the approximation strategy and update subset are empirical choices. First test on a 12-layer Transformer trained on synthetic correlated binary strings and WikiText-103. Compare standard tau-leaping at equal model evaluations. Success means lower perplexity or total-variation error at the same number of evaluations, and especially correction of initially wrong parallel token updates rather than error accumulation at large Delta.