Achieve up to 1.50x speedup in diffusion language models by computing head-specific sparse attention patterns once during early denoising steps and reusing them across all subsequent iterations, while preserving full attention in critical early phases to maintain generation quality and accuracy.
SparseD: Sparse Attention for Diffusion Language Models
version
0.0.2
engine
skillxiv-v0.0.2-claude-opus-4.6
license
MIT
url
https://arxiv.org/abs/2509.24014
keywords
["sparse attention","diffusion language models","computational efficiency","denoising steps","head-specific patterns","temporal consistency","block-wise pooling","FlashAttention optimization"]
description
Achieve up to 1.50x speedup in diffusion language models by computing head-specific sparse attention patterns once during early denoising steps and reusing them across all subsequent iterations, while preserving full attention in critical early phases to maintain generation quality and accuracy.
Accelerate Diffusion Language Models with Sparse Attention Patterns
Problem Context
Diffusion language models (DLMs) offer a compelling alternative to autoregressive models by generating tokens in parallel rather than sequentially. However, they inherit the fundamental challenge of transformer architectures: attention's quadratic complexity O(n²) in sequence length. For a DLM performing k denoising steps on a sequence of length n, the computational cost becomes O(k·n²), making long-context inference prohibitively expensive.
Standard sparse attention techniques borrowed from autoregressive models (sliding-window, block-sparse patterns) fail to capture the unique characteristics of diffusion inference, where attention patterns evolve during the denoising process and early steps critically influence final generation quality.
Core Concept
SparseD leverages three empirical observations about diffusion language models to enable efficient sparse attention:
Head-Specific Patterns: Different attention heads exhibit distinct sparsity structures (column-wise focusing, sliding-window behavior, or mixed patterns), requiring individualized sparse masks rather than uniform patterns across all heads.
Temporal Consistency: Attention patterns within each head remain remarkably stable across denoising steps. Once a head's sparse pattern emerges in early steps, it persists throughout the remaining iterations.
Early-Step Criticality: Initial denoising steps disproportionately impact generation quality. Applying sparse attention too early degrades output significantly, but later steps tolerate sparse computation without quality loss.
SparseD capitalizes on these insights through pattern pre-computation and reuse: compute full attention during initial "critical" steps to extract head-specific sparse patterns, then reuse these masks across all remaining steps without recalculation. This one-time pattern computation cost is amortized across many denoising iterations.
Architecture Overview
The SparseD pipeline consists of four integrated components:
Isolated Selection: For each attention head independently, compute the full attention score matrix during early steps and select the top ρ% highest-scoring query-key pairs using block-wise average pooling. This creates a head-specific sparse mask.
Block-Wise Pooling: Partition the query and key sequences into fixed-size blocks. Compute average attention scores within each block to identify globally important positions without materializing the full n×n attention matrix.
Pattern Freezing: After the skip% initial denoising steps complete, freeze the computed head-specific sparse masks. These masks remain constant for all subsequent denoising iterations.
Sparse Masking Application: During later denoising steps, apply the precomputed sparse masks to restrict attention computation only to selected query-key pairs, implementing the sparse attention via PyTorch's FlexAttention mechanism.
Implementation
Step 1: Block-Wise Importance Pooling
The foundation of SparseD lies in identifying important token positions without computing full O(n²) attention matrices. This uses hierarchical pooling.
import torch
import torch.nn.functional as F
defcompute_block_importance(attention_scores, block_size=128):
"""
Compute block-level importance from raw attention scores.
Args:
attention_scores: [batch, heads, seq_len, seq_len] attention matrix
block_size: size of blocks for pooling
Returns:
block_importance: [batch, heads, num_blocks] importance scores per block
"""
batch, heads, seq_len, _ = attention_scores.shape
num_blocks = (seq_len + block_size - 1) // block_size
# Pad if necessary
padded_len = num_blocks * block_size
if padded_len != seq_len:
pad_amount = padded_len - seq_len
attention_scores = F.pad(attention_scores, (0, pad_amount, 0, pad_amount))
# Reshape into blocks and compute mean importance
blocked = attention_scores.view(batch, heads, num_blocks, block_size, num_blocks, block_size)
# Take mean across block dimensions to get per-block importance
block_importance = blocked.mean(dim=(3, 5)) # [batch, heads, num_blocks, num_blocks]return block_importance
Step 2: Head-Specific Sparse Pattern Selection
After identifying block-level importance, select top-ρ% blocks for each head to create the sparse mask.
defselect_sparse_pattern(block_importance, selection_ratio=0.3):
"""
Select top-ρ% blocks per head to form sparse attention mask.
Args:
block_importance: [batch, heads, num_blocks, num_blocks]
selection_ratio: fraction of blocks to retain (ρ)
Returns:
sparse_mask: [batch, heads, num_blocks, num_blocks] boolean mask
"""
batch, heads, num_blocks_q, num_blocks_k = block_importance.shape
# Reshape to treat each head independently
importance_flat = block_importance.view(batch * heads, num_blocks_q, num_blocks_k)
# Find threshold for top-ρ% blocks
num_keep = max(1, int(num_blocks_k * selection_ratio))
# For each query block, select top-k key blocks
sparse_masks = []
for b_h inrange(batch * heads):
mask = torch.zeros(num_blocks_q, num_blocks_k, dtype=torch.bool, device=importance_flat.device)
for q_idx inrange(num_blocks_q):
scores = importance_flat[b_h, q_idx, :]
topk_indices = torch.topk(scores, num_keep)[1]
mask[q_idx, topk_indices] = True
sparse_masks.append(mask)
sparse_mask = torch.stack(sparse_masks).view(batch, heads, num_blocks_q, num_blocks_k)
return sparse_mask
Step 3: Expand Block Mask to Token-Level Granularity
Convert the block-level sparse mask back to token-level positions for attention computation.
defexpand_block_mask_to_tokens(sparse_block_mask, block_size=128, seq_len=None):
"""
Expand block-level sparse mask to token-level attention mask.
Args:
sparse_block_mask: [batch, heads, num_blocks_q, num_blocks_k] boolean
block_size: original block size used
seq_len: actual sequence length (before padding)
Returns:
attention_mask: [batch, heads, seq_len, seq_len] boolean
"""
batch, heads, num_blocks_q, num_blocks_k = sparse_block_mask.shape
# Expand blocks to tokens
attention_mask = torch.zeros(
batch, heads,
num_blocks_q * block_size,
num_blocks_k * block_size,
dtype=torch.bool,
device=sparse_block_mask.device
)
for b inrange(batch):
for h inrange(heads):
for q_block inrange(num_blocks_q):
for k_block inrange(num_blocks_k):
if sparse_block_mask[b, h, q_block, k_block]:
q_start = q_block * block_size
q_end = q_start + block_size
k_start = k_block * block_size
k_end = k_start + block_size
attention_mask[b, h, q_start:q_end, k_start:k_end] = True# Trim to actual sequence length if neededif seq_len isnotNone:
attention_mask = attention_mask[:, :, :seq_len, :seq_len]
return attention_mask
Step 4: Pattern Reuse Across Denoising Steps
During early denoising steps, compute and cache sparse patterns. Reuse them for all subsequent steps.
classSparseAttentionCache:
"""Cache and manage sparse attention patterns across denoising steps."""def__init__(self, skip_ratio=0.2, selection_ratio=0.3, block_size=128):
self.skip_ratio = skip_ratio # Fraction of steps to use full attentionself.selection_ratio = selection_ratio # Sparsity level ρself.block_size = block_size
self.sparse_masks = {} # Cache: (batch_id, head_id) -> maskself.pattern_frozen = Falsedefshould_use_full_attention(self, step, total_steps):
"""Determine if current step should use full attention."""
skip_threshold = int(total_steps * self.skip_ratio)
return step < skip_threshold
defget_or_compute_mask(self, attention_scores, step, total_steps, seq_len):
"""
Get cached sparse mask or compute new one if in early steps.
Args:
attention_scores: [batch, heads, seq_len, seq_len] full attention
step: current denoising step (0-indexed)
total_steps: total number of denoising steps
seq_len: sequence length
Returns:
mask: sparse attention mask, or None if full attention should be used
"""ifself.should_use_full_attention(step, total_steps):
# Compute patterns during early steps
block_importance = compute_block_importance(attention_scores, self.block_size)
sparse_block_mask = select_sparse_pattern(block_importance, self.selection_ratio)
sparse_mask = expand_block_mask_to_tokens(sparse_block_mask, self.block_size, seq_len)
self.sparse_masks[step] = sparse_mask
returnNone# Use full attentionelse:
# Reuse cached pattern from first post-skip stepifnotself.pattern_frozen:
skip_threshold = int(total_steps * self.skip_ratio)
if skip_threshold inself.sparse_masks:
self.pattern_frozen = Trueifself.pattern_frozen and skip_threshold inself.sparse_masks:
returnself.sparse_masks[skip_threshold]
returnNone
Step 5: Integration with Forward Pass
Apply sparse attention during the actual attention computation using PyTorch's FlexAttention.
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
defapply_sparse_attention(
query, key, value,
sparse_mask,
block_size=128):
"""
Apply sparse attention using FlexAttention.
Args:
query: [batch, heads, seq_len, dim]
key: [batch, heads, seq_len, dim]
value: [batch, heads, seq_len, dim]
sparse_mask: [batch, heads, seq_len, seq_len] boolean
block_size: block size for FlexAttention
Returns:
output: [batch, heads, seq_len, dim] attention output
"""
batch, heads, seq_len, dim = query.shape
# Convert boolean mask to FlexAttention block mask format
block_mask = create_block_mask(sparse_mask, batch, heads, seq_len, seq_len)
# Compute sparse attention
output = flex_attention(query, key, value, block_mask=block_mask)
return output
defdiffusion_forward_with_sparse_attention(
model, input_ids, noise_level, step, total_steps,
sparse_cache
):
"""
Forward pass for diffusion model with optional sparse attention.
Args:
model: diffusion language model
input_ids: token indices
noise_level: diffusion noise level
step: current denoising step
total_steps: total denoising iterations
sparse_cache: SparseAttentionCache instance
Returns:
logits: output predictions
"""# Get embeddings
hidden = model.embed(input_ids)
# Add noise level information
hidden = model.add_noise_embedding(hidden, noise_level)
seq_len = hidden.shape[1]
# Apply transformer blocks with sparse/full attentionfor layer in model.transformer_blocks:
# Standard pre-norm setup
norm_hidden = layer.norm1(hidden)
# Attention computation
q, k, v = layer.attention.qkv(norm_hidden)
# Check if we should use sparse attention
use_full_attention = sparse_cache.should_use_full_attention(step, total_steps)
if use_full_attention:
# Compute full attention and extract/cache patterns
scores = torch.matmul(q, k.transpose(-2, -1)) / (q.shape[-1] ** 0.5)
sparse_mask = sparse_cache.get_or_compute_mask(scores, step, total_steps, seq_len)
attn_output = torch.matmul(
torch.softmax(scores, dim=-1), v
)
else:
# Use cached sparse pattern
sparse_mask = sparse_cache.get_or_compute_mask(None, step, total_steps, seq_len)
if sparse_mask isnotNone:
attn_output = apply_sparse_attention(q, k, v, sparse_mask)
else:
# Fallback to full attention
scores = torch.matmul(q, k.transpose(-2, -1)) / (q.shape[-1] ** 0.5)
attn_output = torch.matmul(torch.softmax(scores, dim=-1), v)
# Projection and residual
hidden = hidden + layer.attention.out_proj(attn_output)
# FFN
hidden = hidden + layer.ffn(layer.norm2(hidden))
logits = model.output_projection(hidden)
return logits
Models with stable attention patterns: Dream, LLaDA, and similar diffusion models show consistent head-specific patterns.
Tasks tolerating small quality trade-offs: Some domains (summarization, paraphrasing) show negligible degradation.
When NOT to Use SparseD
Very short sequences (<2k tokens): Sparse attention overhead exceeds quadratic attention cost.
Quality-critical tasks requiring perfect preservation: Machine translation, instruction-following, or code generation where even 0.1% accuracy loss is unacceptable.
Early-stage denoising dominance: If your model heavily depends on the first 5% of denoising steps, applying sparsity too late may still hurt.
Heterogeneous attention patterns: Models where different samples exhibit completely different attention structures won't benefit from frozen masks.
Single-pass inference: SparseD amortizes cost across steps; single-pass inference (e.g., length-1 generation) sees no benefit.
Common Pitfalls and Mitigations
Setting skip too low: Applying sparse attention in early critical steps causes 3-5% accuracy degradation. If you see sharp quality loss, increase skip to 0.25-0.30.
Frozen patterns on distribution shift: If your model receives inputs very different from training distribution, cached patterns may become stale. Periodically recompute if input style changes significantly.
Block size misalignment: Using block_size that doesn't divide sequence length evenly causes padding overhead. Ensure block_size divides your typical seq_len evenly.
FlexAttention compilation overhead: First inference includes compilation cost (~100-200ms). Always run a warm-up pass to amortize this cost.
Memory spikes during pattern computation: Early steps still compute full attention. Ensure GPU memory accommodates batch_size × heads × seq_len² × 4 bytes temporarily.
Attention head heterogeneity assumptions: SparseD assumes patterns stabilize; if heads are highly noisy or unstable, quality may degrade. Validate on a small sample first.
