Source-linked AI summary
Sink-Aware Pruning for Diffusion Language Models
Aidar Myrzakhan, Tianyi Li, Bowei Guo, Shengkun Tang, Zhiqiang Shen
TL;DR
DLMs require efficient pruning because iterative denoising repeatedly updates the sequence, while existing AR-derived heuristics assume attention sinks should be preserved. The paper measures higher sink-position variance in DLMs, proposes Sink-Aware Pruning to downweight and prune unstable sinks, and reports a better quality-efficiency trade-off than prior baselines under matched compute.
Problem
Existing pruning recipes are adapted from AR Transformers and assume their attention behaviors and important tokens transfer unchanged to diffusion generation.
Method
Sink-Aware Pruning estimates sink variance across denoising timesteps, identifies unstable sinks, and uses downweighted activations in pruning criteria to guide their removal.
Results
Sink-Aware Pruning outperforms strong prior pruning baselines under the same compute and improves the quality-efficiency trade-off without retraining.
Takeaways & Limitations
Attention sinks are not universally must-keep tokens; pruning rules should align with generation dynamics rather than transfer AR heuristics directly.
Takeaways & Limitations
Sink statistics rely on a fixed calibration distribution, and broader validation on larger multimodal and long-context settings remains needed.
Abstract
from arXiv · showhide
Diffusion Language Models (DLMs) incur high inference cost due to iterative denoising, motivating efficient pruning. Existing pruning heuristics largely inherited from autoregressive (AR) LLMs, typically preserve attention sink tokens because AR sinks serve as stable global anchors. We show that this assumption does not hold for DLMs: the attention-sink position exhibits substantially higher variance over the full generation trajectory (measured by how the dominant sink locations shift across timesteps), indicating that sinks are often transient and less structurally essential than in AR models. Based on this observation, we propose ${\bf \texttt{Sink-Aware Pruning}}$, which automatically identifies and prunes unstable sinks in DLMs (prior studies usually keep sinks for AR LLMs). Without retraining, our method achieves a better quality-efficiency trade-off and outperforms strong prior pruning baselines under matched compute. Our code is available at https://github.com/VILA-Lab/Sink-Aware-Pruning.
1 Introduction
DLMs update the full sequence through iterative denoising, but pruning heuristics often transfer AR assumptions that attention sinks should always be preserved. Sink-Aware Pruning measures sink-position variance and prunes unstable DLM sinks, improving pruning efficiency under matched compute.
- Motivation: DLMs repeatedly update the full token sequence across denoising timesteps, increasing inference compute and memory costs.This motivates efficient pruning for practical deployment.
- Motivation: AR-derived pruning recipes assume that attention behaviors and important tokens transfer unchanged to diffusion-style generation.AR methods commonly preserve sinks because they behave as stable global anchors.
- Sink dynamics: Sink-position variance is substantially higher in DLMs than AR LLMs because dominant sink locations shift across diffusion timesteps.The analysis tracks sink positions over the full generation process, using each AR token-generation step or DLM diffusion timestep.
- Sink dynamics: DLM sinks are often transient: they attract attention at some timesteps, then lose prominence as different sinks take over during denoising.This contrasts with the persistent sink positions typical of AR generation.
- Method: Sink-Aware Pruning estimates sink variance, identifies unstable sink candidates, and prunes them while preserving informative attention pathways.The method is generation-paradigm aware: it encourages sink pruning in high-variance DLMs rather than applying AR sink preservation.
- Results: Sink-Aware Pruning improves the quality-efficiency trade-off over strong prior pruning baselines under matched compute.The paper presents this as a retraining-free route to accelerating DLM inference.
2 Related Work
Related work establishes DLMs as iterative or parallel alternatives to AR generation, while attention sinks and post-training pruning motivate adapting pruning criteria to diffusion-specific sink behavior. Prior masked-DLM studies report step-dependent sinks and relative robustness to masking top sinks.
- Diffusion language models: DLMs include continuous and discrete formulations that diffuse over embeddings, logits, or categorical token spaces.Recent discrete DLMs can approach strong AR baselines while supporting parallel decoding.
- Attention sinks: Attention sinks are small sets of often-early tokens that absorb disproportionate attention across many layers and heads despite limited semantic value.Prior work attributes this behavior to Softmax normalization when queries lack strong matches.
- Attention sinks: Masked DLM studies find sink locations can emerge, shift, or vanish during denoising, unlike the more stable sinks of autoregressive models.This motivates treating sink behavior as dependent on generation dynamics.
- Attention sinks: Masked DLMs are relatively robust to masking top sinks, with prior studies reporting only minor performance drops.This differs from autoregressive settings, where removing sinks can severely hurt performance.
- Pruning: Post-training pruning removes unimportant weights from pretrained models using lightweight importance metrics without full retraining.The paper’s sink-aware pipeline modifies activations before applying existing Wanda or SparseGPT criteria.
3 Sink-Aware Pruning for DLMs
This section contrasts sink behavior in diffusion and autoregressive models, then uses those differences to motivate sink-aware importance estimation for pruning.
- Pruning baselines: Wanda combines weight magnitude with the norm of corresponding input activations, while SparseGPT prunes through layer-wise reconstruction using second-order information.Both are post-training pruning criteria that operate without full retraining; Wanda avoids gradient or Hessian computation, whereas SparseGPT uses an empirical Hessian approximation.
- Sink variance: Temporal variance measures sink-location shifts across steps, with near-zero values indicating stationary sinks and larger values indicating migration.The analysis tracks attention over generation steps, defined as diffusion timesteps for DLMs and token emissions for AR models.
- Sink variance: AR models have high spatial but near-zero temporal variance, whereas DLMs have lower spatial variance and temporal variance orders of magnitude larger.AR sinks remain concentrated on fixed early positions, while DLM sink locations change across the denoising trajectory.
- Sink variance: DLM sink trajectories drift from earlier to later positions as denoising advances, with wide step-to-step variability.The reported pattern links early denoising with global structure and later denoising with local token-level refinement.
- Sink-aware importance reweighting: The method uses attention mass aggregated across layers and heads to identify sink positions whose attention substantially exceeds that of other positions.A token is treated as a sink when its aggregated attention mass is unusually high; the resulting scores support sink-aware pruning decisions.
- Sink-aware importance reweighting: Sink-aware pruning suppresses sink-token contributions in activations before applying existing pruning criteria such as Wanda or SparseGPT.The masked activation scales each token-position row, and sink-masked inputs are then used for importance scoring or Hessian construction.
4 Experiments
Experiments evaluate Sink-Aware Pruning across multiple diffusion language models, benchmarks, pruning methods, sparsity levels, and both unstructured and structured settings. The method generally improves over baselines, especially under aggressive compression, though gains are not uniform.
- Experimental Setup: Evaluations cover LLaDA, Dream, and LLaDA1.5 across eight benchmarks, using Wanda, SparseGPT, and magnitude-based pruning baselines.Experiments include unstructured sparsity ratios of 25%, 50%, and 75%, plus structured pruning settings.
- Unstructured Pruning: Sink-aware pruning consistently matches or surpasses corresponding baselines across unstructured pruning levels.The reported comparison spans the evaluated models and benchmarks under identical evaluation protocols.
- Unstructured Pruning: 50%–75% sparsity produces the most pronounced gains, while advantages at 25% sparsity are narrower and occasionally marginal.At 75% sparsity, all methods degrade substantially, but sink-aware pruning remains among the top-performing approaches.
- Structured Pruning: Sink-aware pruning consistently outperforms the structured baseline, with the margin growing as pruning ratios increase from 0.3 to 0.5.Structured pruning removes entire heads or layers, making pruning decisions more disruptive than individual-weight removal.
- Visualization and Analysis: Figure 6 shows broadly positive improvements over Wanda and SparseGPT across models and sparsity levels, with larger gains under more aggressive compression.The most notable improvements appear at 75% sparsity, particularly for LLaDA-1.5 under Wanda.
- Visualization and Analysis: At 50% unstructured sparsity, sink-aware pruning shows partial alignment with sink strength for Wanda, weaker correspondence for SparseGPT, and smaller head-level sparsity changes for SparseGPT.Figures 7 and 8 compare head-level mask disagreement and signed sparsity differences between the methods.
- Discussion: Improvements are not uniform across all settings, with sink-aware variants sometimes close to or slightly below strong baselines in low-sparsity regimes.The authors attribute this variation to differences in sink instability across layers, tasks, and model families.
5 Conclusion
The paper argues that pruning rules should account for generation dynamics rather than transfer attention-sink heuristics directly from autoregressive models. Sink-Aware Pruning discounts unstable DLM sinks and improves or matches strong baselines without retraining.
- Conclusion: Sink positions in DLMs are substantially less stable than in autoregressive models, challenging the common sink-preservation heuristic.Across multiple DLM families, benchmarks, and sparsity levels, the method improves or matches strong baselines while producing a better quality-efficiency trade-off.
- Conclusion: Sink-Aware Pruning discounts unstable sinks when estimating pruning importance and aligns pruning rules with diffusion generation dynamics.The conclusion presents this alignment as the paper’s central practical implication.
Limitations
The study identifies calibration distribution, post-training evaluation, and validation scope as limitations that may constrain robustness and generality.
- Limitations: Sink statistics rely on a fixed calibration distribution, the study mainly evaluates pruning without recovery finetuning, and broader multimodal and long-context validation remains needed.The authors suggest post-pruning adaptation, larger multimodal and long-context evaluations, timestep-adaptive policies, and joint quantization as future directions.
A Results on Additional Models
Additional results include MMaDA, a multimodal unified diffusion language model, extending evaluation beyond the main experiments on LLaDA, Dream, and LLaDA-1.5.
- Additional Models: Table 5 reports pruning results on MMaDA, complementing the main-paper results on LLaDA, Dream, and LLaDA-1.5.MMaDA provides an additional multimodal DLM evaluation setting.
B Evaluation Details
The evaluation uses official model implementations and Fast-DLLM within lm-evaluation-harness, with identical configurations across models and baselines. Generation and loglikelihood tasks use specified lengths, denoising steps, and Monte Carlo sampling protocols.
- Official LLaDA and Dream implementations support loglikelihood evaluations, while Fast-DLLM enables fast inference for GSM8K generation.
- All models and baselines are evaluated under identical configurations using lm-evaluation-harness.
- Generation tasks use 256-token generation and block lengths with 256 denoising steps.
- Loglikelihood benchmarks use 128 Monte Carlo samples, except MMLU, which uses one sample following official LLaDA conventions.
C Model Architecture Details
The evaluated models use decoder-only Transformer backbones as denoising networks for discrete diffusion. LLaDA, Dream-7B, and MMaDA-8B differ in their base configurations and initialization sources.
- All models use decoder-only Transformer backbones as denoising networks in discrete diffusion.
- LLaDA uses a LLaMA-style architecture with RMSNorm, SwiGLU blocks, and RoPE positional encoding.
- Dream-7B adopts the base configuration of Qwen2.5-7B.
- MMaDA-8B follows the LLaDA configuration for text generation and is initialized from an LLaDA-8B checkpoint.
D Benchmark Descriptions
The study evaluates text generation and reasoning across established benchmarks covering factual knowledge, science, commonsense, reading comprehension, expert knowledge, and mathematical reasoning. The listed suite includes MMLU, ARC-C, PIQA, WinoGrande, HellaSwag, RACE, GSM8K, and GPQA.
- The benchmark suite covers text generation and reasoning across diverse established tasks.
- MMLU measures factual knowledge and reasoning across 57 academic subjects using multiple-choice questions.
- ARC-C tests grade-school science reasoning beyond surface pattern matching, while GPQA uses expert-written graduate-level questions in biology, physics, and chemistry.
- PIQA, WinoGrande, and HellaSwag evaluate physical commonsense, coreference resolution, and commonsense continuation inference, respectively.
- RACE evaluates multi-sentence reading comprehension from English examinations, and GSM8K tests multistep mathematical reasoning with exact-match final-answer evaluation.