The
Diffusion
Revolt
Runs
Locally

Every language model in production today generates text the same way: one token at a time, strictly left to right, never reconsidering what came before. The field treated this as a law of nature. LLaDA proves it was always an assumption. And not a necessary one.

This post covers the full picture: the theoretical challenge LLaDA mounts against autoregressive orthodoxy, the exact architecture and sampling algorithm, how MLX's unified-memory design maps to LLaDA's inference structure, and where the implementation does and does not perform. The full code lives in the GitHub repository, linked throughout.

Autoregressive models (ARMs) factorize the probability of a sentence as a product of conditionals: p(x₁) × p(x₂|x₁) × ... × p(xₙ|x₁...xₙ₋₁). The chain rule of probability makes this mathematically exact. But it bakes in a directional constraint: token i can only see tokens 1 through i−1. The attention mask enforces this structurally. It is not a soft bias, it is a hard architectural ceiling.

LLaDA's central argument is precise: the intelligence in LLMs comes not from autoregression itself, but from the deeper goal both ARMs and diffusion models pursue — approximating the true language distribution via maximum likelihood. The chain-rule factorization is one valid road to that destination. It is not the only road. Masked diffusion models reach the same destination through a different mechanism, with a different set of structural properties.

The paper was accepted as an oral presentation, a distinction awarded to fewer than 1% of submissions. The field apparently considered this a genuinely landmark result, not an incremental variation. Whether that consensus holds as follow-up work accumulates is the more interesting question.

LLaDA is a Masked Diffusion Language Model. It has two processes: a forward process that corrupts text by masking tokens, and a reverse process that reconstructs text by iteratively unmasking them. Both are trained jointly; only the reverse process runs at inference.

Forward process. During pretraining, a sentence of L tokens is taken, a masking ratio t ~ Uniform(0,1) is sampled, and each token is independently replaced by a [MASK] token with probability t. A bidirectional Transformer then predicts all masked tokens simultaneously, conditioned on the unmasked ones. The loss is cross-entropy over masked positions, forming a variational lower bound (ELBO) on the model's log-likelihood.

Reverse process. At inference, the response field is initialized as all masks. The model runs N denoising steps from t=1 down to t=0. Each step:

The most important implementation detail the paper buries: LLaDA's Transformer is structurally near-identical to LLaMA 3 8B. Same hidden dimension (4096), same 32 layers, same RMSNorm, same SwiGLU activations, same RoPE positional embeddings. There is exactly one architectural change from a standard decoder:

In mlx-d, the entire patch is a single line:

# The single architectural delta from LLaMA → LLaDA:
# Remove the causal attention mask for bidirectional attention.
import mlx_lm.models.llama as llama_mod

llama_mod.create_attention_mask = lambda *a, **kw: None

Traditional ML frameworks treat CPU and GPU as separate devices with separate physical memory. On Apple Silicon, this distinction does not exist. The CPU, GPU, and Neural Engine share one physical memory pool. MLX is designed from the ground up around this reality.

Zero-copy operations. An MLX array lives once in memory. Any compute unit operates on it in place with no transfers, no synchronization overhead. For LLaDA's repeated full forward passes over the full sequence, this eliminates the dominant bottleneck that exists on discrete-GPU systems.

Lazy computation and graph compilation. MLX evaluates lazily, so arrays materialize only when needed. Applying mx.compile() to the forward function yields consistent throughput gains on subsequent calls. Always run one warmup pass before benchmarking: first-pass shader compilation is a one-time cost.

The zero KV cache line deserves attention. AR models must cache all previous key-value pairs to avoid recomputing attention from scratch at each generation step; the cache grows linearly with sequence length. LLaDA doesn't generate left-to-right; no prefix accumulates. Memory usage is fixed and predictable regardless of response length.

The full implementation is in the repository. What follows explains the non-obvious decisions that determine whether a working implementation produces coherent text or noise.

The sampling loop. The standard mlx_lm.generate function implements autoregressive sampling and is structurally incorrect for LLaDA. A custom loop is required:

# Confidence via logsumexp trick: avoids materializing the
# full (1, seq_len, vocab_size) softmax tensor (~130MB/step).
x0_logits = mx.take_along_axis(
    logits, mx.expand_dims(x0, axis=-1), axis=-1,
).squeeze(-1)
lse = mx.logsumexp(logits, axis=-1)
confidence = mx.exp(x0_logits - lse)

# Commit top-k via argpartition (O(n) vs O(n log n) sort)
partitioned = mx.argpartition(conf_flat, kth=conf_flat.shape[0] - k)
top_indices = partitioned[-k:]

Remasking strategy. Two options: random remasking and low-confidence remasking. For the Instruct model, low-confidence remasking is measurably better. Random remasking occasionally re-masks tokens the model was confident about, wasting steps. Low-confidence remasking only re-masks positions the model is genuinely uncertain about.

Semi-autoregressive block mode. The recommended configuration splits the response into blocks (default: 32 tokens) processed left-to-right, with tokens within each block unmasked in parallel. 64 passes over 128 tokens becomes 4 block passes of 32 tokens each.

# Install
pip install mlx-d

# Generate text
mlx-d generate "What is masked diffusion?" --steps 64 --block-length 32

# Run benchmarks
mlx-d benchmark

# Speculative diffusion decoding
mlx-d spec-decode "Explain LLaDA." --draft-steps 8

The benchmark picture is honest. LLaDA-8B-Instruct is competitive with LLaMA 3 8B on general knowledge, definitively better on bidirectional constraint tasks, and currently weaker on multi-step mathematical reasoning.

The math weakness has a precise diagnosis. The literature now calls it the Flexibility Trap: confidence-based token selection during denoising systematically produces generation orders that approximate left-to-right sequential behavior.

The d1 framework (RL for diffusion LMs) reached approximately 84.5% on GSM8K with LLaDA by replacing confidence-based selection with policy-gradient optimization. The model weights are not the bottleneck. The sampling algorithm is.

The highest near-term value from LLaDA is not as a drop-in replacement for production AR systems. It is as a targeted component inside existing pipelines, deployed where its specific structural properties address bottlenecks that AR models handle poorly.

mlx-d includes a complete speculative diffusion decoding implementation: drafter-verifier framework, acceptance criterion, and throughput benchmarks. The voice pipeline that motivated this project is available as an optional install (pip install mlx-d[voice]).

The Bigger Picture

LLaDA establishes something the field will be unwinding for years: the generation algorithm and the model architecture are separable concerns. A Transformer without a causal mask, trained with the right objective, achieves competitive language understanding at 8B scale. The autoregressive mechanism is not intelligence. It is one implementation of the goal.

MLX establishes something equally important on the hardware side: unified memory architectures remove the dominant inference bottleneck for non-standard models. Running frontier research locally is no longer purely a function of hardware budget.

The boundaries between AR and diffusion, between cloud and local, between research prototype and production pipeline are all softer than they appeared twelve months ago. The implementation is in the repository.