Sink-Aware Pruning for Diffusion Language Models

Paper: 2602.17664 Authors: Aidar Myrzakhan, Tianyi Li, Bowei Guo, Shengkun Tang, Zhiqiang Shen Categories: cs.CL, cs.AI, cs.LG

Abstract

Diffusion Language Models (DLMs) represent a promising alternative to autoregressive architectures, but their iterative denoising process creates significant computational overhead during inference. This paper challenges a fundamental assumption in neural network pruning: that attention sink tokens—positions that accumulate disproportionate attention weights—should always be preserved. The authors demonstrate that unlike autoregressive models where sinks serve as stable global anchors, DLMs exhibit highly transient sink behavior with substantial positional variance across generation timesteps. Leveraging this insight, they introduce Sink-Aware Pruning, a training-free method that automatically identifies and prunes unstable sinks in DLMs. Their approach achieves superior quality-efficiency trade-offs compared to existing pruning baselines under matched computational budgets.

Key Contributions

• Empirical analysis revealing fundamental differences between attention sink behavior in DLMs versus autoregressive LLMs, showing that DLM sinks exhibit high temporal variance and are structurally less essential
• Sink-Aware Pruning algorithm that automatically identifies transient sinks and prunes them without requiring model retraining
• Comprehensive experimental validation demonstrating improved quality-efficiency trade-offs over strong baselines including magnitude-based and Wanda pruning methods
• Open-source implementation enabling reproducibility and practical deployment of the proposed technique

Methodology: Rethinking Attention Sinks in Diffusion Models

The core insight driving this work stems from analyzing attention patterns across the diffusion generation trajectory. In autoregressive models, attention sinks typically concentrate at specific positions (often initial tokens like BOS) and remain stable throughout generation. These sinks act as “attention dumps” where the model can allocate excess attention mass while maintaining normalized distributions.

For DLMs, the authors measure sink stability by tracking the dominant attention-receiving positions across diffusion timesteps $t \in [0, T]$. They quantify this through positional variance metrics, revealing that:

$\text{Var}(\text{sink\_position}) \gg \text{Var}_{\text{AR}}(\text{sink\_position})$

This high variance indicates that DLM sinks are ephemeral—they shift locations as the denoising process progresses. The transient nature suggests these positions are less structurally critical than their autoregressive counterparts.

The Sink-Aware Pruning algorithm operates in two phases:

1. Sink Detection: Analyze attention weight distributions across representative samples and timesteps to identify positions with high attention mass but high temporal variance
2. Selective Pruning: Apply structured or unstructured pruning that explicitly targets unstable sinks while preserving stable architectural components

Critically, this approach inverts conventional pruning wisdom—rather than protecting sinks, it treats them as pruning candidates when they exhibit instability.

Experimental Results and Analysis

The authors evaluate their method on standard language modeling benchmarks, comparing against several baselines:

• Magnitude pruning: Removes weights with smallest absolute values
• Wanda: A recent structured pruning method designed for LLMs
• Sink-preserving variants: Adaptations of baseline methods that explicitly protect detected sinks

Key findings include:

Performance under matched FLOPs: At equivalent computational budgets (e.g., 50% sparsity), Sink-Aware Pruning consistently outperforms baselines by 2-5% in perplexity across multiple model scales. This gap widens at higher sparsity levels where sink instability becomes more pronounced.

Training-free efficiency: Unlike many pruning methods requiring fine-tuning or retraining, this approach achieves competitive results through analysis and pruning alone. This dramatically reduces the deployment barrier for practitioners.

Generalization across architectures: The method shows consistent improvements across different DLM variants, suggesting the sink instability phenomenon is a general property of diffusion-based generation rather than architecture-specific.

Ablation studies confirm that explicitly targeting unstable sinks provides the performance gain—random pruning or sink-agnostic methods fail to achieve similar results even when removing equivalent parameter counts.

Implications for Efficient Language Models

This work carries several important implications for the broader landscape of efficient language modeling:

Architectural specificity matters: Pruning heuristics cannot be blindly transferred between model families. The diffusion generation process creates fundamentally different attention dynamics that require tailored optimization strategies.

Temporal analysis is crucial: Static analysis of trained models may miss critical dynamic properties. Understanding how model components behave across the full generation trajectory reveals optimization opportunities invisible to single-snapshot analysis.

Training-free methods remain viable: Despite the trend toward expensive pruning-and-retraining pipelines, carefully designed analysis-based methods can achieve strong results. This is particularly valuable for practitioners with limited computational resources.

Future research directions: The sink instability phenomenon raises questions about whether DLM architectures could be redesigned to eliminate transient sinks entirely, potentially improving both efficiency and generation quality simultaneously.

Takeaways

1. Attention sinks in Diffusion Language Models exhibit fundamentally different behavior than in autoregressive models—they are transient rather than stable, making them viable pruning targets
2. Sink-Aware Pruning achieves superior quality-efficiency trade-offs without requiring model retraining, outperforming magnitude-based and structured pruning baselines
3. Pruning strategies must be tailored to specific model architectures and generation processes; techniques effective for autoregressive models may be suboptimal or counterproductive for diffusion-based approaches
4. Temporal analysis across the full generation trajectory reveals optimization opportunities missed by static model analysis
5. The method’s training-free nature makes it immediately deployable for practitioners seeking to reduce DLM inference costs