sparse-attention-reasoning

name	sparse-attention-reasoning
title	SeerAttention-R: Sparse Attention Adaptation for Long Reasoning
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2506.08889
keywords	["sparse attention","reasoning decoding","efficiency","token selection"]
description	Learn sparse attention patterns for reasoning model decoding via self-distilled gating, achieving 9x speedup at 90% sparsity while maintaining reasoning quality.

SeerAttention-R: Sparse Attention for Reasoning

Core Concept

SeerAttention-R enables efficient long reasoning by learning which tokens are most important for each attention computation. A lightweight gating mechanism learns sparsity patterns during training, allowing models to focus computational resources on relevant tokens. The approach maintains near-lossless reasoning performance while achieving 9x speedup.

Architecture Overview

• Self-distilled gating mechanism: Learn token importance without external labels
• Plug-in design: Integrates with existing models without parameter modification
• Large sparse blocks: 64/128 token blocks maintain coherence while enabling speedups
• Minimal retraining: Effective training on just 400M tokens
• Optimized kernels: TileLang for near-theoretical speedup on H100
• Autoregressive compatible: Seamlessly works with sequential token generation

Implementation

Step 1: Design Sparse Attention Gate

Create lightweight selector for important tokens:

class SparseAttentionGate(torch.nn.Module):
    def __init__(self, hidden_dim: int = 768,
                 sparsity: float = 0.9):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.sparsity = sparsity

        # Gating network (lightweight)
        self.gate = torch.nn.Sequential(
            torch.nn.Linear(hidden_dim, 256),
            torch.nn.ReLU(),
            torch.nn.Linear(256, 1)
        )

        # Temperature for sparsity control
        self.temperature = torch.nn.Parameter(
            torch.tensor(1.0)
        )

    def forward(self, query: torch.Tensor,
               key: torch.Tensor) -> torch.Tensor:
        """Compute sparse mask indicating important keys."""

        batch_size, query_len, _ = query.shape
        key_len = key.shape[1]

        # Compute importance scores
        scores = self.gate(key)  # (batch, key_len, 1)

        # Create mask based on sparsity threshold
        # Keep top (1 - sparsity) fraction of tokens
        k = max(1, int(key_len * (1 - self.sparsity)))

        # Top-k selection
        topk_scores, topk_indices = torch.topk(
            scores.squeeze(-1),
            k,
            dim=-1
        )

        # Create sparse mask
        mask = torch.zeros(batch_size, key_len,
                          device=scores.device)
        mask.scatter_(, topk_indices, )

         mask, topk_indices

     () -> :
        

        
        scores = .gate(key) / .temperature

        
        mask_soft = torch.sigmoid(scores)

        
        mask = (mask_soft > ).()

         mask, mask_soft

Step 2: Implement Self-Distillation Training

Train gating without external labels using model self-distillation:

class SelfDistillationTrainer:
    def __init__(self, model_with_gate):
        self.model = model_with_gate
        self.optimizer = torch.optim.Adam(
            self.model.gate.parameters(),
            lr=1e-4
        )

    def compute_self_distillation_loss(self,
                                      logits_full: torch.Tensor,
                                      logits_sparse: torch.Tensor,
                                      temperature: float = 4.0
                                      ) -> torch.Tensor:
        """KL divergence between full and sparse attention logits."""

        # Compute soft probabilities
        full_probs = torch.softmax(
            logits_full / temperature,
            dim=-1
        )
        sparse_probs = torch.softmax(
            logits_sparse / temperature,
            dim=-1
        )

        # KL divergence
        kl_loss = torch.nn.functional.kl_div(
            torch.log(sparse_probs),
            full_probs,
            reduction='mean'
        )

        return kl_loss

    def train_step(self, input_ids: torch.Tensor,
                  max_tokens: int = 4096) -> dict:
        """Single self-distillation training step."""

        batch_size, seq_len = input_ids.shape

        # Forward pass with full attention (teacher)
        with torch.no_grad():
            output_full = self.model(
                input_ids,
                use_sparse=False,
                output_hidden_states=True
            )
            logits_full = output_full.logits

        
        output_sparse = .model(
            input_ids,
            use_sparse=,
            output_hidden_states=
        )
        logits_sparse = output_sparse.logits

        
        loss = .compute_self_distillation_loss(
            logits_full,
            logits_sparse
        )

        
        .optimizer.zero_grad()
        loss.backward()
        .optimizer.step()

         {: loss.item()}

Step 3: Integrate Sparse Attention into Model

Add sparse attention as plug-in to reasoning model:

class ReasoningModelWithSparseAttention(torch.nn.Module):
    def __init__(self, base_model,
                 sparsity: float = 0.9,
                 block_size: int = 128):
        super().__init__()
        self.base_model = base_model
        self.sparsity = sparsity
        self.block_size = block_size

        # Add gates for each attention head
        num_heads = base_model.config.num_attention_heads
        self.sparse_gates = torch.nn.ModuleList([
            SparseAttentionGate(
                hidden_dim=base_model.config.hidden_size,
                sparsity=sparsity
            )
            for _ in range(num_heads)
        ])

    def forward(self, input_ids: torch.Tensor,
               use_sparse: bool = True) -> dict:
        """Forward pass with optional sparse attention."""

        if not use_sparse:
            # Use standard model
            return self.base_model(input_ids)

        # Apply sparse attention via hooks
        sparse_attention_fn = self._create_sparse_attention_hook(
            use_sparse=True
        )

        # Register hooks for each attention layer
        hooks = []
        for layer in self.base_model.model.layers:
            hook = layer.self_attn.register_forward_hook(
                sparse_attention_fn
            )
            hooks.append(hook)

        
        output = .base_model(input_ids)

        
         hook  hooks:
            hook.remove()

         output

     ():
        

         ():
              use_sparse:
                 output

            attn_output, attn_weights = output

            
            

             (attn_output, attn_weights)

         sparse_attention_hook

     () -> :
        

        seq_len = input_ids.shape[]

        
         block_start  (, seq_len, .block_size):
            block_end = (block_start + .block_size, seq_len)
            block_ids = input_ids[:, block_start:block_end]

            
             torch.no_grad():
                hidden = .base_model.get_hidden(block_ids)

                
                masks = []
                 gate  .sparse_gates:
                    mask, _ = gate(hidden, hidden)
                    masks.append(mask)

        output = .base_model(input_ids)
         output

Step 4: Optimized Kernel Implementation

Implement efficient sparse attention computation:

class OptimizedSparseAttentionKernel:
    """Optimized kernel for sparse attention (pseudocode)."""

    def __init__(self):
        # Would use TileLang or CUDA kernels in production
        pass

    def sparse_attention_forward(self, Q: torch.Tensor,
                                K: torch.Tensor,
                                V: torch.Tensor,
                                sparse_mask: torch.Tensor,
                                block_size: int = 128
                                ) -> torch.Tensor:
        """Compute attention only on sparse tokens."""

        # Q: (batch, num_heads, query_len, head_dim)
        # K, V: (batch, num_heads, key_len, head_dim)
        # sparse_mask: (batch, num_heads, key_len)

        batch_size, num_heads, query_len, head_dim = Q.shape

        # Compute scores only for non-masked tokens
        scores = torch.matmul(Q, K.transpose(-2, -1))
        scores = scores / torch.sqrt(torch.tensor(head_dim))

        # Apply sparse mask
        mask_value = torch.tensor(float('-inf'))
        masked_scores = scores.masked_fill(
            sparse_mask.unsqueeze(2) == 0,
            mask_value
        )

        # Compute attention weights
        attn_weights = torch.softmax(masked_scores, dim=-1)

        # Output (only use non-masked values)
        output = torch.matmul(attn_weights, V)

        return output

Step 5: Evaluate on Reasoning Tasks

Measure speedup and quality retention:

def evaluate_sparse_reasoning(model,
                             benchmark_dataset: list,
                             sparsity: float = 0.9
                             ) -> dict:
    """Benchmark sparse attention on reasoning."""

    import time

    results = {
        "full_attention": {"accuracy": 0.0, "time": 0.0},
        "sparse_attention": {"accuracy": 0.0, "time": 0.0}
    }

    # Test full attention
    start = time.time()
    full_correct = 0

    for sample in benchmark_dataset:
        with torch.no_grad():
            output = model(
                sample["input_ids"],
                use_sparse=False
            )
        full_correct += (output.predictions == sample["label"]).sum()

    results["full_attention"]["accuracy"] = (
        full_correct / len(benchmark_dataset)
    )
    results["full_attention"]["time"] = time.time() - start

    # Test sparse attention
    start = time.time()
    sparse_correct = 0

    for sample in benchmark_dataset:
        with torch.no_grad():
            output = model(
                sample["input_ids"],
                use_sparse=True
            )
        sparse_correct += (output.predictions == sample["label"]).sum()

    results[][] = (
        sparse_correct / (benchmark_dataset)
    )
    results[][] = time.time() - start

    
    speedup = (results[][] /
              results[][])

    results[] = speedup
    results[] = (
        results[][] /
        results[][]
    )

     results

Practical Guidance

Sparsity Levels: 90% sparsity achieves 9x speedup while maintaining near-lossless reasoning performance. Adjust based on latency requirements.

Block Size: Large blocks (64/128 tokens) maintain coherence while enabling efficient computation. Smaller blocks lose context; larger blocks reduce efficiency gains.

Self-Distillation: Training with full model as teacher ensures sparse model preserves reasoning quality. This eliminates need for labeled data.

Minimal Retraining: 400M tokens sufficient for effective sparse gate training. Don't overtrain—risk of distributional shift.

Hardware Optimization: Speedup depends on optimized kernels. Generic sparse operations may be slower than dense—use TileLang or custom CUDA for H100 speedups.

When to Apply: Use SeerAttention-R when inference latency is critical for reasoning models, or when deploying on resource-constrained hardware.

Reference

SeerAttention-R learns to identify and focus on important tokens for reasoning via lightweight gating mechanisms. Key insight: self-distillation from the full model ensures sparse attention preserves reasoning quality while reducing computation. Achieves near-theoretical 9x speedup at 90% sparsity on modern hardware.