Sink-Aware Pruning for Diffusion Language Models

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Myrzakhan, Aidar, Li, Tianyi, Guo, Bowei, Tang, Shengkun, Shen, Zhiqiang
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866915808139018240
author Myrzakhan, Aidar
Li, Tianyi
Guo, Bowei
Tang, Shengkun
Shen, Zhiqiang
author_facet Myrzakhan, Aidar
Li, Tianyi
Guo, Bowei
Tang, Shengkun
Shen, Zhiqiang
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2602_17664
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Sink-Aware Pruning for Diffusion Language Models
Myrzakhan, Aidar
Li, Tianyi
Guo, Bowei
Tang, Shengkun
Shen, Zhiqiang
Computation and Language
Artificial Intelligence
Machine Learning
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.
title Sink-Aware Pruning for Diffusion Language Models
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2602.17664