Source-linked AI summary

ReLaDrift: Recurrent Language Generation via Drifting for Test-Time Scaling

Yuki Otsuka, Daisuke Oba, Naoaki Okazaki

arXiv:2610.23225v1cs.CL

TL;DR

The one-step Drifting Model lacks a mechanism for improving outputs with additional test-time computation. ReLaDrift recurrently refines a draft in continuous latent space, training only at a sampled terminal round. On Sudoku Hard and GSM8K, it exhibits test-time scaling and competitive or superior performance to diffusion and flow language models across inference budgets.

  • Problem

    The one-step Drifting Model does not improve when iterated, leaving open how to train it to refine its own outputs and trade test-time computation for accuracy.

    Source: §1 · p. 1
  • Method

    ReLaDrift repeatedly applies a shared Drifting Model to refine a noise-generated draft in continuous latent space, training with loss only at a sampled terminal round.

    Source: §1 · p. 2
  • Results

    ReLaDrift exhibits test-time scaling on Sudoku Hard and GSM8K and outperforms diffusion and flow language models across NFE=4–180 on Sudoku Hard and up to NFE=16 on GSM8K.

    Source: §1 · p. 2, §6 · p. 8
  • Takeaways & Limitations

    ReLaDrift provides drifting-based generation with a test-time computation–quality trade-off on Sudoku Hard and GSM8K.

    Source: §1 · p. 2, §6 · p. 8

    On GSM8K, the advantage over parallel generators depends on a verifier and pretrained CodeT5 embeddings, and results use a single training seed.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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

1 Introduction

The one-step Drifting Model is fast but lacks a mechanism for improving outputs with more test-time computation, and simply iterating it does not improve accuracy. ReLaDrift instead refines drafts recurrently, achieving competitive or superior performance to diffusion and flow language models across inference budgets on Sudoku Hard and GSM8K.

1 Introduction

  • Simply iterating a one-step Drifting Model yields no accuracy gains on Sudoku Hard or GSM8K.The one-step formulation is trained to generate from noise rather than refine its own outputs.Source: §1 · p. 1
  • ReLaDrift generates an initial draft from noise, then iteratively refines it in a continuous latent space using a shared model.Training samples the number of rounds and applies the drifting loss only to the final output, stopping gradients through earlier rounds.Source: §1 · p. 2
  • ReLaDrift achieves competitive or superior performance to diffusion and flow language models across inference budgets on Sudoku Hard and GSM8K.The tasks have verifiable answers, enabling direct measurement of test-time scaling.Source: §1 · p. 2

2 Related Work

Related work contrasts iterative refinement in diffusion and flow models with the original Drifting Model’s one-step design. ReLaDrift combines a complete one-step draft with recurrent refinement, without a predefined noise-to-data trajectory.

Iterative generation vs. one-step drifting

  • The original Drifting Model is designed for one-step generation and has no built-in mechanism to benefit from additional inference steps.It is trained to generate from noise rather than refine its own outputs.Source: §2 · p. 2

Drifting-based methods

  • Drift Flow Matching enables test-time scaling through finer transport steps, while ReLaDrift first generates a complete draft and then refines it recurrently without a predefined noise-to-data trajectory.Drift Flow Matching is described as bridging drifting and flow matching (Ma et al., 2026).Source: §2 · p. 2

Recurrent refinement

  • Recurrent refinement has been studied in parallel language generation and diffusion, flow-based generation, and weight-tied recurrent models.The cited examples include Lee et al. (2018), Helbling et al. (2026), and Dehghani et al. (2019).Source: §2 · p. 2

3 Background: Drifting Model

The Drifting Model maps noise directly to generated samples and learns through a drift field that combines attraction toward real samples with repulsion from generated samples. Its field guides training but is not needed at inference.

3 Background: Drifting Model

  • The Drifting Model maps noise to generated samples in one step and learns to match the target data distribution without adversarial training.During training, a drift field computed from real and generated samples guides the generated distribution toward the target distribution.Source: §3 · p. 2
  • The drift field combines attraction toward real samples with repulsion from generated samples, and the training loss regresses outputs toward the stop-gradient field target.The field’s expectations are estimated from mini-batches of real positive and generated negative samples; it is not needed at inference.Source: §3 · p. 2
  • Figure 2 depicts ReLaDrift’s recurrent source formulation: noise creates the first draft, later rounds refine it, and the decoder is applied once to the final draft.The caption states that the number of test-time rounds may exceed the training depth and that training loss is applied at the terminal round only.Source: Figure 2 · p. 3

4 Methods

ReLaDrift operates on continuous representations of generated tokens and repeatedly applies a shared Drifting Model to refine drafts. Training samples a terminal round and stops gradients through earlier rounds; at inference, the number of rounds is freely chosen, including beyond training depth.

4.1 Continuous Representation of Token Sequences

  • A fixed encoder maps token sequences to continuous vectors, and a frozen decoder maps the model’s continuous output back to tokens in parallel.Only solution positions are generated, and the decoder adds one pass without reintroducing sequential generation.Source: §4.1 · p. 3

4.2 Recurrent Application of the Drifting Model

  • Each round uses the same network and passes drafts between rounds as continuous latents, without decoding and re-encoding them.The first round maps noise to a draft; later rounds take the previous draft, with shared refinement parameters and a round index.Source: §4.2 · p. 3
  • The source formulation uses noise only in the first round and then feeds each draft into the next; additive and conditional alternatives use fresh noise and require an initial draft.The main experiments use the source formulation; the alternatives differ in how noise and drafts enter the Transformer.Source: §4.2 · pp. 3–4

4.3 Training

  • Each update samples a round count and applies the drifting loss only at the terminal round, stopping gradients through preceding rounds rather than backpropagating through time.Sampling the terminal round trains the shared network at different depths; earlier rounds require forward passes, while only terminal-round activations are stored.Source: §4.3 · pp. 4–5

4.4 Generation and Test-Time Scaling

  • At inference, the number of rounds is freely chosen, and the decoder is applied once to the final draft.Rounds beyond training depth reuse the map trained at the last refinement round, with the round embedding clamped to the last trained round.Source: §4.4 · p. 5

5 Experiments

Across Sudoku Hard and GSM8K, ReLaDrift improves accuracy with additional inference rounds, outperforming compared diffusion and flow models on Sudoku and leading at matched NFE up to 16 on GSM8K. The experiments compare accuracy across inference budgets and, for GSM8K, also measure wall-clock generation time.

5 Experiments

  • Evaluation measures accuracy across inference budgets with fixed weights, counts each ReLaDrift round as one NFE, and uses one generation trajectory per problem without inference-time verification or candidate selection.Because NFE does not capture per-step cost, the study also compares GSM8K wall-clock generation time.Source: §5 · p. 5

5.1 Sudoku; 5.2 GSM8K

  • On Sudoku Hard, ReLaDrift improves monotonically across all three seeds, reaches 94.0% at NFE=180, and outperforms compared diffusion and flow models across NFEs; on GSM8K it rises from 13.2% at NFE=1 to 29.4% at NFE=16 and leads at matched NFE up to 16.The iterated one-step Drifting Model shows no accuracy improvement on either task.Source: §5.1 · p. 6, §5.2 · pp. 7–8

5.1 Sudoku

  • All three Sudoku input formulations scale with test-time computation, while Source has the highest accuracy and lowest seed variability at NFE ≥8; greater training depth improves multi-round accuracy even beyond Rmax = 5.Greater training depth costs additional forward passes, while backpropagation remains limited to the terminal round.Source: §5.1 · p. 6

5.2 GSM8K

  • For GSM8K’s multiple valid programs, training adds verified generated programs to per-problem positive queues and excludes them from repulsion; verification and queues are used only during training.Verification compares program outputs with a pseudo-reference formed by voting over the given programs’ answers.Source: §5.2 · p. 7
  • ReLaDrift takes 24 ms per round versus Duo’s 95 ms per step on one H200; it reaches 29.4% accuracy in 13 seconds, while Duo takes 63 seconds to reach comparable 30.6% accuracy.The wall-clock advantage is more pronounced than the matched-NFE comparison.Source: Figure 3 · p. 7, §5.2 · p. 8
  • On GSM8K, increasing the number of given positives substantially improves accuracy, and test-time scaling persists without the verifier and positive queue; Rmax = 5 performs best, though inference rounds beyond training depth still help.The authors suggest that longer refinement exposure and less frequent supervision at each depth may have task-dependent effects under a fixed training budget.Source: §5.2 · p. 8

6 Conclusion

ReLaDrift turns the one-step Drifting Model into a recurrent generator that refines drafts in continuous latent space. It scales with test-time computation on Sudoku Hard and GSM8K, while incurring higher per-update training cost.

  • ReLaDrift outperforms the compared models across NFE=4–180 on Sudoku Hard and up to NFE=16 on GSM8K.The conclusion positions drifting-based generation as a promising alternative for test-time-scalable parallel language generation.Source: §6 · p. 8
  • The model refines its own drafts in continuous latent space, applying the drifting loss at a sampled terminal round and avoiding backpropagation through time.Source: §6 · p. 8
  • ReLaDrift incurs approximately 3.3× the per-update training cost of a one-pass method.The cost reflects recurrent forward passes and drift-field estimation from multiple generated samples; the experiments limit this to n=2 samples per problem.Source: §6 · p. 8

Limitations

The GSM8K advantage over parallel generators depends on a verifier and pretrained CodeT5 embeddings, although test-time scaling persists without them.

  • The GSM8K advantage over parallel generators depends on a verifier and pretrained CodeT5 embeddings.Source: Limitations · p. 9
  • The verifier requires executable outputs and uses voted pseudo-references that may be incorrect.Source: Limitations · p. 9
  • GSM8K results use a single training seed because of limited training resources.The authors identify broader task and model-scale evaluation, verifier-free training with multiple valid targets, and narrowing the autoregressive-model gap as future work.Source: Limitations · p. 9

A.1 Encoding of the Prompt and the Target

The appendix encodes prompts and targets differently for contextual CodeT5 embeddings and context-free one-hot Sudoku codes. These choices determine which positions are conditioned on the prompt and generated.

  • For GSM8K, the target latent encodes the problem and solution together, while the prompt latent encodes only the problem and is placed at prompt positions.The contextual target representations at solution positions therefore incorporate the problem statement.Source: §A.1 · p. 11
  • GSM8K uses frozen CodeT5-small contextual embeddings of 512 dimensions and pads sequences to 576 tokens; generation and loss operate only after the prompt.Noise and drafts enter only at program, EOS, and padding positions, and the generated program is read through its first EOS.Source: §A.1 · p. 11
  • Sudoku uses one-hot codes for its 12 tokens and generates all 89 answer positions, including given digits and row separators.The puzzle is not part of the target because each token code is context-independent; evaluation uses exact match, and decoding is not restricted to digits.Source: §A.1 · pp. 11–12
  • The appendix also reports GSM8K experiments using context-free CodeT5 input embeddings and a random Gaussian embedding table.Source: §A.1 · p. 12

A.2 Model Architecture and Input Layout

The architecture uses round-specific inputs and outputs to generate from noise initially and refine a prior draft thereafter. It also describes additive and conditional input formulations, plus the shared Transformer backbone and round embedding.

  • The source formulation projects inputs into the model dimension and sums them, using separate noise and draft projections for initial generation and refinement rounds.The additive formulation instead adds projected draft to projected noise; its initial draft is zero on GSM8K and the blank-token code on Sudoku.Source: §A.2 · pp. 12–13
  • In GSM8K, prompt positions receive the prompt latent while later positions receive noise or draft; in Sudoku, noise and draft cover only the 89 answer positions.The source-formulation Sudoku input has 180 positions, with outputs taken from the final 89 positions.Source: §A.2 · p. 12
  • The round embedding ϕ(r) supplies the current round index, with r starting at 0; its constant coordinate keeps sine and cosine terms at the same four frequencies.Source: §A.2 · p. 12
  • The model backbone is bidirectional pre-normalization Transformer blocks with scale-only LayerNorm, GELU feed-forward networks, and RoPE on queries and keys.No attention mask is used, and padding positions participate in GSM8K attention.Source: §A.2 · p. 12
  • The output projection uses Winit at r=0 and Wrefine at r≥1; the round index selects the corresponding input and output projections.The round projection and Winit/Wrefine are initialized to zero, while Wout is normally initialized with zero bias.Source: §A.2 · p. 13
  • The conditional formulation concatenates the prior-draft and current-generation sequences before the model, and is reported only on Sudoku.It increases attention length to 269 positions, compared with 180 for the source and additive formulations.Source: §A.2 · p. 13

A.3 Drift Field and Loss

The drift field attracts generated samples toward positives and repels them from other generated samples, with scaled, multi-temperature weights. A stop-gradient target converts this field into a loss applied only at the terminal round.

Scaled coordinates

  • Distances and fields use scaled coordinates restricted to post-prompt positions, covering all 89 Sudoku answer positions.The scale is computed from mini-batch distances, with δ = 10^-3 preventing division by a vanishing scale.Source: §A.3 · pp. 13–14

Attractive drift

  • Attraction uses weights normalized over positives only, concentrating on the nearest positive with a small bandwidth.Sudoku has a unique target, while TinyGSM uses between K and K + 4 queued positives.Source: §A.3 · p. 14

Repulsive drift

  • Repulsion combines three temperatures—0.02, 0.05, and 0.2—to capture both local and global structure among generated samples.The lowest temperature concentrates weights on the nearest generated sample; the highest spreads them across samples.Source: §A.3 · p. 14

Drift field and loss

  • The final field sums RMS-normalized fields across temperatures, using coefficients (λ+, λ−) = (2, 1) in experiments unless noted otherwise.The original Drifting Model uses (1, 1).Source: §A.3 · p. 14
  • A stop-gradient target adds the drift to each scaled sample, and the loss is applied only at the sampled terminal round.Gradients flow only through the scaled generated sample; preceding rounds do not receive gradients.Source: §A.3 · pp. 14–15

A.4 Coefficients of the Attractive and Repulsive Drifts

The coefficient comparison favors (2, 1) with two generated samples on both tasks, while the four-sample outcome differs between Sudoku and GSM8K. The authors relate this pattern to whether the drift supplies a restoring force toward positives.

Sudoku

  • The authors selected (2, 1) on separate Sudoku development data, used it on both tasks, and chose n = 2 to limit training cost.The selection did not use S-FLM benchmark evaluation scores or GSM8K test scores; n = 2 approximately halves training cost relative to n = 4 in the source formulation.Source: §A.4 · p. 15
  • With n = 4, both coefficient choices train, and (1, 1) achieves higher Sudoku accuracy from NFE=8 onward.The comparison is in Table 5's Sudoku Hard benchmark ablation.Source: §A.4 · p. 15, Table 5 · p. 16

GSM8K

  • On GSM8K, (1, 1) remains at 0.0% at every evaluated NFE for n = 2 and n = 4, while (2, 1) with n = 2 reaches 29.4% at 250,000 updates.At 40,000 updates, (2, 1) with n = 2 reaches 16.1% at NFE=16; the (1, 1) runs were stopped after evaluations at 20,000 and 40,000 updates.Source: §A.4 · p. 15

Sudoku

  • At n = 2, (1, 1) collapses to 0.0% at every evaluated NFE for all three seeds, while (2, 1) reaches 92.1% at NFE=16.Table 5 reports Sudoku Hard accuracy at 20,000 updates for the source formulation.Source: §A.4 · p. 15, Table 5 · p. 16

A possible reason: the restoring force toward the positives

  • With (1, 1), the n = 2 repulsive factor exceeds 1 and pushes samples farther apart; with (2, 1), it preserves separation while the restoring force pulls samples toward positives.The authors report training collapse with (1, 1) on both tasks in this setting.Source: §A.4 · p. 16
  • For n = 4, the (1, 1) target still lacks an explicit restoring term, which the authors conjecture contributes to GSM8K collapse.They note that generated samples can be attracted to different positives on GSM8K.Source: §A.4 · p. 16

A.5 Training Procedure of the Recurrent Model

Training samples a recurrence length for each problem and applies the drifting loss at its terminal round. Earlier rounds run without gradients, so backpropagation is limited to the terminal-round outputs.

Step 1: sampling the number of rounds

  • A round count is sampled from 1 through Rmax; GSM8K draws it per problem, while Sudoku shares one draw across the mini-batch.For GSM8K, each intermediate round batches trajectories that have not reached their terminal round.Source: §A.5 · p. 16

Step 2: encoding the prompt and the positives

  • Prompt and positive latents are cached because the encoder is frozen; in TinyGSM, only newly queued programs are encoded during training.GSM8K prompt and positive latents are standardized with fixed constants.Source: §A.5 · p. 16

Step 3: intermediate rounds without gradient

  • Intermediate rounds are stop-gradient forwards: the first transports noise, later rounds refine the previous output, and no decoding or re-encoding occurs between rounds.No activations are stored, so these passes cost only forward computation.Source: §A.5 · p. 17

Step 4: terminal round with gradient

  • At the terminal round, trajectories produce generated samples with gradients; with the verifier, samples are decoded and executed and verified programs enter the positive queue.When the sampled round count is one, the terminal round is the first round.Source: §A.5 · p. 17

Steps 5–9: drift field, loss, and update

  • The drift field, stop-gradient targets, and loss are computed from terminal samples and positives, with gradients flowing only through terminal outputs.The gradient norm is clipped to 1 before an AdamW step and EMA update.Source: §A.5 · p. 17

A.6 Training Step with the Verifier and the Positive Queue

This appendix describes verifier-assisted training and iterative generation, including decoder and positive-data construction and experimental and computational settings; the verifier relies on agreement with a voted pseudo-reference rather than ground truth. It also reports training and generation costs, supplementary Sudoku and GSM8K results, ablations, and tests of embedding extrapolation and single-round model iteration.

A.6 Training Step with the Verifier and the Positive Queue

  • The verifier uses the unique largest answer group among K positives as a pseudo-reference, and disables verification for tied or nontransitive groups.Answers are grouped within 10^-3 of each group’s minimum; agreement with the selected reference does not establish ground-truth correctness.Source: §A.6 · p. 18
  • Verified programs are executed at the terminal round, added to the positive queue, and excluded from repulsion after normalization while remaining in the loss.On Sudoku, the queue never grows because every verified answer is the given positive; verification acts only through the repulsion mask.Source: §A.6 · p. 18

A.7 Generation Procedure

  • Generation encodes the prompt once, refines the draft N times with the same model, clamps the round embedding past training depth, and decodes only the final draft.Source: §A.7 · p. 18

A.8 Decoder

  • 99.98% of tokens and 98.18% of programs are recovered exactly on clean CodeT5 latents, but these results do not measure latents generated by ReLaDrift.The separate bidirectional Transformer decoder is trained before the Drifting Model and remains frozen during its training and at test time.Source: §A.8 · pp. 18, 20

A.9 Generating the Positives

  • The training set contains 1,588,781 problems with 12 kept solutions each, and all 12 return the same value for 62.7% of problems.Solutions are not filtered for correctness, and neither generated nor original solutions are guaranteed correct.Source: §A.9 · p. 20

A.10 Experimental Settings

  • At NFE=16, DFM reaches 31.9%, 35.2%, and 33.4% with n_g=4, 8, and 16, versus 29.2% with n_g=2.These are means over three seeds; the authors say the gains are within seed standard deviation and do not change DFM’s ranking against ReLaDrift.Source: §A.10 · pp. 21–22
  • At NFE=180, FRM with the SDE sampler reaches 79.4% versus 77.2% for FMLM+, while FMLM+ leads parallel generators from NFE 4–128.The authors report that their Sudoku runs of S-FLM and MDLM are within 2.4 points of published accuracies, while FMLM+ is about 4 points below published results.Source: §A.10 · p. 22

A.11 Training and Generation Cost

  • With expected depth 3 and two source trajectories, the source formulation runs six forward passes per update, only two with gradients, and costs about 3.3 times a one-pass method.The additive formulation uses four forward passes and costs about 2.7 times a one-pass method.Source: §A.11 · p. 22
  • On GSM8K, an update takes 0.56, 0.65, and 0.76 s for K = 1, 4, and 12, respectively, totaling about 39, 45, and 53 h for 250,000 updates.Generating the 1,319 test programs takes ReLaDrift 6.8 s at NFE=4 and 13 s at NFE=16, compared with 5,200 s for S-FLM at NFE=1,024.Source: §A.11 · p. 25

A.12 Additional Sudoku Results

  • With n = 2, the Sudoku verifier improves accuracy by 1.1–1.7 points at NFE=2–5, but results are within 0.3 points at NFE=1 and NFE ≥8.At NFE=16, accuracy is 92.0% with the verifier versus 92.1% without it.Source: §A.12 · p. 25

A.13 Results without the Clamp of the Round Embedding

  • Without the round-embedding clamp, some Sudoku runs collapse, while all three source-formulation seeds with the clamp remain at 88–93% accuracy at NFE=8 and 16.The collapse includes all three Rmax = 3 seeds reaching 0.0% at NFE ≥8; the clamp costs little on GSM8K and prevents collapse on Sudoku.Source: §A.13 · pp. 25–26

A.14 GSM8K Ablations and Training Progress

  • On GSM8K, removing the verifier and positive queue lowers NFE=16 accuracy from 29.4% to 20.0%, while equal drift coefficients (1, 1) collapse to 0.0%.Increasing attraction bandwidth to 0.1× the median lowers accuracy to 10.3%, 24.0%, and 26.3% at NFE=1, 4, and 16, respectively; context-free embeddings reach 17.7% or 7.1% at NFE=16.Source: §A.14 · p. 26
  • On GSM8K, the additive formulation stays within about one point of the source formulation, whereas training with Rmax = 8 or 16 lowers accuracy at every NFE.The authors attribute the depth penalty to each round receiving loss for only a 1/Rmax fraction of problems when the terminal round is sampled uniformly.Source: §A.14 · p. 26

A.15 Iterating the Single-Round Drifting Model at Test Time

  • Iterating an untrained-for-refinement single-round model does not improve accuracy: Sudoku remains 4.9–5.0% at NFE=2–128, and GSM8K remains 7.6–7.8% at NFE=2–1,024.The model feeds each decoded and re-encoded output back as noise, but has no round embedding or refinement projections; on Sudoku, about 1,800 of 2,000 puzzles stop changing by NFE=8.Source: §A.15 · p. 27

A.16 All GSM8K Results of ReLaDrift

  • Appendix A.16 lists ReLaDrift GSM8K test accuracy at every evaluated checkpoint and NFE using the clamped round embedding.For NFE ≤5, the clamp does not change the round embedding.Source: §A.16 · p. 27

A.17 Generated Samples

  • Among 1,319 GSM8K test problems, 221 are wrong at NFE=1 and correct at NFE=16, while 7 change from correct to wrong.The passage presents generated ReLaDrift samples, including an example that is correct at NFE=1 and another that becomes correct at NFE=4.Source: §A.17 · pp. 27, 29
Loading 2610.23225v1…