Sink-Aware Pruning for Diffusion Language Models

Aidar Myrzakhan; Tianyi Li; Bowei Guo; Shengkun Tang; Zhiqiang Shen

cs.CLFeb 19, 2026

Sink-Aware Pruning for Diffusion Language Models

Authors: Aidar Myrzakhan, Tianyi Li, Bowei Guo, Shengkun Tang, Zhiqiang Shen

Organizations: VILA Lab, MBZUAI

Abstract

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.

Figures & tables

Figure 1: Illustration of attention sink behaviors in Diffusion Language Models. Ours Sink-Aware Pruning reduces sink variance by downscaling unstable sinks.

Figure 2: Attention sink heatmap dynamics across generation steps for AR LLM (LLaMA-3-8B) and DLM (LLaDA). For each model, we show 3 different generation stages (25, 50, and 75% of the total generation steps) and plot the attention mass received by each token position (y-axis) across all heads/layers (x-axis). In LLaMA, the sink position (deep-blue vertical band) is stable across steps, while in LLaDA, the sink position shifts significantly across diffusion steps, indicating higher sink variance. The step in AR model refers to the generation process.

Figure 3: Overview of Sink-Aware Pruning . Given input activations, we compute per-token attention mass aggregated across all layers and heads (Step 1), identify sink tokens via a threshold-based criterion, and derive a soft down-weighting factor $\omega=1-s$ . The original activation $X$ is then suppressed at sink positions to produce a new activation $\tilde{X}$ (Step 2), which is substituted into existing pruning criteria, Wanda or SparseGPT, to compute sink-aware importance scores (Step 3). Final pruning decisions are made based on the updated scores (Step 4).

Panel of Figure 4 — Figure 4: Attention sink variance for diffusion LMs (LLaDA, Dream) and autoregressive LMs (Llama 3.1, Qwen 2.5).

Figure 5: Sink position across generation/denoising steps for diffusion and AR LMs. Shaded regions denote

\pm

std across runs.

Sparsity	Method	Avg.	MMLU	ARC-C	PIQA	WinoG	GSM8K	HellaSwag	GPQA	RACE
Dense	Base	57.93	65.97	43.00	74.10	69.30	69.29	72.70	30.40	38.70
0.25	Wanda	57.43	65.20	43.94	75.30	68.59	66.03	71.95	29.85	38.55
	Sink-Aware	57.62	65.41	43.52	74.97	68.59	68.16	72.30	29.70	38.32
	SparseGPT	57.23	65.16	43.09	74.43	67.56	67.17	72.10	30.25	38.10
	Sink-Aware	57.68	65.33	43.09	74.37	69.53	68.58	71.98	30.12	38.40
0.50	Wanda	52.70	61.43	39.08	72.63	64.56	57.01	67.52	27.15	32.20

Table 1: Pruning results on LLaDA across 8 benchmarks.

Sparsity	Method	Avg.	MMLU	ARC-C	PIQA	WinoG	GSM8K	HellaSwag	GPQA	RACE
Dense	Base	60.94	68.25	55.20	74.27	67.72	67.50	73.30	36.60	44.70
0.50	Wanda	51.58	60.90	41.64	68.12	60.38	49.13	65.40	30.85	36.20
	Sink-Aware	51.68	60.84	42.83	68.82	59.75	49.66	65.15	30.52	35.85
	SparseGPT	54.40	64.13	44.88	70.18	63.54	52.84	68.72	32.48	38.42
	Sink-Aware	54.58	64.46	44.97	70.18	64.09	53.15	68.35	32.90	38.58
0.75	Wanda	30.38	23.96	20.65	53.43	52.17	1.06	45.12	23.50	23.15

Table 2: Pruning results on Dream across 8 benchmarks.

Sparsity	Method	Avg.	MMLU	ARC-C	PIQA	WinoG	GSM8K	HellaSwag	GPQA	RACE
Dense	Base	58.59	64.07	50.43	73.61	68.90	65.73	74.43	31.75	39.87
0.25	Wanda	58.09	63.30	51.37	74.81	68.19	62.47	73.68	31.20	39.72
	Sink-Aware	58.22	63.41	51.54	74.97	68.27	62.70	73.55	31.30	40.05
	SparseGPT	58.44	63.62	52.05	74.97	68.51	63.08	73.92	31.55	39.85
	Sink-Aware	58.47	63.59	52.05	74.97	68.68	63.23	73.78	31.62	39.80
0.50	Wanda	53.78	59.04	47.27	72.69	66.85	51.63	69.55	28.35	34.82

Table 3: Pruning results on LLaDA1.5 across 8 benchmarks.

Structured Pruning 0.3
Model	PIQA	WinoG	ARC-E	ARC-C
LLaDA (Base)	0.7942	0.7388	0.7420	0.4437
LLaDA-structure	0.6834	0.6630	0.6907	0.3780
Sink-Aware (Ours)	0.6955	0.6740	0.7175	0.3820
Structured Pruning 0.5
LLaDA-structure	0.5898	0.5572	0.4853	0.2039

Table 4: Structured pruning results on LLaDA.

Figure 6: Delta vs. baseline on average accuracy. Average accuracy change (

\Delta

Avg, in percentage points) from applying Sink-Aware pruning on top of each baseline ( Wanda / SparseGPT ) at different sparsity levels. Bars report

\Delta=\text{Avg}(\texttt{Sink-Aware})-\text{Avg}(\text{baseline})

; positive values indicate improved performance retention after pruning. Results are shown for LLaDA, Dream, and LLaDA1.5 over 8 benchmarks.

Figure 7: Sink strength and per-head mask disagreement (XOR) between baseline and sink-aware pruning. The left subfigure shows the sink strength of the original LLaDA-8B model, measured as the average attention each head allocates to sink tokens. The center and right show the fraction of weights per head whose pruning decision differs (XOR) between the baseline and our method, using Wanda (center) and SparseGPT (right), respectively.

Figure 8: Per-head sparsity difference between baseline and sink-aware pruning. The left and right subfigures show the signed per-head sparsity difference for Wanda and SparseGPT, respectively. Red indicates that the sink-aware variant prunes more aggressively in a given head, while blue indicates it preserves more weights.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Sparsity	Method	Avg.	MMLU	ARC-C	PIQA	WinoG	GSM8K	HellaSwag	GPQA	RACE
Dense	Base	35.63	33.69	24.32	60.50	50.83	9.48	52.50	25.20	28.50
0.25	Wanda	35.20	32.89	23.98	60.28	50.59	8.72	52.25	24.80	28.12
	Sink-Aware	35.03	32.99	24.15	60.17	49.64	8.87	51.80	24.72	27.90
	SparseGPT	35.48	33.30	24.57	60.39	50.91	9.10	52.05	25.15	28.35
	Sink-Aware	35.45	33.30	24.83	60.34	50.67	9.17	52.10	25.00	28.20
0.50	Wanda	33.91	31.14	23.12	59.79	51.14	5.38	49.72	24.18	26.80

Appendix

Table 5: Pruning results on MMaDA across 8 benchmarks.

Model	Params	#Layers	$d_{\text{model}}$
LLaDA-8B	8.02B	32	4096
LLaDA-1.5-8B	8.02B	32	4096
Dream-7B	7B	28	3584
MMaDA-8B	8.02B	32	4096

Appendix

Table 6: Backbone configurations.

LLM Pruning

Large language model pruning removes less useful weights or components to reduce memory and inference costs while preserving model capabilities.

See all 166 papers

Learnability-Informed Fine-Tuning of Diffusion Language Models
May 21, 2026Shubham Parashar, Atharv Chagi, Jacob Helwig +5Supervised Fine-Tuning Diffusion Language Models
Department of Computer Science and Engineering, Texas A&M University, College Station, TX, USA · Department of Electrical and Computer Engineering, Texas A&M University, College Station, TX, USA
ALoDLM: Adaptively Looped Diffusion Language Models
Oct 3, 2026Liancheng Fang, Zhuowei Li, Youngeun Kim +10Adaptive Computation for Diffusion Models Diffusion Language Models
University of Illinois Chicago · Amazon AGI · Korea University
Window-Diffusion: Accelerating Diffusion Language Model Inference with Windowed Token Pruning and Caching
Jan 28, 2026Fengrui Zuo, Zhiwei Ke, Yiming Liu +3Diffusion Language Model Inference Token Pruning
School of Computer Science and Technology, University of Science and Technology of China · Suzhou Institute of Advanced Research, University of Science and Technology of China

May 21, 2026J/K move · Enter open · S save

Learnability-Informed Fine-Tuning of Diffusion Language Models

Shubham Parashar, Atharv Chagi, Jacob Helwig, Lakshmi Jotsna+4

Department of Computer Science and Engineering, Texas A&M University, College Station, TX, USA · Department of Electrical and Computer Engineering, Texas A&M University, College Station, TX, USA

Supervised Fine-Tuning Diffusion Language Models Diffusion Language Model Training

Open paper AskPDF arXiv

We aim to improve the reasoning capabilities of diffusion language models (DLMs). While SFT is a popular post-training recipe for autoregressive models, its use in DLMs faces challenges and can even hurt performance, though the underlying causes remain understudied. Our analysis reveals that vanilla SFT overlooks learnability, namely what and when tokens are learned. Specifically, rare tokens are difficult to learn when most of the input is masked, whereas it is straightforward and thus of little value to learn common tokens when most of the input is unmasked. Motivated by our analysis, we propose LIFT, an efficient SFT-based post-training algorithm for DLMs. LIFT learns easy tokens when most of the input is masked and hard tokens when more context is available, thus aligning the training with the information available at different diffusion time steps. Our results show that LIFT outperforms existing SFT baselines across six reasoning benchmarks, achieving up to a 3x relative gain on AIME'24 and AIME'25. Our code is publicly available at https://github.com/divelab/LIFT.

Abstract

LLM Pruning

Explore similar work

Learnability-Informed Fine-Tuning of Diffusion Language Models

ALoDLM: Adaptively Looped Diffusion Language Models

Window-Diffusion: Accelerating Diffusion Language Model Inference with Windowed Token Pruning and Caching