Sink-Aware Pruning for Diffusion Language Models

TL;DR

This paper introduces Sink-Aware Pruning, a novel method for reducing the inference cost of Diffusion Language Models by identifying and pruning unstable attention sinks, leading to better efficiency without retraining.

Contribution

It reveals that attention sinks in DLMs are often transient, unlike in AR models, and proposes a pruning method that leverages this insight for improved performance.

Findings

Outperforms prior pruning baselines in quality-efficiency trade-off

Pruning unstable sinks improves inference efficiency without retraining

Attention sink variance is higher in DLMs than in AR models

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.