🤖 AI

Prefix-Adaptive Block Diffusion for Efficient Document Recognition

This paper proposes the Prefix-Adaptive Block Diffusion Model (PA-BDM), which enhances efficient document recognition by replacing fixed block boundaries with dynamic prefix commitment and causal denoising to resolve information flow inconsistencies, thereby achieving superior accuracy and a 71.6% increase in inference throughput compared to existing models.

Original authors: Mingxu Chai, Ziyu Shen, Chenyu Liu, Kaidi Zhang, Jiazheng Zhang, Dingwei Zhu, Zhiheng Xi, Ruoyu Chen, Jun Long, Jihua Kang, Tao Gui, Qi Zhang

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mingxu Chai, Ziyu Shen, Chenyu Liu, Kaidi Zhang, Jiazheng Zhang, Dingwei Zhu, Zhiheng Xi, Ruoyu Chen, Jun Long, Jihua Kang, Tao Gui, Qi Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to transcribe a complex, handwritten math formula or a dense table from a document image into a computer-readable format.

The Old Way (The "Rigid Block" Problem)
Think of current efficient AI models (called Block Diffusion Models) as a team of scribes working in shifts. They divide the page into fixed-size chunks, say 32 words at a time.

The Bottleneck: They work on all 32 words simultaneously (which is fast), but they can't save their progress or "lock in" any of those words until the entire chunk of 32 is finished.
The Confusion: Inside that chunk, the scribes can look at each other's work from left to right and right to left. While this seems helpful, it creates a mess when they move to the next chunk. The next chunk only knows what came before it (left-to-right), but the previous chunk was a jumbled mix of directions. This inconsistency makes it hard to get the structure of things like math formulas or tables right.
The Waste: As they finish words inside a chunk, the "parallel" work slows down because fewer words are left to guess. It's like a factory assembly line that stops being efficient once half the products are made.

The New Solution: PA-BDM (The "Adaptive Scribe")
The authors propose a new method called Prefix-Adaptive Block Diffusion (PA-BDM). Here is how it changes the game using simple analogies:

1. The "Candidate Range" vs. The "Fixed Box"

Instead of treating a block of 32 words as a rigid box that must be filled completely before moving on, PA-BDM treats it as a maximum candidate range.

Analogy: Imagine you are filling a bucket with water. The old method says, "You must fill the whole bucket before you can pour any water into the storage tank." PA-BDM says, "Fill the bucket as much as you can, and as soon as you are sure a cup of water is clean and safe, pour it into the storage tank immediately."
Result: The AI doesn't wait for the whole block to finish. It grabs the "reliable" part (the prefix) and saves it instantly.

2. Causal Flow (The "One-Way Street")

The old method let scribes look backward and forward within a block, which confused the order of things. PA-BDM forces everyone to look only forward (from left to right), just like reading a book.

Analogy: In the old system, a scribe in the middle of a sentence could peek at the end of the sentence to guess the middle. This worked for casual chat but failed for strict structures like math equations where order matters. PA-BDM enforces a strict "one-way street" rule, ensuring the AI learns the correct sequence every time.

3. Progressive Prefix Commitment (PPC) – The "Confidence Check"

This is the engine of the new system. As the AI guesses the next batch of words, it checks its own confidence.

How it works: If the AI is 99% sure about the first 10 words of a 32-word block, it "commits" (locks in) those 10 words immediately. It then throws away the remaining 22 guesses and starts a fresh batch of 32 guesses based on those 10 locked words.
The Benefit: This resets the "parallel space." Instead of working on a shrinking list of 22, then 10, then 5 words, the AI gets to work on a fresh, full set of 32 words again and again. This keeps the speed high.

4. Confidence-Gated Structural Loss (CSL) – The "Strict Teacher"

During training, the AI learns by trying to fix mistakes. The old method forced the AI to learn from all mistakes, even if the beginning of the sentence was already shaky.

The Fix: PA-BDM uses a "confidence gate." If the AI is unsure about the start of a sentence, the teacher (the training algorithm) stops grading the rest of the sentence. It only grades the part the AI is confident about.
Analogy: Imagine a teacher grading a math test. If the student gets the first step wrong, the teacher doesn't waste time grading the final answer because it's based on a wrong premise. The teacher waits until the student gets the first step right before grading the rest. This prevents the AI from learning "noisy" or bad patterns.

The Results

The paper claims that this new approach, PA-BDM, is a massive upgrade:

Speed: It is 71.6% faster than the previous best diffusion model (MinerU-Diffusion) and about 8 times faster than standard autoregressive models (which write one word at a time).
Accuracy: It is more accurate at recognizing complex structures like math formulas, tables, and diagrams because it respects the strict order of tokens.
Efficiency: It uses memory just as efficiently as the older models but gets much more work done in the same amount of time.

In short, PA-BDM stops the AI from waiting for "perfect blocks" and instead lets it lock in "good enough" progress immediately, keeping the assembly line moving at top speed without losing accuracy.

Technical Summary: Prefix-Adaptive Block Diffusion for Efficient Document Recognition

1. Problem Statement

Document parsing requires converting document images into machine-readable formats (e.g., LaTeX, HTML, Markdown). While Autoregressive Models (ARMs) have achieved substantial progress, their strictly token-by-token generation paradigm limits inference efficiency. To address this, recent research has explored parallel generation paradigms, specifically Block Diffusion Models (BDMs). BDMs generate blocks autoregressively while denoising tokens within each block in parallel, offering flexible-length output and KV caching.

However, the authors identify two critical limitations in standard BDMs when applied to structure-sensitive document recognition:

Fixed Block-Level Commitment: Standard BDMs bind denoising and KV cache commitment to fixed block boundaries. Tokens cannot be cached until the entire block is completed, preventing the prompt reuse of reliable intermediate predictions. Furthermore, as intra-block denoising proceeds, the number of remaining masked tokens decreases, causing the effective parallel decoding space to shrink.
Inconsistent Information Flow: Standard BDMs employ bidirectional denoising within a block but maintain left-to-right autoregression between blocks. This creates inconsistent information flow. For tasks requiring precise reconstruction of discrete token sequences (like formulas and tables), this bidirectional context can conflict with the strict structural boundaries and token order required, leading to structural generation errors.

2. Methodology: Prefix-Adaptive Block Diffusion Model (PA-BDM)

The authors propose PA-BDM, a framework that redefines the role of the "block" from a fixed generation unit to a maximum candidate generation range. The core components are:

A. Causal Intra-Block Denoising

Unlike standard BDMs that use bidirectional attention within a block, PA-BDM employs causal denoising (prefix-to-suffix) inside each candidate block.

Mechanism: Tokens within a block attend only to previous blocks and their left-side context within the current block.
Benefit: This aligns the intra-block information flow with inter-block autoregression, ensuring that reliable prefixes generated within a block are valid for KV-cache reuse.

B. Confidence-gated Structural Loss (CSL)

To support causal denoising during training, PA-BDM introduces CSL to manage supervision noise.

Problem: In causal denoising, later masked tokens depend on earlier tokens. If early tokens are uncertain (high entropy), supervising the entire suffix introduces noisy gradients.
Solution: CSL calculates the confidence of predicted tokens. It identifies the first position in the masked suffix where confidence drops below a threshold $\tau$ . Supervision is applied only up to this "learning frontier." Positions to the right of this frontier are excluded from the loss calculation.
Training Strategy: This encourages the model to build low-entropy, reliable structural prefixes before extending to longer continuations.

C. Progressive Prefix Commitment (PPC)

During inference, PA-BDM utilizes PPC to dynamically manage generation and caching.

Process:
1. The model predicts a candidate range of size $D$ (e.g., 32 tokens).
2. PPC scans the predictions from left to right and identifies the longest contiguous prefix where all tokens exceed the confidence threshold $\tau$ .
3. This reliable prefix is committed to the KV cache immediately.
4. The unresolved suffix is discarded, and the next candidate range is reset starting from the newly committed prefix.
Benefit: This restores a large parallel decoding space at every step (avoiding the shrinking space of standard BDMs) and enables timely KV-cache reuse, significantly improving throughput.

3. Key Contributions

Identification of Limitations: The paper pinpoints fixed block-level commitment and inconsistent intra/inter-block information flow as the primary bottlenecks for BDMs in structure-sensitive document recognition.
PA-BDM Framework: A novel framework that treats block size as a maximum candidate range rather than a fixed commitment unit, utilizing causal intra-block modeling to enable adaptive generation.
PPC Mechanism: A decoding strategy that dynamically commits reliable prefixes, enabling immediate KV-cache reuse and resetting the candidate range to recover parallelism.
CSL Objective: A training loss that adapts supervision based on prefix confidence, reducing noisy gradients from unstable prefix states.

4. Experimental Results

The authors evaluated PA-BDM (3B parameters) against various baselines, including DiffusionVL, MinerU-Diffusion, and autoregressive models (e.g., Qwen2.5-VL, Dolphinv2) on text, formula, table, and diagram recognition benchmarks.

Accuracy: PA-BDM achieves the best performance among diffusion-based models on most benchmarks. It shows particularly strong gains on complex formula recognition (e.g., CPE and HWE subsets), outperforming DiffusionVL and MinerU-Diffusion.
Efficiency: PA-BDM improves inference throughput by 71.6% over the 2.5B MinerU-Diffusion baseline (267.2 tokens/sec vs. 155.7 tokens/sec). It achieves approximately 8× higher throughput than compared ARM recognizers.
Memory: Despite the efficiency gains, PA-BDM maintains peak GPU memory comparable to other 3B-scale models.
Ablation Studies:
- Causal vs. Bidirectional: Causal intra-block modeling outperforms bidirectional modeling on structured tasks (formulas, tables), suggesting bidirectional context can lead to structural invalidity.
- CSL: Models trained with CSL show higher accuracy and throughput than those with standard Cross-Entropy or random masking, indicating CSL successfully learns longer reliable prefixes.
- PPC: Adaptive commitment significantly outperforms fixed-length commitment strategies, balancing accuracy and parallelism effectively.

5. Significance and Claims

The paper claims that PA-BDM offers a superior speed–accuracy trade-off for document parsing compared to existing diffusion-based baselines. By treating blocks as candidate ranges and committing reliable prefixes, PA-BDM resolves the tension between parallel generation efficiency and the strict structural requirements of document recognition.

The authors emphasize that their approach is particularly effective for structure-sensitive tasks (LaTeX, HTML, Mermaid) where token order and structural boundaries are critical. They note that while the method improves efficiency significantly, the current evaluation is primarily conducted on English-centric, open-source datasets. The generalization to more complex multilingual scenarios (e.g., Chinese formulas, mixed layouts) remains an area for future study, as these scenarios involve richer character sets and more complex structural dependencies that may challenge reliable-prefix modeling.

The work suggests that reliable-prefix modeling is an effective direction for improving the efficiency of diffusion-based document recognition, moving beyond fixed-block constraints to adaptive, confidence-driven generation.