💬 NLP

Sink-Aware Pruning for Diffusion Language Models

This paper introduces Sink-Aware Pruning, a method that improves the efficiency of Diffusion Language Models by identifying and removing transient attention sink tokens—a strategy distinct from autoregressive models where sinks are typically preserved—thereby achieving superior quality-efficiency trade-offs without retraining.

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

Published 2026-02-20

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Two Different Ways of Writing

Imagine you have two different types of writers trying to write a story:

The Autobiographer (Autoregressive Models): This writer writes one word at a time, left to right. Once they write a word, they never change it. They rely heavily on the very first few words (the "sink") to set the tone and keep the story on track. If you remove those first words, the whole story falls apart.
The Sculptor (Diffusion Language Models): This writer starts with a block of noise (static) and slowly chisels away the bad parts to reveal the final statue. They look at the entire picture at every single step, refining it from a rough draft to a masterpiece.

The Problem: The "Don't Touch the Foundation" Rule

In the world of AI, researchers have a golden rule for compressing models (making them smaller and faster): "Never cut the 'Attention Sinks'."

What is an Attention Sink? Think of it as a "magnet" in the model. In the Autobiographer (AR) models, there are specific tokens (usually the very first ones) that act like a heavy anchor. They hold the whole structure together. If you cut them out, the model collapses. So, standard pruning tools are programmed to protect these anchors at all costs.

The Mistake: Researchers tried to use this same "Protect the Anchor" rule on the Sculptor (Diffusion) models. They assumed the Sculptor also needed a permanent, unchanging anchor to hold the story together.

The Discovery: The Anchor is Moving!

The authors of this paper did a deep dive and found something surprising: The Sculptor doesn't have a permanent anchor.

In the Autobiographer: The "magnet" stays in the same spot (the beginning) the whole time. It's a stable anchor.
In the Sculptor: The "magnet" is transient. It jumps around!
- At the beginning of the process (when the image is just noise), the model might focus on the first few words to get the general shape.
- In the middle, it might focus on a word in the middle of the sentence to fix a grammar error.
- At the end, it might focus on the last word to polish the punctuation.

Because the "anchor" keeps moving, treating it like a permanent foundation is a mistake. If you try to protect a moving target, you end up keeping useless weights and cutting the ones that actually matter at that specific moment.

The Solution: "Sink-Aware Pruning"

The authors propose a new strategy called Sink-Aware Pruning. Here is how it works, using a metaphor:

Imagine you are cleaning a house (the AI model) to make it smaller.

Old Method: You see a heavy, dusty sofa in the living room (the "Sink"). You assume it's the most important piece of furniture, so you refuse to throw it away, even if it's broken or in the wrong room. You end up keeping the sofa and throwing away the useful kitchen tools.
New Method (Sink-Aware): You realize the "sofa" moves around. Sometimes it's in the living room, sometimes in the kitchen, sometimes it's just a pile of dust.
1. Track the movement: You watch where the "sofa" (the attention sink) goes over time.
2. Identify the fakes: If the sofa keeps jumping around and never stays in one place, you realize it's not a structural pillar; it's just temporary clutter.
3. Cut the clutter: You confidently throw away the weights associated with these moving, unstable "sinks."

By removing these "moving anchors," the model becomes much lighter and faster without losing its ability to write good stories.

Why This Matters

It's Faster: By cutting the weights that aren't actually needed, the model runs much quicker.
It's Smarter: It stops blindly following rules from the past (AR models) and adapts to how Diffusion models actually work.
No Retraining Needed: You don't have to teach the model a new way to think. You just trim the fat based on how it behaves, and it works better immediately.

The Takeaway

The paper teaches us that one size does not fit all. Just because a rule works for a writer who writes one word at a time (Autoregressive) doesn't mean it works for a sculptor who refines the whole picture at once (Diffusion).

By realizing that "Attention Sinks" in Diffusion models are like shifting sand rather than solid rocks, the authors found a way to build smaller, faster, and smarter AI models without breaking them.

1. Problem Statement

Diffusion Language Models (DLMs) generate text through iterative denoising processes, which incur significantly higher inference costs (compute and memory) compared to autoregressive (AR) models due to the need for multiple forward passes. To make DLMs practical, model pruning (removing redundant parameters) is essential.

However, existing pruning heuristics are largely inherited from AR models. A dominant heuristic in AR pruning is the "Attention Sink Preservation" rule:

AR Context: In AR models, specific early tokens (e.g., BOS, system prompts) act as stable "attention sinks" that consistently attract disproportionate attention mass. These sinks serve as global anchors to stabilize the residual stream and propagate conditioning information. Consequently, AR pruning methods explicitly protect these tokens.
The Gap: The authors argue that this assumption does not hold for DLMs. In DLMs, the model updates the entire token sequence at every diffusion timestep. The authors hypothesize that attention sinks in DLMs are transient and unstable, shifting positions as the denoising trajectory progresses from high-noise global structure formation to low-noise local refinement. Blindly preserving AR-style sinks in DLMs may lead to suboptimal pruning decisions, retaining unstable parameters that do not contribute to final quality.

2. Methodology: Sink-Aware Pruning

The proposed method, Sink-Aware Pruning, is a post-training pruning strategy that adapts importance scoring to the specific dynamics of diffusion generation. It consists of four main steps:

A. Analysis of Sink Variance

The authors first quantify sink stability by tracking the position of dominant attention sinks across the generation trajectory.

Metric: They define Sink Variance (both spatial and temporal).
- Temporal Variance: Measures how much the sink location shifts across diffusion timesteps.
- Observation: AR models exhibit near-zero temporal variance (sinks are fixed). DLMs exhibit orders of magnitude higher temporal variance, with sink positions drifting significantly throughout the denoising process.

B. Sink Score Estimation

To identify which sinks are "unstable" and candidates for pruning, the method computes a Sink Score ( $\bar{\phi}$ ) for each token position:

Attention Mass Aggregation: At each timestep $t$ , the total attention mass received by token $j$ is calculated across all layers and heads.
Sink Detection: A token is identified as a sink if its attention mass significantly exceeds the average of other tokens.
Soft Scoring: A sigmoid function converts this thresholded condition into a differentiable soft sink score $\phi_t(j)$ .
Averaging: The scores are averaged over a set of calibration timesteps to obtain a step-invariant sink estimate $\bar{\phi}(j)$ .

C. Importance Reweighting (The Core Mechanism)

Instead of hard-coding a rule to "keep" sinks, the method dynamically down-weights them during the importance scoring phase:

Down-weighting Factor: A factor $\omega_j = 1 - \bar{\phi}(j)$ $ω_{j} = 1 - \overset{ˉ}{ϕ} (j)$ is computed.
- If a token is a stable, high-variance sink (in the DLM context), $\bar{\phi}$ is high, and $\omega$ approaches 0.
- If a token is not a sink, $\omega$ approaches 1.
Activation Masking: The original input activations $X$ are scaled by $\omega$ to produce a masked activation $\tilde{X}$ :
$\tilde{X}_{j,:} = \omega_j \cdot X_{j,:}$
Pruning Integration: This masked activation $\tilde{X}$ $\tilde{X}$ replaces the standard $X$ $X$ in existing pruning algorithms (specifically Wanda and SparseGPT).
- For Wanda: The importance score becomes $|W_{ij}| \cdot \|\tilde{X}_{\cdot j}\|_2$ .
- For SparseGPT: The empirical Hessian is computed using $\tilde{X}$ , altering the reconstruction error estimation to prioritize semantically meaningful positions over transient sink positions.

3. Key Contributions

Discovery of Transient Sinks: The paper establishes that attention sinks in DLMs are fundamentally different from AR models; they are temporally unstable and shift across diffusion steps, challenging the "always keep sinks" heuristic.
Sink-Aware Pruning Algorithm: A novel, training-free pruning framework that automatically identifies and down-weights unstable sinks in DLMs, allowing for more aggressive pruning of redundant attention pathways.
Paradigm-Specific Adaptation: Demonstrates that pruning strategies must be aligned with generation dynamics (causal AR vs. iterative diffusion) rather than being blindly transferred.

4. Experimental Results

The method was evaluated on multiple DLMs (LLaDA, Dream, LLaDA-1.5, MMaDA) across standard benchmarks (MMLU, ARC, GSM8K, etc.) and compared against strong baselines (Wanda, SparseGPT, Magnitude Pruning).

Performance Gains: Sink-Aware Pruning consistently outperforms or matches strong baselines across all sparsity levels (25%, 50%, 75%).
High Sparsity Advantage: The improvements are most pronounced at high sparsity (50%–75%). At these levels, naive AR-based heuristics fail significantly, whereas Sink-Aware Pruning maintains better quality-efficiency trade-offs.
- Example: On LLaDA at 75% sparsity, Sink-Aware (Wanda) achieved an average accuracy of 30.94% vs. 29.99% for standard Wanda.
Structured Pruning: The method also yields consistent gains in structured pruning (removing entire heads/layers), with margins increasing as pruning ratios become more aggressive.
Generalization: The approach works effectively across different base pruning algorithms (Wanda and SparseGPT), indicating it acts as a general-purpose correction mechanism.

5. Significance

Theoretical Insight: The paper fundamentally shifts the understanding of attention mechanisms in generative models, proving that "stability" is not a universal property of attention sinks but is dependent on the generation paradigm (AR vs. Diffusion).
Practical Impact: It provides a critical tool for deploying DLMs. By enabling more aggressive pruning without retraining, it significantly reduces the inference cost and memory footprint of DLMs, making them more viable for real-world applications.
Future Direction: It suggests that future model compression techniques must account for the dynamic nature of the generation process, moving beyond static importance metrics derived from AR assumptions.

In conclusion, Sink-Aware Pruning effectively bridges the gap between the unique dynamics of Diffusion Language Models and the need for efficient inference, proving that discarding "unstable" attention sinks is not only safe but beneficial for DLM compression.