Skip to main content Inicio Creadores adu2021 skillxiv dynamic-mask-sparse-attention
dynamic-mask-sparse-attention Sparse attention mechanism combining content-aware and position-aware sparse patterns through dynamic masking. Achieves 10x speedup while maintaining model quality on long-context benchmarks through hardware-friendly implementation.
Ir a la instalación Skills Marketplace Descubre y explora habilidades de IA creadas por la comunidad.
Instalar con Codex o Claude Copia este prompt, pégalo en Codex, Claude u otro asistente, y deja que revise la página de la skill y la instale por ti.
Mostrar detalles del prompt Un comando directo omite el prompt de revisión. Revisa el origen antes de ejecutarlo.
npx skills add https://github.com/ADu2021/skillXiv --skill dynamic-mask-sparse-attentionEl comando permanece en una sola línea. Desplázate horizontalmente para revisarlo antes de copiarlo.
¿Prefieres una copia local? Descarga los archivos que SkillsMP tiene disponibles ahora.
Descargar Zip Descargando... Más de este repositorio meaningful-kebab-case-name Convert arXiv papers into ready-to-use agent skills using category-aware extraction. First classifies the paper into one or more of 11 research categories, then applies a specialized extraction pipeline for each category — because different types of papers produce different types of usable knowledge. A single paper can yield multiple skills if it spans categories. Use this skill whenever the user wants to turn a paper into a skill, extract practical techniques from research, build a skill library from papers, convert arXiv papers into reusable agent instructions, or batch-process multiple papers into skills. Also trigger when someone asks about extracting actionable knowledge from papers, making research practical for LLM agents, or systematically converting academic contributions into structured agent capabilities.
action-quantization-behavior-cloning Establish regret bounds for behavior cloning with discretized actions combining statistical error and quantization error terms. Prove smoothness requirements for safe quantizer design, show that learning-based quantizers fail these requirements, and propose model-based augmentation to reduce error dependence from H² to H.
adaptive-lora-personalized-ranks Dynamically allocate LoRA ranks per-layer during fine-tuning instead of using fixed uniform ranks. Learn optimal rank for each layer and subject via variational framework with discretized exponential distribution, reducing memory footprint while maintaining fidelity and text-alignment.
Ocupaciones relacionadas SOC
Basado en la clasificación ocupacional SOC
name dynamic-mask-sparse-attention title Trainable Dynamic Mask Sparse Attention for Long Context version 0.0.2 engine skillxiv-v0.0.2-claude-opus-4.6 license MIT url https://arxiv.org/abs/2508.02124 keywords ["sparse-attention","long-context","efficient-transformers","dynamic-masking","attention-mechanism"] description Sparse attention mechanism combining content-aware and position-aware sparse patterns through dynamic masking. Achieves 10x speedup while maintaining model quality on long-context benchmarks through hardware-friendly implementation.
Trainable Dynamic Mask Sparse Attention
Dynamic Mask Sparse Attention (DMA) addresses the quadratic complexity bottleneck in standard self-attention for long-context language models. By combining position-aware and content-aware sparse attention through differentiable dynamic masking, DMA achieves dramatic speedups while preserving or improving model performance.
Core Concept
The fundamental insight is that not all token pairs need full attention computation. DMA:
Uses position-aware patterns for efficient fixed computation patternsApplies content-aware masks to identify and focus on critical informationDynamically learns masking through differentiable gradient flowMaintains hardware efficiency through careful implementationPreserves expressiveness by combining both sparse strategies
Architecture Overview
The framework consists of:
Content-Aware Mask Generator : Uses value vectors to identify important tokensPosition-Aware Pattern Layer : Applies efficient fixed sparsity patternsDynamic Mask Combiner : Merges content and position masksHardware-Optimized Attention : Efficient kernel for sparse computationTraining Infrastructure : End-to-end differentiable learning
Implementation Steps
Step 1: Implement content-aware mask generation
Learn which tokens are important for attention:
import torch
import torch.nn as nn
from typing import Tuple , Optional
import math
class ContentAwareMaskGenerator (nn.Module):
"""Generates attention masks based on content importance"""
def __init__ (self, hidden_size: int , num_heads: int , sparsity: float = 0.9 ):
super
self
self
self
self
self
self
1
self
1.0
def
forward
self, values: torch.Tensor,
queries: torch.Tensor
"""
Generate content-aware attention mask.
Args:
values: (batch, num_heads, seq_len, head_dim)
queries: (batch, num_heads, seq_len, head_dim)
Returns:
Mask of shape (batch, num_heads, seq_len, seq_len)
"""
1
1
1
2
max
1
int
1
self
for
in
range
for
in
range
for
in
range
1.0
return
def
compute_loss
self, mask: torch.Tensor
"""
Regularization loss to encourage sparsity.
"""
abs
sum
return
Step 2: Implement position-aware sparse patterns
Create efficient fixed sparsity patterns based on position:
class PositionAwarePattern (nn.Module):
"""Fixed sparse attention patterns based on position"""
def __init__ (self, pattern_type: str = 'local' ):
super ().__init__()
self .pattern_type = pattern_type
def get_local_pattern (self, seq_len: int ,
window_size: int = 64 ) -> torch.Tensor:
"""
Local attention: each token attends to nearby tokens.
Args:
seq_len: Sequence length
window_size: Local window size
Returns:
Mask of shape (seq_len, seq_len)
"""
mask = torch.zeros(seq_len, seq_len, dtype=torch.bool )
for i in range (seq_len):
start = max (0 , i - window_size // 2 )
end = min (seq_len, i + window_size // 2 )
mask[i, start:end] = True
return mask
def get_strided_pattern (self, seq_len: int ,
stride: int = 8 ) -> torch.Tensor:
"""
Strided attention: each token attends to every stride-th token.
Reduces quadratic complexity to linear while maintaining coverage.
"""
mask = torch.zeros(seq_len, seq_len, dtype=torch.bool )
for i in range (seq_len):
mask[i, ::stride] = True
mask[i, max (0 , i-2 ):min (seq_len, i+3 )] = True
return mask
def get_dilated_pattern (self, seq_len: int ,
dilation: int = 4 ) -> torch.Tensor:
"""
Dilated attention: attend with dilated receptive field.
Captures long-range dependencies efficiently.
"""
mask = torch.zeros(seq_len, seq_len, dtype=torch.bool )
for i in range (seq_len):
mask[i, max (0 , i-4 ):min (seq_len, i+5 )] = True
for offset in range (-seq_len, seq_len, dilation):
j = i + offset
if 0 <= j < seq_len:
mask[i, j] = True
return mask
def forward (self, seq_len: int ,
device: torch.device = None ) -> torch.Tensor:
"""Get sparse pattern for given sequence length"""
if self .pattern_type == 'local' :
pattern = self .get_local_pattern(seq_len)
elif self .pattern_type == 'strided' :
pattern = self .get_strided_pattern(seq_len)
elif self .pattern_type == 'dilated' :
pattern = self .get_dilated_pattern(seq_len)
else :
pattern = torch.ones(seq_len, seq_len, dtype=torch.bool )
if device is not None :
pattern = pattern.to(device)
return pattern
Step 3: Combine masks into dynamic sparse attention
Merge content and position masks for efficient computation:
class DynamicMaskAttention (nn.Module):
"""Combines content and position masks for sparse attention"""
def __init__ (self, hidden_size: int , num_heads: int ,
position_pattern: str = 'local' ,
sparsity: float = 0.9 ):
super ().__init__()
self .hidden_size = hidden_size
self .num_heads = num_heads
self .head_dim = hidden_size // num_heads
self .content_mask_gen = ContentAwareMaskGenerator(
hidden_size, num_heads, sparsity
)
self .position_pattern = PositionAwarePattern(position_pattern)
self .query_proj = nn.Linear(hidden_size, hidden_size)
self .key_proj = nn.Linear(hidden_size, hidden_size)
self .value_proj = nn.Linear(hidden_size, hidden_size)
self .out_proj = nn.Linear(hidden_size, hidden_size)
def forward (self, hidden_states: torch.Tensor,
attention_mask: Optional [torch.Tensor] = None ) -> Tuple [torch.Tensor, dict ]:
"""
Compute sparse attention with dynamic masking.
Args:
hidden_states: (batch, seq_len, hidden_size)
attention_mask: Optional padding mask
Returns:
(output, attention_stats)
"""
batch_size, seq_len, hidden_size = hidden_states.shape
query = self .query_proj(hidden_states)
key = self .key_proj(hidden_states)
value = self .value_proj(hidden_states)
query = query.view(batch_size, seq_len, self .num_heads, self .head_dim)
query = query.transpose(1 , 2 )
key = key.view(batch_size, seq_len, self .num_heads, self .head_dim)
key = key.transpose(1 , 2 )
value = value.view(batch_size, seq_len, self .num_heads, self .head_dim)
value = value.transpose(1 , 2 )
scores = torch.matmul(query, key.transpose(-2 , -1 )) / math.sqrt(self .head_dim)
content_mask = self .content_mask_gen(value, query)
pos_pattern = self .position_pattern(seq_len, device=hidden_states.device)
pos_pattern = pos_pattern.unsqueeze(0 ).unsqueeze(0 )
combined_mask = content_mask & pos_pattern
mask_value = torch.finfo(scores.dtype).min
scores = scores.masked_fill(~combined_mask, mask_value)
attention_weights = torch.softmax(scores, dim=-1 )
attention_weights = attention_weights.masked_fill(~combined_mask, 0.0 )
context = torch.matmul(attention_weights, value)
context = context.transpose(1 , 2 )
context = context.contiguous().view(batch_size, seq_len, hidden_size)
output = self .out_proj(context)
stats = {
'mask_density' : combined_mask.float ().mean().item(),
'content_density' : content_mask.float ().mean().item(),
'pos_density' : pos_pattern.float ().mean().item(),
'sparsity_loss' : self .content_mask_gen.compute_loss(combined_mask)
}
return output, stats
Step 4: Implement hardware-optimized sparse kernels
Create efficient implementations for sparse attention computation:
class SparseAttentionKernel :
"""Hardware-optimized sparse attention computation"""
@staticmethod
def sparse_matmul (query: torch.Tensor,
key: torch.Tensor,
mask: torch.Tensor,
scaling_factor: float ) -> torch.Tensor:
"""
Efficient sparse matrix multiplication for attention.
Uses mask to avoid computing attention for masked positions.
"""
batch_size, num_heads, seq_len, head_dim = query.shape
scores = torch.matmul(query, key.transpose(-2 , -1 )) / math.sqrt(head_dim)
mask_float = mask.float ()
mask_value = torch.finfo(scores.dtype).min
scores = scores.masked_fill(~mask, mask_value)
attention = torch.softmax(scores, dim=-1 )
return attention
@staticmethod
def compute_flops_reduction (seq_len: int ,
mask_density: float ) -> float :
"""
Estimate FLOPs reduction from sparsity.
Args:
seq_len: Sequence length
mask_density: Fraction of attention computed
Returns:
FLOPs reduction ratio (1.0 = no reduction, lower = more sparse)
"""
full_flops = seq_len * seq_len
sparse_flops = seq_len * seq_len * mask_density
reduction = full_flops / sparse_flops
return reduction
Step 5: Integrate into training loop
Train the sparse attention mechanism end-to-end:
class DMATrainer :
"""Trains Dynamic Mask Attention with full architecture"""
def __init__ (self, model, attention_module: DynamicMaskAttention,
learning_rate: float = 1e-4 ):
self .model = model
self .attention = attention_module
self .optimizer = torch.optim.AdamW(
model.parameters(),
lr=learning_rate
)
def training_step (self, batch: Dict ) -> Dict :
"""
Single training step with sparse attention.
Args:
batch: Contains 'input_ids', 'labels'
Returns:
Training metrics
"""
input_ids = batch['input_ids' ]
labels = batch.get('labels' , input_ids)
hidden_states = self .model.embed(input_ids)
output, att_stats = self .attention(hidden_states)
logits = self .model.lm_head(output)
lm_loss = torch.nn.functional.cross_entropy(
logits.view(-1 , logits.size(-1 )),
labels.view(-1 )
)
sparsity_loss = att_stats['sparsity_loss' ]
total_loss = lm_loss + 0.01 * sparsity_loss
self .optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(self .model.parameters(), 1.0 )
self .optimizer.step()
return {
'lm_loss' : lm_loss.item(),
'sparsity_loss' : sparsity_loss.item(),
'total_loss' : total_loss.item(),
'mask_density' : att_stats['mask_density' ],
'flops_reduction' : SparseAttentionKernel.compute_flops_reduction(
input_ids.size(1 ),
att_stats['mask_density' ]
)
}
def train_epoch (self, dataloader, num_epochs: int = 3 ):
"""Train for multiple epochs"""
for epoch in range (num_epochs):
total_metrics = {}
for batch_idx, batch in enumerate (dataloader):
metrics = self .training_step(batch)
for key, value in metrics.items():
if key not in total_metrics:
total_metrics[key] = []
total_metrics[key].append(value)
if batch_idx % 100 == 0 :
print (f"Epoch {epoch} , Batch {batch_idx} : Loss={metrics['total_loss' ]:.4 f} , "
f"Density={metrics['mask_density' ]:.3 f} , "
f"Speedup={metrics['flops_reduction' ]:.1 f} x" )
avg_metrics = {k: sum (v) / len (v) for k, v in total_metrics.items()}
print (f"\nEpoch {epoch} Summary:" )
for key, value in avg_metrics.items():
print (f" {key} : {value:.4 f} " )
Practical Guidance When to use Dynamic Mask Sparse Attention:
Long-context language models (>4K tokens)
Scenarios requiring 10x+ efficiency improvement
Models where attention is major bottleneck
Tasks needing both quality and speed
Hardware with good sparse computation support
When NOT to use Dynamic Mask Sparse Attention:
Short sequences where dense is already fast
Tasks requiring full attention (e.g., alignment)
Hardware without sparse optimizations
When quality degradation is unacceptable
sparsity: 0.85-0.95 typical (85-95% of attention masked)position_pattern: 'local' for general, 'strided' for efficiencysparsity_loss_weight: 0.001-0.01 for regularizationwindow_size (local): 64-256 typical Expected characteristics:
Speedup: 7-10x on long sequences
Quality: >95% of dense attention performance
Mask density: 10-15% typical (85-90% sparse)
Training overhead: ~20% from mask generation
Long-context (4K tokens): 10x speedup
Very long (16K tokens): 8-10x speedup
Short context (1K): 2-3x speedup
Memory: 50-60% reduction on long sequences
Reference Trainable Dynamic Mask Sparse Attention. arXiv:2508.02124
