GeoBlock: Inferring Block Granularity from Dependency Geometry in Diffusion Language Models

1. Problem Statement

Diffusion Language Models (DiffLMs) offer an alternative to autoregressive (AR) generation by enabling parallel token refinement. Block diffusion extends this by updating contiguous blocks of tokens in parallel while maintaining causal ordering between blocks. However, a critical bottleneck exists: block size selection.

• Current Limitations: Existing strategies rely on fixed schedules or heuristic signals (e.g., token confidence, entropy, or volatility). These methods treat block size as a hyperparameter or a function of local uncertainty.
• The Core Issue: These heuristics fail to account for the dependency geometry of the text.
  • Causal Regions: Some text segments have strong sequential dependencies (e.g., logical reasoning steps) requiring small, sequential updates.
  • Parallel Regions: Other segments are semantically cohesive with dense mutual interactions, allowing large parallel updates.
• The Gap: Using a uniform or confidence-based block size often leads to premature refinement in causal regions (causing errors) or inefficient small updates in parallel regions. The paper argues that block granularity should be determined by the underlying structural dependency of the tokens, not just their individual uncertainty.

2. Methodology: GeoBlock

The authors propose GeoBlock, a training-free, geometry-aware framework that dynamically infers block boundaries during decoding based on the model's internal attention patterns.

A. Theoretical Foundation: Dependency Geometry

GeoBlock views the decoding process through the lens of dependency geometry. It posits that a valid block for parallel refinement must be a self-contained dependency unit.

• Strong Internal Coupling: Tokens within the block must rely heavily on each other.
• Past Anchoring: Tokens must be conditioned on the already resolved history.
• Minimal Future Leakage: Tokens should not depend heavily on unresolved future tokens.

B. The Algorithm

GeoBlock operates at a "frontier" position $y$ (where tokens $1:y$ are partially refined and $y+1:L$ are masked). It evaluates candidate split points $x < y$ to define a candidate block $C = \{x+1, \dots, y\}$.

1. Attention as Dependency Proxy: The method uses the self-attention matrix $A$ from the model as an observable proxy for dependency structure.

2. Fused Attention: To ensure robustness, GeoBlock aggregates attention from selected layers and heads (typically mid-to-high semantic layers) into a unified attention map.

3. Closure Score Calculation: For a candidate block $C$, it computes three metrics based on accumulated attention mass:

   • $S_{C \to C}$: Internal coupling within the block.
   • $S_{C \to H}$: Conditioning on the history (past).
   • $S_{C \to F}$: Dependency leakage to the future (unresolved tokens).

   The Closure Score is defined as:
   $$\text{Score}(x) = \frac{S_{C \to C} + \alpha S_{C \to H}}{S_{C \to C} + \alpha S_{C \to H} + S_{C \to F}}$$
   Where $\alpha$ balances internal cohesion and past anchoring. A high score indicates the block is structurally self-contained and safe to update jointly.

4. Boundary Selection (Right-Shift Rule): Instead of picking the strict maximum score (which might be too conservative due to noise), GeoBlock selects the rightmost candidate boundary $x^*$ such that $\text{Score}(x^*) \geq \text{Score}_{\max} - \delta$. This allows for maximal parallelism while staying within a tolerance of optimal structural stability.

C. Computational Efficiency

• Training-Free: Requires no additional model training or fine-tuning.
• Zero Forward Pass Overhead: It operates entirely on attention matrices already computed during the standard decoding step.
• Linear Complexity: Dependency calculations are performed within a bounded frontier window, adding negligible computational cost (approx. 11% additional Number of Function Evaluations).

3. Key Contributions

1. Structural-Geometric Perspective: Reframes block diffusion decoding from a heuristic scheduling problem to a structural inference problem based on dependency geometry.
2. GeoBlock Framework: Introduces a novel, training-free method to infer adaptive block boundaries directly from attention-derived dependency structures.
3. Empirical Validation: Demonstrates that dependency geometry is a superior principle for block selection compared to confidence-based heuristics, achieving state-of-the-art accuracy-efficiency trade-offs.

4. Experimental Results

The authors evaluated GeoBlock on LLaDA-1.5 and Dream-v0-Base-7B across four benchmarks: GSM8K (Math), MATH, IFEval (Instruction Following), HumanEval, and MBPP (Code).

• Accuracy vs. Efficiency: GeoBlock consistently outperformed vanilla fixed-block decoding and dynamic confidence-based baselines (like AdaBlock).
  • On GSM8K and IFEval, GeoBlock achieved the highest accuracy with comparable or slightly increased NFE (Number of Function Evaluations).
  • On HumanEval and MBPP, it showed competitive or improved performance, particularly in medium block size regimes where reliable multi-token commitment is critical.
• Block Size Dynamics: GeoBlock naturally inferred block sizes between 13–19 tokens on average, adapting to the task. It formed larger blocks in semantically cohesive regions and smaller blocks in reasoning-heavy sequences.
• Ablation Studies:
  • Anchoring ($\alpha$): Including past conditioning ($\alpha > 0$) was crucial for stability.
  • Tolerance ($\delta$): A moderate right-shift tolerance ($\delta=0.1$) provided the best trade-off, allowing block expansion without sacrificing accuracy.
  • Layer Selection: Mid-to-high semantic layers provided the most reliable dependency signals.

5. Significance and Impact

• Paradigm Shift: Moves the field away from treating block size as a static hyperparameter or a function of token uncertainty, toward a structure-aware approach.
• Reliability: By respecting the causal geometry of language, GeoBlock prevents the "premature refinement" errors common in diffusion models, making them more viable for complex reasoning tasks.
• Practicality: As a training-free, drop-in replacement for existing schedulers, it offers an immediate path to improving the efficiency and quality of Diffusion Language Models without architectural changes.

In summary, GeoBlock proves that the "geometry" of token dependencies, observable via attention mechanisms, is the key to unlocking efficient and accurate parallel decoding in Diffusion Language Models.