TiDAR: Think in Diffusion, Talk in Autoregression

In plain words

Language models usually write one token at a time. Each step has to reload most of the model from memory, so the speed is limited by memory traffic rather than by arithmetic. Diffusion-style language models can fill several positions at once, which is faster, but when they do so their text tends to get worse.

TiDAR tries to get both. One model guesses several upcoming tokens in parallel, the way a diffusion model would. In the same pass it checks the guesses from the previous step in the usual left-to-right way, and it keeps the ones that match. The trick is an attention mask that lets one forward pass treat the accepted text, the old guesses and the blank slots for the next guesses differently. The model starts from an ordinary autoregressive language model and is trained to do both jobs. The authors report several times more tokens per second than the autoregressive model it came from, with quality close to that model on coding and math tests.

TL;DR

TiDAR is a sequence-level hybrid of a diffusion-style drafter and an autoregressive verifier, sharing one transformer. A training mask makes the clean sequence causal and a masked copy bidirectional within each block, so both output heads come from one pass §3.1 fig. 3. Each decoding step runs one forward pass that verifies the previous drafts and pre-drafts the next block §3.2 fig. 1. The author-reported result is 4.71x (1.5B) and 5.91x (8B) tokens per second over the matching AR baselines, at batch size 1 on one H100 §4.3 fig. 4.

Problem

Autoregressive decoding generates one token per forward pass and is memory-bound at small batch sizes. The paper argues that extra token slots in the same pass cost little latency until the workload becomes compute-bound §1 fig. 1.

Masked diffusion language models can predict several tokens per step. Their best quality, however, comes from decoding one token per denoising step, and decoding several tokens independently per step lowers quality §1 §2.1. The paper also notes that diffusion models have lacked exact KV caching because of bidirectional attention §2.1.

Speculative decoding usually uses a separate, weaker drafter, or extra heads whose drafting is sequential and waits for verification §2.2 tab. 1. The paper's stated assumptions are that the drafting capacity should come from the base model itself and that drafting and verification should share one forward pass §2.2.

Method

Overview. One network is trained so that it gives both an autoregressive (joint, causal) distribution and a diffusion (marginal) distribution from one forward pass §3.1. At inference, each step verifies the previous draft with the autoregressive head and pre-drafts the next block with the diffusion head §3.2 fig. 1.

Training (§3.1, fig. 3).

1. Build the input from the clean sequence of length S, followed by a copy of the same length in which every token is a mask token §3.1 fig. 3. The sequence length is therefore doubled §3.1.
2. Use one attention pattern. The clean part is causal. The mask part attends to the clean prefix, and inside the mask part attention is bidirectional within each block fig. 3 §3.1. The text of §3.1 describes the same idea with only the last block bidirectional and the rest causal; see Pitfalls.
3. Labels: for the causal part the target is shifted by one position (next-token prediction). For the masked part the target is the token at the same position §3.1.
4. Loss: the paper's displayed loss is a weighted sum of a next-token term and a diffusion term, divided by (1 + α), with α the loss balancing factor §3.1 (the equation is unnumbered in the source). α = 1 in most runs §3.3.
5. Every token of the masked part is a mask token. The paper does not use random partial corruption, which it reports as easier to balance and better for train–test consistency §3.1 §4.4.4 tab. 5.

# TiDAR training loss for one sequence  [§3.1]
# Input:  tokens (B, S) clean ids; alpha (loss balance)
# Output: scalar loss

def tidar_loss(tokens, alpha):
    # 1. Append a copy where every token is the mask token  [§3.1]
    masked = full_like(tokens, mask_id)
    inputs = concat([tokens, masked], dim=1)     # (B, 2S)

    # 2. One pass with the hybrid causal/bidirectional mask  [§3.1] [fig. 3]
    logits = model(inputs, mask=hybrid_mask)     # (B, 2S, V)
    ar_logits = logits[:, :S]                    # clean part, causal
    diff_logits = logits[:, S:]                  # masked part

    # 3. Next-token loss: position i predicts token i+1  [§3.1]
    ar_loss = cross_entropy(ar_logits[:, :-1], tokens[:, 1:])

    # 4. Diffusion loss: masked position i predicts token i  [§3.1]
    diff_loss = cross_entropy(diff_logits[:, :-1], tokens[:, :-1])

    # 5. Weighted combination, as in the displayed loss  [§3.1]
    return (alpha * ar_loss + diff_loss) / (1 + alpha)

Generation (§3.2, fig. 1).

1. Prefill: encode the prompt causally, and in the same way draft one block in parallel with bidirectional attention §3.2 app. B fig. 7.
2. Each decoding step is one forward pass over three parts §1 fig. 1: the cached prefix, the drafts proposed at the previous step (causal), and mask slots for the next step's pre-drafts.
3. Verify: the autoregressive prediction for the current step is compared with the previous drafts. Drafts that match are accepted §3.2.
4. Cache: KV entries of the forwarded causal tokens are kept, and the entries of rejected tokens are evicted later §3.3 fig. 1.
5. Pre-draft: the diffusion part proposes the next step's block conditioned on every possible accepted length. The block for the accepted length is passed on §1 fig. 1.
6. Drafting uses one diffusion step, not several denoising steps. The paper reports one step gives drafts whose quality is high enough for a high acceptance rate §3.2.

# One TiDAR decoding step: one forward pass  [§3.2] [fig. 1]
# Input:  cache (KV of the accepted prefix), drafts (block_len,)
# Output: accepted tokens, drafts for the next step

def decode_step(cache, drafts):
    # 1. One pass over: cached prefix, drafts (causal), mask slots  [§3.3]
    ar_logits, diff_logits = model(drafts, cache, mask=decode_mask)

    # 2. AR prediction for each draft, from the current step  [§3.2]
    ar_preds = argmax(ar_logits)                # predicts drafts  [§1]

    # 3. Accept the longest prefix where drafts match  [§3.2]
    accepted = longest_matching_prefix(drafts, ar_preds)

    # 4. Keep KV of accepted tokens; evict rejected ones  [§3.3] [fig. 1]
    cache = commit_and_evict(cache, accepted)

    # 5. Pre-drafts were built for each possible accepted length  [§1]
    pre_drafts = argmax(diff_logits)            # layout not stated
    next_drafts = pre_drafts[len(accepted)]     # index is inferred
    return accepted, next_drafts

Inference optimisation (§3.3, fig. 3 right, fig. 7). The inference mask is pre-built once with size q_len × (q_len + max_sequence_len), where the paper gives q_len = block_len · (1 + block_len) in a footnote. Each step slices it, after the draft part and the prefix are reordered so that the slice is valid §3.3 fig. 3. The paper states that no inference hyperparameters need tuning apart from the draft length and the optional way of combining the two heads §3.3 §4.4.3. The draft length can be changed at decoding time without retraining §5.

Mixing the two heads (§4.4.3, fig. 6). When the two heads disagree at the same position, the paper proposes taking the argmax of β · logits_AR + (1 − β) · logits_diff and comparing it with the draft §4.4.3. The value of β used for the reported runs is not stated. Table 2's 8B rows use "Trust AR" and "Trust Diff" tab. 2. Figure 6 varies β and α, but its values are only in the figure fig. 6.

Hyperparameters

Results

Author-reported. Setup: continual pretraining from Qwen2.5 1.5B and Qwen3 8B with 50B and 150B tokens §4.1; single H100, batch size 1, native PyTorch with FlexAttention §4.3 §5; baselines include Qwen2.5, Qwen3, Llama3.2, SmolLM2, Dream, LLaDA, Block Diffusion trained with the same recipe, and EAGLE-3 speculative decoding on Qwen3-8B-Instruct §4.1 §4.3.

The per-point throughput numbers of fig. 4 and the per-point curves of fig. 5 are not given in the text; only the averages above are.

Table 2: generative (coding and math), one token per NFE for diffusion baselines §4.2.1 tab. 2.

Superscripts: * = Qwen2.5 1.5B init, † = Qwen3 4B init, ‡ = Qwen3 8B init tab. 2.

Table 3: likelihood (knowledge and commonsense) §4.2.2 tab. 3. The diffusion baselines use Monte Carlo likelihood. Rows marked "MC128" are evaluated with 128 steps. TiDAR uses its causal mask tab. 3.

Table 4: decoding strategies at 1.5B (ablation) §4.4.2 tab. 4. "T/NFE" is the average tokens per forward pass.

Table 5: masking strategy at 1.5B (ablation) §4.4.4 tab. 5.

Table 6: benchmark configuration app. A tab. 6.

The paper states that HumanEval and HumanEval+ use "standard post-processing" for all models, without giving the procedure app. A.

Other reported points: the speed-up is measured in native PyTorch and the authors say custom kernels, KV cache management and scheduling could change the ratios §4.3. The paper notes that speculative decoding can guarantee exactly the same output as the base model, so the two should be used for different purposes §4.3. The paper reports a higher acceptance rate than the public EAGLE-3 weights and a better conversion from tokens per NFE to tokens per second §4.3.

Paper vs code

No official code.

Searches: the paper text and the arXiv metadata link no code or project page §1 §4 [metadata]. GitHub repository search for tidar, for TiDAR diffusion autoregression, and for think in diffusion talk in autoregression returned two third-party repositories, which were not used: irfannaqieb/TiDAR (no license; "small implementation") and fish4terrisa-MSDSM/tidar (AGPL-3.0; labelled unofficial). Searches with org:NVlabs and org:nvidia returned nothing.

Pitfalls

• Loss range. The displayed loss sums the next-token and diffusion terms over i = 1 … S−1, each divided by S−1 §3.1. The text says the diffusion loss is applied to every token of the masked section, so that both losses have the same number of terms §3.1. The pseudocode follows the displayed loss.
• Mask shape. §3.1 says only the last block of the masked copy is bidirectional and the prefix is causal. The captions of fig. 2 and fig. 3 say blocks are bidirectional within themselves and block-causal across blocks. These give different masks when the sequence has more than one block. The paper does not say which one training used.
• Acceptance rule. The text says drafts are accepted when they match the autoregressive prediction, and the ablation compares an argmax with the draft §3.2 §4.4.3. The paper does not say whether acceptance is greedy or sampled, what happens at the first mismatch, or whether the autoregressive token at the mismatch is emitted, as in standard speculative decoding.
• Pre-draft layout. The paper says pre-drafts are conditioned on all possible accepted lengths §1 and that one is selected §1 fig. 1. The tensor shape and indexing are not stated; the pseudocode's indexing is inferred.
• Two predictions per position. The autoregressive and diffusion heads both predict the same position at inference §4.4.3. Which one is trusted for verification changes the 8B generative average (Trust Diff 65.31% vs Trust AR 63.90%) tab. 2. The paper says the values of β used are not in the text.
• Training cost. Doubling the sequence length during training is a stated cost §3.1 §5. The paper leaves long-context extension to future work §5. Long-context behaviour is not evaluated.
• Training data. The paper gives token counts (50B and 150B) and the batch size, but does not name the training corpus or its mixture §4.1.
• Inference speed. The 4.71× and 5.91× figures use batch size 1 and native PyTorch with FlexAttention §4.3 §5. They may not carry over to batched serving or custom kernels §5.
• Baseline mismatch. The EAGLE-3 comparison uses public weights for Qwen3-8B-Instruct, because no base-model weights exist §4.3 fig. 4. The other comparisons use base models.
• Likelihood comparison. TiDAR uses its causal mask for likelihood, while the diffusion baselines use Monte Carlo estimates tab. 3. The paper calls the comparison aligned with AR likelihoods, but the estimators differ.
• Figures without numbers. Per-point throughput fig. 4, Pareto curves fig. 5 and the β and α sweep fig. 6 appear only as figures. The text does not give their values.
• Exact output. The paper says quality is guaranteed by autoregressive rejection sampling §4.4.3. It does not say the output is identical to the autoregressive model's own decoding.

Lineage

• Builds on: Block Diffusion (Arriola et al., 2025), with the change of keeping only the decoding block bidirectional §3.1 §2.1; Set Block Decoding (Gat et al., 2025) for next-token loss on the prefix §3.1; one-step diffusion drafting (Liu et al., 2025) §3.2; the pre-draft idea from Apple's MTP work (Samragh et al., 2025) §3.2; speculative decoding (Leviathan et al., 2023) §2.2; FlexAttention (Dong et al., 2024) §3.3.
• Initialised from: Qwen2.5 1.5B and Qwen3 8B base models §4.1.
• Compared against: Qwen2.5, Qwen3, Llama3.2 and SmolLM2 §4.1; Dream and LLaDA §4.2; Block Diffusion under the same recipe §4.1; EAGLE-3 on Qwen3-8B-Instruct §4.3.
• Superseded by: not stated in the content this wiki holds.