ReLaDrift: Recurrent Language Generation via Drifting for Test-Time Scaling
Organizations: Institute of Science Tokyo · AIST · NII LLMC
Abstract
Drifting Models generate a sample with a single function evaluation, making them attractive for fast parallel language generation, but their one-step formulation does not support iterative refinement. We propose Recurrent Language Generation via Drifting (ReLaDrift), which uses a shared Drifting Model to generate an initial draft from noise and iteratively refine it in a continuous latent space. We train this recurrence by sampling the number of rounds and applying the drifting loss only at the final round, stopping gradients through preceding rounds so that only one round requires backpropagation. Unlike the one-step Drifting Model, whose accuracy does not improve when iterated, ReLaDrift enables a trade-off between test-time computation and generation quality, as diffusion and flow language models do. On Sudoku Hard and GSM8K, it achieves competitive or superior performance compared with these iterative-generation approaches across a range of inference budgets. Our code is available at https://github.com/Yuki-Ots/ReLaDrift
Figures & tables
| Method | NFE=1 | 4 | 5 | 8 | 16 | 32 | 64 | 128 | 180 |
|---|---|---|---|---|---|---|---|---|---|
| AR (sequential) | Greedy: ; sampling: | ||||||||
| MDLM | |||||||||
| Duo | |||||||||
| S-FLM | |||||||||
| FMLM+ | |||||||||
| FRM (Euler) | |||||||||
| Setting | NFE=1 | 4 | 5 | 8 | 16 | 32 | 64 | 128 | 180 |
|---|---|---|---|---|---|---|---|---|---|
| Default | |||||||||
| Input formulation (Default: Source) | |||||||||
| - Additive | |||||||||
| - Conditional | |||||||||
| Training depth (Default: ) | |||||||||
| - | |||||||||
| Method | NFE | ||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| 1 | 4 | 5 | 8 | 16 | 32 | 64 | 128 | 256 | 512 | 1024 | |
| AR (sequential) | 72.3 | ||||||||||
| MDLM | 0.3 | 1.3 | 1.9 | 5.2 | 12.1 | 18.8 | 27.2 | 28.7 | 29.7 | 32.0 | 32.1 |
| Duo | 0.5 | 5.2 | 6.8 | 12.5 | 23.1 | 30.6 | 34.7 | 34.9 | 36.3 | 37.1 | 35.9 |
| S-FLM | 0.5 | 10.5 | 13.3 | 16.2 | 17.7 | 19.9 | 20.2 | 20.5 | 20.3 | 20.4 | 20.8 |
| FMLM+ | 0.0 | 2.6 | 4.9 | 10.1 | 20.5 | 24.9 | 26.0 | 25.6 | 25.5 | 25.5 | 25.5 |
| Setting | NFE | |||
|---|---|---|---|---|
| 1 | 4 | 8 | 16 | |
| Default | 13.2 | 27.3 | 29.0 | 29.4 |
| Number of given positives (default 12) | ||||
| - | 10.4 | 22.7 | 24.4 | 24.9 |
| - | 1.0 | 8.7 | 10.7 | 12.2 |
| Training depth (default 5) | ||||
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Setting | NFE=1 | 4 | 5 | 8 | 16 | 32 | 64 | 128 | 180 |
|---|---|---|---|---|---|---|---|---|---|
| , (default) | 7.3 1.0 | 78.9 0.6 | 82.5 0.8 | 88.4 0.4 | 92.1 0.4 | 93.1 0.6 | 93.5 0.7 | 93.9 0.7 | 94.0 0.7 |
| , | 7.6 2.2 | 78.1 1.4 | 81.7 1.3 | 87.4 1.5 | 91.1 1.2 | 92.1 1.1 | 92.6 1.2 | 92.7 1.2 | 92.9 1.4 |
| , | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 |
| , | 6.5 1.7 | 76.0 2.6 | 80.7 2.2 | 88.0 1.2 | 93.5 0.6 | 95.9 0.4 | 97.1 0.3 | 98.0 0.2 | 98.3 0.2 |
| Item | ReLaDrift (ours) | S-FLM | FMLM+ | FRM (our implementation) | DFM (our implementation) |
|---|---|---|---|---|---|
| Sequence, vocabulary | 180 tokens, vocabulary 12, fixed prefix 91 | same | same | same | same |
| State | continuous vector (12-dim one-hot code) | learned embedding on the sphere | one-hot + Gaussian noise | one-hot + carry | one-hot + Gaussian noise (linear interpolant ) |
| Encoder / decoder | one-hot code / inner product with the codes + argmax | learned embedding on the sphere / of the output layer | one-hot / argmax | one-hot / argmax | one-hot / argmax |
| Transformer | 9 layers, width 512, 8 heads, FFN 2,048, RoPE, dropout 0.1, round embedding added to the input | DiT on the sphere, 8 layers, width 512, 8 heads | DiT, 8 layers, width 512, 8 heads, two-time adaLN | DiT, 8 layers, width 512, 8 heads, adaLN | DiT of FMLM+, 8 layers, width 512, 8 heads, two-time adaLN ( , ) |
| Parameters | 28.39M | 28.56M | 28.61M | 28.57M | 28.61M |
| Optimizer | AdamW 3e-4 (warmup 2,500, then constant), , weight decay 0, clip 1 | same (official) | same (official) | same; optimizer and EMA restarted at Stage B | same |
| Item | ReLaDrift (ours) | AR | S-FLM | FMLM+ |
|---|---|---|---|---|
| Positives | first of the 12 solutions (fixed) + up to 4 verified samples (positive queue) | 12 (problem, solution) pairs per problem, shuffled | same | same |
| Tokenizer, length | CodeT5 (vocabulary 32,100), 576 | SmolLM-135M (vocabulary 49,153), 512 | same | same |
| Latent | CodeT5 contextualized embedding, 512 dim (frozen) | discrete tokens | learned embedding on the sphere | one-hot + Gaussian noise |
| Transformer | DiT blocks of S-FLM (bidirectional, no adaLN), 12 layers, width 768, 12 heads, FFN 3,072, RoPE, round embedding | DiT, 12 layers, width 768, 12 heads (causal) | S-arch, 12 layers, width 768, 12 heads | DiT, 12 layers, width 768, 12 heads, two-time adaLN |
| Parameters | core 87.4M (+ frozen encoder 35.3M, decoder 42.2M = 164.9M) | 160.5M | 163.0M | 168.0M |
| Optimizer | AdamW , weight decay 0, clip 1 (same as the compared methods) | AdamW , weight decay 0 (= Adam), clip 1 (official) | same | same |
| Condition | NFE=1 | 2 | 4 | 5 | 8 | 16 | 32 | 64 | 128 | 180 |
|---|---|---|---|---|---|---|---|---|---|---|
| Source (main) | 7.3 1.0 | 54.7 0.3 | 78.9 0.6 | 82.5 0.8 | 88.4 0.4 | 92.1 0.4 | 93.1 0.6 | 93.5 0.7 | 93.9 0.7 | 94.0 0.7 |
| Additive | 7.1 2.2 | 54.8 3.9 | 77.6 2.6 | 81.5 3.1 | 86.2 3.6 | 89.5 3.1 | 90.5 3.0 | 91.0 2.8 | 91.6 2.5 | 91.6 2.8 |
| Conditional | 7.4 0.9 | 54.3 3.3 | 77.6 1.7 | 81.1 1.1 | 86.1 1.2 | 89.6 1.8 | 90.7 1.7 | 91.2 2.0 | 91.6 2.3 | 91.7 2.2 |
| Single-round Drifting Model | 5.1 0.8 | 4.9 0.8 | 4.9 0.8 | 5.0 0.8 | 4.9 0.8 | 4.9 0.8 | 4.9 0.8 | 4.9 0.8 | 4.9 0.8 | 4.9 0.8 |
| Source, w/ verifier | 7.0 1.0 | 56.4 0.6 | 80.0 1.0 | 83.7 0.7 | 88.4 0.7 | 92.0 0.6 | 93.0 0.6 | 93.4 0.8 | 93.6 0.8 | 93.7 0.8 |
| Source, | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 | 0.0 0.0 |
| Condition | Round embedding | NFE=8 | 16 | 32 | 64 | 128 |
|---|---|---|---|---|---|---|
| Source | extrapolated | 39.4 40.2 | 46.4 41.4 | 54.4 47.2 | 57.1 49.5 | 58.0 50.2 |
| clamped | 88.4 0.4 | 92.1 0.4 | 93.1 0.6 | 93.5 0.7 | 93.9 0.7 | |
| Additive | extrapolated | 38.6 40.3 | 33.5 19.0 | 49.0 26.8 | 55.6 29.6 | 58.4 32.3 |
| clamped | 86.2 3.6 | 89.5 3.1 | 90.5 3.0 | 91.0 2.8 | 91.6 2.5 | |
| Conditional | extrapolated | 74.7 18.5 | 87.6 1.5 | 89.5 2.5 | 90.1 2.7 | 90.9 3.4 |
| clamped | 86.1 1.2 | 89.6 1.8 | 90.7 1.7 | 91.2 2.0 | 91.6 2.3 |
| Round embedding | NFE=8 | 16 | 32 | |
|---|---|---|---|---|
| extrapolated | 11.1 | 13.3 | 14.7 | |
| clamped | 10.7 | 12.2 | 12.9 | |
| extrapolated | 24.0 | 23.4 | 24.9 | |
| clamped | 24.4 | 24.9 | 25.2 | |
| extrapolated | 29.0 | 29.4 | 30.1 | |
| clamped | 29.0 | 29.4 | 29.4 |
| Setting | NFE=1 | 4 | 5 | 8 | 16 | 32 | 64 | 128 | 256 | 512 | 1024 |
|---|---|---|---|---|---|---|---|---|---|---|---|
| , verifier, (default) | 13.2 | 27.3 | 27.4 | 29.0 | 29.4 | 29.4 | 29.4 | 29.5 | 29.5 | 29.4 | 29.4 |
| without verifier and queue | 4.5 | 17.0 | 17.3 | 18.5 | 20.0 | 20.3 | 20.8 | 20.7 | 20.7 | 20.6 | 20.7 |
| , verifier (stopped at 40,000 updates) | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | – | – | – | – |
| median | 10.3 | 24.0 | 24.5 | 24.6 | 26.3 | 26.8 | 26.8 | 27.0 | 27.0 | 27.0 | 27.0 |
| median (40,000 updates) | 0.0 | 0.2 | 0.2 | 0.2 | 0.4 | 0.4 | 0.4 | – | – | – | – |
| CodeT5 embedding table | 6.7 | 15.1 | 16.0 | 16.7 | 17.7 | 18.0 | 18.2 | 18.1 | 18.1 | 18.1 | 18.2 |
| ReLaDrift ( ), NFE=16 | |||||||
|---|---|---|---|---|---|---|---|
| Updates | AR | S-FLM | FMLM+ | default | w/o ver. | ||
| 20k | 48.1 | 0.7 | 3.1 | 9.1 | 6.7 | 0.0 | 5.5 |
| 40k | 56.7 | 2.4 | 6.7 | 16.1 | 11.9 | 0.0 | 12.7 |
| 60k | 59.8 | 4.7 | 9.0 | 19.6 | 13.9 | – | 17.2 |
| 80k | 62.5 | 5.9 | 11.2 | 22.1 | 16.6 | – | 20.4 |
| 100k | 62.5 | 8.0 | 16.9 | 22.8 | 17.6 | – | 20.6 |
| Updates | NFE=1 | 2 | 4 | 5 | 8 | 16 | 32 | 64 | |
|---|---|---|---|---|---|---|---|---|---|
| 1 | 20k | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
| 40k | 0.0 | 0.3 | 0.2 | 0.3 | 0.4 | 0.4 | 0.5 | 0.5 | |
| 60k | 0.2 | 0.3 | 0.5 | 0.6 | 0.8 | 1.1 | 1.1 | 1.1 | |
| 80k | 0.3 | 0.4 | 1.0 | 1.0 | 1.1 | 1.8 | 2.1 | 2.4 | |
| 100k | 0.3 | 0.5 | 1.3 | 1.4 | 2.4 | 3.5 | 4.5 | 5.2 | |
| 120k | 0.3 | 0.9 | 2.4 | 3.1 | 3.9 | 5.1 | 6.1 | 6.5 |
Test-Time Scaling
Test-time scaling studies how extra inference-time computation can improve model answers, decisions, and generated outputs.
See all 209 papers