From LLaDA's own probabilities: sampling 2 masked tokens in 1 step gives "San Angeles" ~47% of the time; one extra step drops it to ~0.2%
I've been measuring LLaDA-8B-Instruct on an A100 (bf16) and wanted to share a small, reproducible result about the step count. It's the factorization issue the paper discusses, but with concrete numbers.
Setup. Chat-templated prompt: Reply with ONLY a two-word major California city (either Los Angeles or San Francisco), nothing else:, followed by two [MASK] slots. One forward pass, softmax at each slot, then renormalized over the two candidate words.
Step 1 (both slots at once).
- Slot 1: "Los" 70.69%, "San" 29.31%
- Slot 2: "Angeles" 56.16%, "Francisco" 43.84%
Each slot is predicted without seeing the other. If both are sampled in the same step, the chance of a city that doesn't exist ("Los Francisco" or "San Angeles") is 0.7069 × 0.4384 + 0.2931 × 0.5616 = 47.45%. With greedy argmax, this prompt gives Los + Angeles, so the problem shows up when sampling.
Step 2 (lock slot 1, re-run). With slot 1 fixed to "Los", P("Angeles" | "Los") = 99.95%. With "San", P("Francisco" | "San") = 99.37%. The expected hybrid rate falls to 0.22%.
Two more prompts showed the same pattern:
| Pair | Hybrid rate, 1 step | Hybrid rate, 2 steps |
|---|---|---|
| New York / Los Angeles | 46.87% | 0.07% |
| Open AI / Deep Mind | 47.23% | 2.34% |
64-token answer, greedy, sweeping steps S. Adjacent-token repetition was:
| S | 1 | 2 | 4 | 8 | 16 | 32 | 64 |
|---|---|---|---|---|---|---|---|
| Repetition | 54.1% | 60.7% | 63.9% | 4.9% | 0.0% | 0.0% | 0.0% |
For this prompt, the quality cliff sat between 4 and 8 tokens per step.
Minimal code for the two-slot check:
# pip install "transformers<5" accelerate
import torch, torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
name = "GSAI-ML/LLaDA-8B-Instruct"
tok = AutoTokenizer.from_pretrained(name, trust_remote_code=True)
model = AutoModel.from_pretrained(name, trust_remote_code=True, torch_dtype=torch.bfloat16).cuda().eval()
MASK = 126336
prompt = "Reply with ONLY a two-word major California city (either Los Angeles or San Francisco), nothing else:"
ids = tok(tok.apply_chat_template([{"role": "user", "content": prompt}], add_generation_prompt=True, tokenize=False),
return_tensors="pt").input_ids.cuda()
x = torch.cat([ids, torch.full((1, 2), MASK, device="cuda")], dim=1)
def p(probs, word): # sum over with/without leading space
return sum(probs[tok.encode(w, add_special_tokens=False)[0]].item() for w in {word, " " + word})
with torch.no_grad():
pr = F.softmax(model(x).logits[0, -2:].float(), -1)
p1 = torch.tensor([p(pr[0], "Los"), p(pr[0], "San")]); p1 /= p1.sum()
p2 = torch.tensor([p(pr[1], "Angeles"), p(pr[1], "Francisco")]); p2 /= p2.sum()
print("1-step hybrid rate:", (p1[0] * p2[1] + p1[1] * p2[0]).item())
Has anyone found a good rule for picking the step count (or a confidence threshold) per task, beyond trial and error?
I covered this, plus the KV-cache side, in a video: https://youtu.be/WNfrA_zpDgc (disclosure: it's my channel). Thanks to the LLaDA team for releasing the weights.