Masked Diffusion Language Modeling on Infinite Jest

Basic formulation

Let x0 be a clean token sequence from the corpus. I split the sequence into a visible prefix and a continuation region. The prefix is left unchanged. The continuation region is corrupted according to a sampled noise rate t.

model input:      visible prefix + corrupted continuation
model target:     clean continuation tokens at corrupted positions
attention mask:   bidirectional
training loss:    denoising cross entropy, not causal LM loss

The denoiser is trained to estimate p(x0 | xt, t) for the corrupted book tokens. During generation, I reverse the process approximately: start from a noisy continuation and repeatedly apply the denoiser until the canvas stabilizes.

This is the part I care about most. A causal LM represents text as product_i p(x_i | x_<i). The diffusion sampler instead keeps a full continuation canvas in memory and revises positions in parallel.

Corruption process

I used mixed discrete corruption. Selected continuation positions are corrupted in one of two ways:

selected token -> [MASK]
selected token -> random vocabulary token
unselected token -> original token

The random-token branch matters because the sampler does not begin from pure [MASK] tokens. It begins from uniform vocabulary noise. A model trained only on mask replacement can learn a fill-in-the-blank task that is easier than the sampling problem. Mixed corruption moves training closer to the inference distribution I actually use.

The mdlm-subs setting removes special tokens and <mask> from the clean prediction space. The target is always a real Infinite Jest token. I reweight the loss by the sampled noise rate with a cap, which keeps high-noise and low-noise cases in the objective without letting very small corruption rates dominate the update.

Relation to D3PM

After building this run, I checked Austin et al.'s D3PM paper. The closest formal analogue is the absorbing-state text diffusion process: tokens either stay unchanged or move into an absorbing mask state, and the network predicts the clean sequence x0 from a corrupted sequence xt.

My current sampler is not an exact D3PM reverse chain. It uses mixed mask and random-token corruption during training, then entropy-based uniform re-noising at sampling time. I treat that sampler as a practical full-canvas refinement heuristic. The clean next baseline is a pure absorbing [MASK] D3PM trained on the same corpus, with loss-weighting and step-count ablations against this checkpoint.

Sampling procedure

The sampler is full-canvas refinement. It initializes the continuation region with uniform token noise. Each reverse step predicts every generated position. Positions with lower uncertainty are retained. Positions with higher uncertainty are replaced with fresh noise and predicted again.

I measure uncertainty from the predicted distribution at each generated position. Low-entropy positions are treated as more stable. High-entropy positions are treated as unresolved and are re-noised.

Self-conditioning adds state between reverse steps. The previous step's predicted token distribution is projected through the embedding table to form an expected token embedding. That embedding is blended into the next denoising pass at generated positions. The model still predicts all positions in parallel. Self-conditioning does not introduce a left-to-right factorization.

Denoising traces

The first trace uses a one-sentence prompt from the corpus. I keep the prompt fixed and sample the continuation by iterative denoising.

I am seated in an office, surrounded by heads and bodies. I have been there for a few short hours. My desk's right next to the door of the room. The door's white, and slightly deformed, and the back of a chair. It's clean, bright, and tidy.

This is conditional generation in the book's learned local style. It is not an attempt to recover the original paragraph.

The second trace uses a different sentence from the corpus. I included it because it shows the model honestly. It can stay near the topic and local style for a short span, then it degrades into brittle phrase structure and acronym-like fragments.

Sampler ablation

I evaluated three samplers on the same checkpoint. The sweep used two prompts, three seeds, and 48 denoising steps.

The quality score is a local diagnostic. It rewards generated length and lexical diversity and penalizes repetition and punctuation collapse. It is not a general language-model benchmark. In this setting I treat it as useful because the ranking matches sample inspection.

Uniform-state refinement outperforms mask-only refinement. This is consistent with the training setup. The model was exposed to random-token corruption, and the sampler uses random-token initialization and re-noising. Self-conditioning gives a smaller gain on this sweep.

Interpretation

I think the experiment supports a narrow claim: a masked language model can be posttrained into a small text diffusion generator by training it to denoise variable-rate mask and uniform-token corruption, then sampling with full-canvas refinement.

The strongest implementation details in this run were mixed corruption, high-noise continuation training, uniform re-noising, entropy-based token retention, and self-conditioning.

The result is not broad fluency. The model can produce short locally coherent continuations that have some of the Infinite Jest distribution in them, but it still repeats, loses syntax, and collapses into malformed fragments. The useful property is structural: generation is represented as refinement of a noisy canvas rather than as next-token prediction.