google

/

diffusiongemma-26B-A4B-it

like 1.11k

Follow

Google 61.1k

A Research Proposal for Long-Range Parallel Text Diffusion with Streaming Commit
Abstract

Current diffusion language models show strong potential for faster text generation because they refine multiple tokens in parallel instead of producing text strictly one token at a time. However, practical systems are still limited by small fixed canvases, block-level sequentiality, and weak long-range planning during streaming generation. For example, DiffusionGemma uses a fixed 256-token canvas and chains multiple canvases for long-form output: a full 256-token block is denoised, appended to the context, the KV cache is updated, and only then the next canvas starts. This makes the model faster than classic token-by-token decoding, but still restricts each active reasoning/generation unit to a small local block.

We propose Progressive Hierarchical Canvas Diffusion (PHCD): a decoding architecture where a large global canvas, for example 4096 tokens, is continuously refined in the background, while a smaller local active window, for example 256 tokens, is refined at higher quality and committed progressively. After the first local window is accepted, the remaining global canvas is not discarded; it is re-conditioned on the committed prefix and refined further in parallel. This creates a streaming system with long-range planning, local high-quality denoising, and continuous future repair.

1. Current Problem
   1.1 Autoregressive decoding is sequential

Standard autoregressive LLMs generate text left-to-right. Each new token depends on the previous generated tokens, so decoding is inherently serial. KV cache reduces recomputation over the prefix, but the model still has to produce the continuation step by step. This creates latency, especially for long answers, code generation, and agentic workflows. DiffuSpec describes this as a core latency bottleneck of autoregressive decoding: each token requires a serial forward step, and speculative decoding only partially reduces this by drafting multiple tokens and verifying them in parallel.

1.2 Current diffusion LMs are too local

Diffusion language models can refine a block of tokens in parallel, but current practical systems often use a small fixed canvas. DiffusionGemma uses Canvas Length 256 and performs block-autoregressive multi-canvas sampling for longer output. In other words, it generates one 256-token canvas, finalizes it, appends it to the context, updates cache, and then starts the next 256-token canvas.

This creates a problem:

The model can refine 256 tokens in parallel,
but it does not maintain a large live future plan.

So for long code, long reasoning, or structured documents, the model still behaves like a sequence of local blocks.

1.3 Existing block diffusion is a strong step, but still block-sequential

Block Diffusion improves over pure diffusion and pure autoregression by combining autoregressive dependencies between blocks with diffusion inside each block. The paper explicitly targets flexible-length generation, KV caching, and parallel token sampling.

But the core remaining limitation is:

Generation is still sequential at the block level.

The model can denoise tokens inside a block in parallel, but the global future is not continuously alive and repairable while the current block is being streamed.

1.4 Speculative diffusion helps, but draft length is fragile

DiffuSpec uses a pretrained diffusion language model as a drafter and a standard autoregressive model as verifier. This is close to the desired direction: diffusion proposes multiple tokens, AR verifies them. However, DiffuSpec also identifies two issues: diffusion drafts form a bidirectional token lattice that may not correspond to a valid left-to-right causal path, and draft length creates a speed-quality tradeoff.

So the current state is:

AR models: high quality, but slow.
Diffusion models: parallel, but local / fixed-canvas.
Block diffusion: better length handling, but block-sequential.
Speculative diffusion: useful, but draft length and causal consistency are hard.
2. Proposed Approach
Progressive Hierarchical Canvas Diffusion

The proposed system uses two diffusion levels running together:

Global Canvas: 4096 tokens

• low-resolution / coarse future plan
• continuously refined in the background
• keeps long-range structure alive

Local Window: 256 tokens

• high-resolution active generation area
• refined more aggressively
• committed progressively to the output

The key idea:

Do not fully denoise 4096 tokens before output.
Do not generate only isolated 256-token blocks.
Instead:
keep 4096 tokens alive as a rough future,
polish the next 256 tokens,
commit a safe prefix,
repair the remaining future in parallel.
3. Inference Loop
Input prompt
↓
Initialize global canvas G = 4096 masked/noisy token positions
↓
Start coarse global denoising over G
↓
Select local active window W = first 256 positions
↓
Run stronger local denoising on W
↓
Verifier / commit controller accepts first N tokens
↓
Append accepted tokens to final output
↓
Condition global canvas on committed prefix
↓
Repair remaining uncommitted canvas
↓
Slide local window forward
↓
Repeat

Important: the system should not necessarily commit the full 256-token local window. It can commit only the safe part.

Text generation: commit 32–128 tokens
Code generation: commit 8–32 tokens
Math / reasoning: commit 4–16 tokens
JSON / strict format: commit 4–16 tokens
4. Why This Should Be Better
4.1 Bigger active reasoning horizon

A 256-token canvas is often enough for short text, but weak for long code, planning, reasoning, and structured outputs. A 4096-token global canvas gives the model a rough view of the future answer:

intro → plan → implementation → tests → conclusion

This should reduce cases where the model starts in one direction and then later needs to contradict or rewrite itself.

4.2 Local quality without global waiting

The model does not need to fully perfect all 4096 tokens before streaming. Only the current local window must become high-confidence.

4096-token canvas = future map
256-token window = current rendered area

This is similar to progressive rendering: first rough global structure, then local high-quality refinement.

4.3 Continuous future repair

After the first 256 tokens are accepted, the remaining 3840 tokens are not thrown away. They are re-conditioned on the committed prefix and refined again.

Before commit:
[0–255 active] [256–4095 rough future]

After commit:
[0–255 fixed] [256–4095 repaired future]

This is the main difference from simple block diffusion.

4.4 Better streaming

The system can stream committed tokens while continuing to refine future tokens. On real hardware this is not “free parallelism”: processes still compete for compute and memory. But it can reduce perceived latency and improve throughput if scheduled correctly.

5. Architecture
   5.1 Core Components
   Shared Transformer Backbone
   ├── Causal Prefill Encoder
   │ └── reads prompt / committed prefix and builds cache
   │
   ├── Global Diffusion Decoder
   │ └── maintains 4096-token coarse canvas
   │
   ├── Local Refinement Decoder
   │ └── focuses compute on the next 256-token active window
   │
   └── Commit / Verifier Head
   └── decides how many tokens are safe to finalize

This could be implemented in two ways:

Option A: Single model, shared backbone

Best long-term direction.

One model
Different masks / heads / schedules
Lower memory overhead
Harder training
Option B: Separate models

Easier for experiments.

Diffusion model = drafter / canvas refiner
AR model = verifier / committer
More memory
Simpler research prototype

DiffuSpec already shows the feasibility of using a diffusion language model as a drafter while keeping a standard AR verifier.

6. Training Objective

The model should be trained explicitly for hierarchical generation, not only standard 256-token block denoising.

6.1 Global coarse loss

Train the model to reconstruct a long masked sequence at low precision:

Prompt + noisy 4096-token canvas → rough target continuation

The goal is not perfect token accuracy. The goal is long-range structure.

6.2 Local refinement loss

Train the active 256-token window to become high quality:

Prompt + committed prefix + rough global canvas + noisy local window
→ clean local window
6.3 Commit consistency loss

Train the model to predict which prefix of the local window is safe to finalize:

accept length N
confidence per token
rollback probability
syntax / format stability
6.4 Future repair loss

After a prefix is committed, train the remaining canvas to repair itself:

committed prefix + old rough future
→ corrected future continuation

This is the crucial part. Without this, the global canvas becomes stale after each commit.

7. Inference Scheduler

A practical scheduler could look like this:

while not done:
global_canvas = coarse_refine(global_canvas, committed_prefix, steps=low_nfe)

local_window = extract_next_window(global_canvas, size=256)
local_window = local_refine(local_window, committed_prefix, global_canvas, steps=high_nfe)

accepted_tokens = verifier_accept(local_window, max_commit=N)

stream(accepted_tokens)
committed_prefix += accepted_tokens

global_canvas = repair_remaining_canvas(
    global_canvas,
    committed_prefix=committed_prefix,
    accepted_tokens=accepted_tokens
)

slide_window()

The number of denoising steps should be adaptive:

Easy text: fewer local steps, larger commit
Code: more local steps, smaller commit
Reasoning: more verifier checks, smaller commit
Strict format: grammar-aware commit
8. Difference from Existing Methods
Method What it does Limitation
Autoregressive LLM generates one token at a time high latency
Speculative decoding small model drafts, large model verifies draft model often still AR/sequential
DiffuSpec diffusion drafts, AR verifies draft length and causal path are hard
Block Diffusion AR over blocks, diffusion inside blocks future is not continuously refined
DiffusionGemma 256-token canvas, multi-canvas sampling active canvas is small
PHCD 4096 global canvas + 256 local window + continuous repair requires special training/scheduler
9. Main Risks
9.1 Compute is not free

Running a 4096 global canvas and 256 local refinement in parallel may increase FLOPs. On a single GPU, “parallel” mostly means pipelined scheduling, not magical free speed.

The approach is only useful if:

global canvas uses cheap coarse steps
local window receives expensive steps
committed prefix reduces repeated work
future repair is cheaper than fresh generation
9.2 Commit errors are dangerous

If the model commits too early, it can lock in bad text. This is especially risky for code, math, JSON, and tool calls.

Solution:

small commit sizes
confidence thresholds
syntax-aware verifier
rollback buffer
optional AR verifier
9.3 Long canvas consistency is hard

A 4096-token rough canvas may contain contradictions. The system needs repair training and consistency objectives, otherwise the future plan becomes noisy and useless.

9.4 KV cache for diffusion is non-trivial

Diffusion models use bidirectional attention over the canvas, which does not naturally fit the classic AR KV-cache pattern. Recent work on dKV-Cache argues that diffusion LMs need cache-like mechanisms and proposes delayed KV caching for denoising, with reported 2–10x acceleration in their setting.

10. Evaluation Plan
    Speed

Measure:

tokens/sec
time-to-first-token
time-to-first-256-tokens
latency per 1K generated tokens
GPU memory
FLOPs per accepted token

Baselines:

AR decoding
AR + speculative decoding
DiffusionGemma-style 256 canvas
Block Diffusion
DiffuSpec-style diffusion drafting
Quality

Evaluate:

HumanEval / MBPP / LiveCodeBench for code
GSM8K / MATH / AIME-style tasks for reasoning
LongBench / Needle-in-Haystack for long context
JSON/schema adherence for structured output
agentic tool-use tasks
Consistency

Specific PHCD metrics:

future repair success rate
commit rollback rate
accepted tokens per refinement step
global-local contradiction rate
long-range plan preservation
11. Minimal Prototype

The first prototype does not need a new model from scratch.

Prototype A
Use existing diffusion LM as global/local drafter.
Use existing AR LLM as verifier.
Simulate 4096 canvas using chunked masked buffers.
Commit 4–32 tokens at a time.
Repair remaining buffer after every commit.

This would test whether the scheduling idea works before training a custom model.

Prototype B

Fine-tune a smaller open model:

base: 1B–7B transformer
global canvas: 1024 first, not 4096
local window: 128 or 256
task: code/text continuation
objective: coarse draft + local refine + repair

Then scale:

1024 → 2048 → 4096 canvas
128 → 256 → 512 local window
12. Expected Outcome

The hypothesis:

A continuously repaired long canvas should improve long-form coherence,
while local high-quality windows preserve streaming quality.

Expected advantages:

better long-range planning than 256-canvas diffusion
lower latency than pure AR decoding
better quality than committing large diffusion blocks blindly
better streaming than full-sequence diffusion
more stable code generation than small isolated canvases
13. Short Summary

Progressive Hierarchical Canvas Diffusion proposes a hybrid decoding strategy where a large global diffusion canvas, such as 4096 tokens, is refined continuously, while a smaller local window, such as 256 tokens, is refined to high quality and committed progressively. Unlike current 256-token canvas diffusion, the future continuation remains alive and repairable after every commit. Unlike pure autoregressive decoding, the model can plan and refine many future tokens in parallel. Unlike simple block diffusion, the uncommitted future is not discarded or regenerated from scratch; it is updated as a persistent latent text plan.