Rethinking Soft Tokens for Parallel Decoding in Diffusion Language Models
Organizations: National University of Singapore
Abstract
Diffusion language models (DLMs) enable parallel generation by predicting and committing multiple tokens at each denoising step, yet they can generate individually plausible but mutually inconsistent tokens. Recent work shows that soft tokens can mitigate this issue by representing uncertain positions with continuous embeddings built from the model's predictive distribution at the previous decoding step. However, although soft tokens are commonly understood as preserving predictive uncertainty, how soft-token feedback improves parallel decoding has not been systematically examined. In this paper, we investigate this question in frozen pretrained DLMs to examine soft-token feedback without the effects of additional training. To construct soft-token inputs in a training-free setting, we identify a geometric mismatch between conventional soft-token construction and the pretrained embedding space. Based on this observation, we propose a training-free, geometry-aware construction of soft tokens. Our analysis of soft-token feedback suggests that uncertainty preservation alone does not fully explain how it reshapes subsequent predictions. To better explain how soft-token feedback improves parallel decoding, we provide empirical evidence that it favors coherent token sequences. Across four pretrained DLMs and four math and code benchmarks, our method outperforms standard parallel decoding and a training-free Euclidean soft-token baseline. Code: https://github.com/kodaikawamura/rethinking-soft-tokens
Figures & tables
| Model | Additive | Multiplicative |
| LLaDA | 0.2099 | 0.0410 |
| LLaDA-1.5 | 0.0496 | 0.0019 |
| LLaDA-2.0 mini | 0.2198 | 0.0835 |
| Dream | 0.0566 | 0.0119 |
| Model | Method | GSM8K | MATH500 | HumanEval | MBPP | ||||||||
| Tokens/step | 2 | 4 | 8 | 2 | 4 | 8 | 2 | 4 | 8 | 2 | 4 | 8 | |
| LLaDA-2.0 mini | Vanilla | 89.61 | 83.09 | 49.81 | 40.60 | 37.40 | 16.20 | 70.73 | 46.95 | 24.39 | 55.80 | 42.80 | 23.80 |
| Euclidean | 89.76 | 85.82 | 51.78 | 43.60 | 37.20 | 18.40 | 69.51 | 50.00 | 22.56 | 57.60 | 43.80 | 23.00 | |
| Ours | 90.60 | 87.40 | 54.20 | 44.60 | 39.00 | 18.80 | 71.95 | 53.05 | 23.17 | 57.80 | 44.20 | 23.40 | |
| LLaDA-1.5 | Vanilla | 72.25 | 65.13 | 43.82 | 20.80 | 17.60 | 15.00 | 38.41 | 24.39 | 12.80 | 34.80 | 26.20 | 17.20 |
| Euclidean | 71.40 | 66.79 | 47.54 | 21.00 | 18.80 | 14.80 | 39.02 | 29.88 | 15.24 | 35.00 | 27.40 | 18.00 | |
| Model | Dataset | Candidate–candidate | Aggregate– [MASK] |
| LLaDA | GSM8K | 0.496 | 0.015 |
| HumanEval | 0.447 | 0.011 | |
| LLaDA-1.5 | GSM8K | 0.543 | 0.014 |
| HumanEval | 0.497 | 0.011 | |
| LLaDA-2.0 mini | GSM8K | 0.724 | 0.043 |
| HumanEval | 0.629 | 0.074 |
| Model | Vanilla | Euclidean | Ours |
| LLaDA | 2.1989 | 0.3415 | 0.2458 |
| LLaDA-1.5 | 2.0327 | 1.9273 | 0.5624 |
| LLaDA-2.0 mini | 2.7338 | 1.5610 | 1.1739 |
| Dream | 4.2733 | 4.1139 | 1.7553 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| SLERP | Linear | |||
| Model | Benchmark | [MASK] norm | Candidate norm | [MASK] norm |
| Dream | MATH-500 | 22.40 | 21.20 | 16.40 |
| HumanEval | 50.61 | 50.00 | 50.00 | |
| LLaDA | MATH-500 | 23.60 | 22.60 | 21.80 |
| HumanEval | 45.12 | 41.46 | 42.07 | |
| LLaDA 1.5 | MATH-500 | 24.00 | 21.80 | 21.20 |
| LLaDA | Dream | |||
| Sampler | GSM8K | HumanEval | GSM8K | HumanEval |
| Confidence thresholding | 29.80 | 45.12 | 79.80 | 55.49 |
| + Ours | 30.20 | 48.78 | 81.60 | 58.54 |
| EB-Sampler | 24.00 | 46.34 | 82.60 | 56.71 |
| + Ours | 26.00 | 48.17 | 83.60 | 58.54 |
| GSM8K | MATH500 | HumanEval | MBPP | ||||||||||
| Model | Method | 2 | 4 | 8 | 2 | 4 | 8 | 2 | 4 | 8 | 2 | 4 | 8 |
| LLaDA-2.0 mini | Vanilla | 95.50 | 88.33 | 53.67 | 88.83 | 82.04 | 37.38 | 83.33 | 56.82 | 28.79 | 83.70 | 64.89 | 34.99 |
| Euclidean | 95.41 | 91.92 | 55.24 | 92.72 | 81.07 | 41.26 | 81.06 | 62.12 | 28.03 | 85.89 | 66.46 | 34.80 | |
| Ours | 96.94 | 93.23 | 58.73 | 94.17 | 85.92 | 42.23 | 85.61 | 64.39 | 28.03 | 85.58 | 66.46 | 35.42 | |
| LLaDA-1.5 | Vanilla | 91.25 | 83.35 | 56.59 | 71.31 | 59.84 | 40.00 | 70.89 | 46.84 | 22.78 | 78.50 | 61.00 | 38.50 |
| Euclidean | 89.99 | 84.62 | 61.01 | 68.85 | 59.02 | 46.72 | 69.62 | 54.43 | 26.58 | 78.50 | 59.50 | 40.50 | |
| Model | Dataset | SLERP | Linear | Discrete MTP |
| Dream | HumanEval | 0.000704 | 0.084962 | 0.001848 |
| MATH-500 | 0.000554 | 0.003955 | 0.001372 | |
| LLaDA-1.5 | HumanEval | 0.035973 | 0.037907 | 0.038491 |
| MATH-500 | 0.015964 | 0.016934 | 0.017006 | |
| LLaDA-8B | HumanEval | 0.000665 | 0.000918 | 0.000977 |
| MATH-500 | 0.053294 | 0.053845 | 0.065190 |
| Lead-in | Alternative A | Alternative B |
| The location is | New York | San Diego |
| The movie series is | Star Wars | Harry Potter |
| The framework is | React Native | Swift UI |
| The dessert is | ice cream | apple pie |
| The pet is | black cat | white dog |
| The meal is | fried rice | tomato soup |
| Alternative A | Alternative B |
| red apple pie | green bean soup |
| silver sports car | black pickup truck |
| summer music festival | winter sports tournament |
| mountain bike trail | coastal hiking path |
| modern art museum | ancient history archive |
| solar energy research | marine biology laboratory |
| Alternative A | Alternative B |
| New York subway station | San Diego beach hotel |
| black leather office chair | white wooden dining table |
| summer music festival ticket | winter sports tournament pass |
| mountain bike repair shop | coastal hiking guide office |
| modern science fiction magazine | classic detective mystery novel |
| international airport security checkpoint | local railway ticket office |
Language Model Decoding
Language model decoding studies how models choose and generate outputs, balancing quality, reliability, diversity, and inference cost.
See all 161 papers