菜单
Accelerate language model generation 3-5x by combining autoregressive and masked token prediction. Works via fine-tuning—no architectural changes needed. Parallel decode non-consecutive tokens with entropy-bounded sampling.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill set-block-decoding-inference-acceleration命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
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.
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.
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.
基于 SOC 职业分类
正在显示 SKILL.md
| name | set-block-decoding-inference-acceleration |
| title | Set Block Decoding is a Language Model Inference Accelerator |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2509.04185 |
| keywords | ["language model","inference acceleration","token prediction","masked generation","diffusion","transformer","decoding"] |
| description | Accelerate language model generation 3-5x by combining autoregressive and masked token prediction. Works via fine-tuning—no architectural changes needed. Parallel decode non-consecutive tokens with entropy-bounded sampling. |
Language model inference is memory-bound: generating each token requires a forward pass through the entire transformer, even though computational density is low. Standard autoregressive decoding generates one token per forward pass, creating a bottleneck for latency-critical applications like real-time chat, coding assistants, and search.
Existing acceleration methods have tradeoffs:
Set Block Decoding (SBD) eliminates these constraints by enabling parallel generation of multiple non-consecutive tokens within a single fine-tuned model, achieving 3-5x reduction in forward passes.
SBD integrates two objectives into a single transformer during training:
During inference, the model samples multiple non-consecutive future tokens in parallel using entropy-bounded sampling, which identifies tokens with low mutual information dependencies and generates them together. This flexibility allows adaptive batch generation without the overhead of speculative decoding or auxiliary heads.
The mathematical foundation: the model learns to decompose the joint conditional distribution P(tokens at positions S | context) into independent or weakly-dependent marginals, enabling parallel sampling.
Training Architecture:
Inference Pipeline:
Memory and Compute:
Start with a standard next-token prediction model (e.g., Llama-3.1 8B). Create training batches with variable-size future token blocks injected as mask tokens.
# Training data preparation
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "meta-llama/Llama-3.1-8B"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Set mask token (add to vocabulary if necessary)
MASK_TOKEN_ID = tokenizer.eos_token_id + 1
def create_sbd_training_batch(input_ids, block_size=None):
"""
Create training example with masked future block.
Args:
input_ids: token sequence of shape (seq_len,)
block_size: size of future block to mask. If None, sample from [2, 16]
Returns:
sbd_input_ids: input sequence with future block masked
past_labels: labels for NTP objective (positions t+1 after context)
masked_labels: labels for MATP objective (positions in masked block)
mask_positions: boolean mask indicating which positions are masked
"""
if block_size is None:
block_size = torch.randint(2, 17, (1,)).item()
seq_len = input_ids.shape[0]
# Sample random position for future block start
max_start = max(1, seq_len - block_size - 1)
block_start = torch.randint(1, max_start + 1, (1,)).item()
block_end = block_start + block_size
# Create masked version
sbd_input_ids = input_ids.clone()
sbd_input_ids[block_start:block_end] = MASK_TOKEN_ID
# Create labels
past_labels = input_ids.clone()
past_labels[block_end:] = -100 # Ignore tokens after masked block in NTP
masked_labels = input_ids[block_start:block_end].clone()
masked_labels_full = torch.full_like(input_ids, -100)
masked_labels_full[block_start:block_end] = masked_labels
mask_positions = torch.zeros_like(input_ids, dtype=torch.bool)
mask_positions[block_start:block_end] = True
return {
'input_ids': sbd_input_ids,
'ntp_labels': past_labels,
'matp_labels': masked_labels_full,
'mask_positions': mask_positions
}
# Test data preparation
sample_ids = torch.tensor(tokenizer.encode("The quick brown fox jumps over the lazy dog."))
batch = create_sbd_training_batch(sample_ids)
Modify the forward pass to compute both NTP and MATP losses. The dual objective trains the model to perform both tasks simultaneously.
# SBD training loss implementation
import torch.nn.functional as F
class SBDTrainingLoss:
def __init__(self, ntp_weight=1.0, matp_weight=1.0):
"""
Weighted combination of NTP and MATP objectives.
Args:
ntp_weight: weight for next-token prediction loss
matp_weight: weight for masked-token prediction loss
"""
self.ntp_weight = ntp_weight
self.matp_weight = matp_weight
def __call__(self, model_output, batch):
"""
Calculate combined SBD loss.
Args:
model_output: model logits of shape (batch_size, seq_len, vocab_size)
batch: dict with 'ntp_labels', 'matp_labels', 'mask_positions'
Returns:
loss: scalar tensor
loss_dict: dict with individual loss components
"""
logits = model_output.logits
vocab_size = logits.shape[-1]
# NTP Loss: standard causal language modeling before masked block
ntp_labels = batch['ntp_labels']
ntp_loss = F.cross_entropy(
logits.view(-1, vocab_size),
ntp_labels.view(-1),
ignore_index=-100
)
# MATP Loss: predict masked positions using bidirectional attention
# During training, the model attends bidirectionally to masked block
matp_labels = batch['matp_labels']
matp_loss = F.cross_entropy(
logits.view(-1, vocab_size),
matp_labels.view(-1),
ignore_index=-100
)
# Combined loss
total_loss = self.ntp_weight * ntp_loss + self.matp_weight * matp_loss
return total_loss, {
'ntp_loss': ntp_loss.item(),
'matp_loss': matp_loss.item(),
'total_loss': total_loss.item()
}
# Training setup
sbd_loss_fn = SBDTrainingLoss(ntp_weight=1.0, matp_weight=1.0)
# In training loop:
# for batch in train_dataloader:
# sbd_batch = create_sbd_training_batch(batch['input_ids'])
# output = model(**sbd_batch)
# loss, loss_dict = sbd_loss_fn(output, sbd_batch)
# loss.backward()
# optimizer.step()
The EB Sampler identifies which masked tokens to unfreeze by sorting positions by conditional entropy and greedily unmasking until entropy budget is exhausted.
# Entropy-Bounded Sampler
import torch
import torch.nn.functional as F
class EntropyBoundedSampler:
def __init__(self, entropy_threshold=1.0):
"""
Sampler that selects tokens to generate based on conditional entropy.
Args:
entropy_threshold: gamma parameter controlling speed-accuracy tradeoff
- Low values (0.5-1.0): fewer tokens per forward pass, closer to autoregressive
- High values (2.0-4.0): more tokens per forward pass, faster but more approximation
"""
self.entropy_threshold = entropy_threshold
def __call__(self, logits, masked_positions):
"""
Select which masked positions to unfreeze.
Args:
logits: model output logits for masked positions, shape (batch, num_masked, vocab_size)
masked_positions: indices of masked tokens in sequence
Returns:
tokens_to_unfreeze: boolean mask of which positions to decode (batch, num_masked)
entropies: conditional entropies of each position (batch, num_masked)
"""
# Compute probability distributions
probs = F.softmax(logits, dim=-1)
# Compute conditional entropy for each position
# H(p) = -sum(p * log(p))
log_probs = F.log_softmax(logits, dim=-1)
entropies = -(probs * log_probs).sum(dim=-1) # (batch, num_masked)
# Sort by entropy (ascending: lowest entropy first)
sorted_entropies, sorted_indices = torch.sort(entropies, dim=-1)
# Greedily unfreeze tokens until cumulative entropy exceeds threshold
cumsum_entropy = torch.cumsum(sorted_entropies, dim=-1)
num_to_unfreeze = (cumsum_entropy <= self.entropy_threshold).sum(dim=-1)
num_to_unfreeze = torch.clamp(num_to_unfreeze, min=1) # Always unfreeze at least 1
# Create selection mask
tokens_to_unfreeze = torch.zeros_like(entropies, dtype=torch.bool)
for b in range(entropies.shape[0]):
tokens_to_unfreeze[b, sorted_indices[b, :num_to_unfreeze[b]]] = True
return tokens_to_unfreeze, entropies
def sample_tokens(self, logits, masked_positions, temperature=1.0):
"""
Sample token IDs for positions to unfreeze.
Args:
logits: model logits for masked positions
masked_positions: indices in original sequence
temperature: sampling temperature (1.0 = no modification, >1.0 = more entropy)
Returns:
sampled_token_ids: sampled tokens for unfrozen positions
unfrozen_mask: boolean mask of unfrozen positions
"""
unfrozen_mask, entropies = self(logits, masked_positions)
# Sample tokens from the distribution
scaled_logits = logits / temperature
probs = F.softmax(scaled_logits, dim=-1)
sampled_tokens = torch.multinomial(
probs.reshape(-1, probs.shape[-1]),
num_samples=1
).reshape(probs.shape[:-1])
# Mask out frozen positions (replace with special token)
sampled_tokens[~unfrozen_mask] = -1 # Placeholder for frozen tokens
return sampled_tokens, unfrozen_mask
# Inference usage
sampler = EntropyBoundedSampler(entropy_threshold=1.5)
# In generation loop:
# - Forward pass with masked block
# - Get logits for masked positions
# - sampled_ids, unfrozen = sampler.sample_tokens(logits, masked_positions)
# - Keep unfrozen tokens, refreeze others
# - Update KV cache and repeat until sequence complete
Complete generation process using the SBD model and entropy-bounded sampler.
# Complete SBD inference loop
def generate_with_sbd(
model,
tokenizer,
prompt,
max_length=256,
entropy_threshold=1.5,
block_size=8,
temperature=1.0
):
"""
Generate text using Set Block Decoding.
Args:
model: SBD-trained language model
tokenizer: tokenizer with MASK_TOKEN_ID configured
prompt: initial prompt string
max_length: maximum generation length
entropy_threshold: gamma parameter for EB sampler
block_size: number of tokens to mask at each step
temperature: sampling temperature
Returns:
generated_text: completed text including prompt
speedup_stats: dict with timing and forward pass counts
"""
import time
device = next(model.parameters()).device
sampler = EntropyBoundedSampler(entropy_threshold=entropy_threshold)
# Encode prompt
input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
seq_len = input_ids.shape[1]
# Tracking
forward_passes = 0
start_time = time.time()
with torch.no_grad():
while seq_len < max_length:
# Determine how many tokens to mask
remaining = max_length - seq_len
current_block_size = min(block_size, remaining)
# Create masked input
masked_input = input_ids.clone()
mask_start = seq_len
mask_end = min(seq_len + current_block_size, max_length)
masked_input = torch.cat([
masked_input,
torch.full((1, current_block_size), MASK_TOKEN_ID, device=device)
], dim=1)
# Forward pass
output = model(masked_input)
logits = output.logits[0, mask_start:mask_end, :]
forward_passes += 1
# Sample using EB sampler
sampled_ids, unfrozen = sampler.sample_tokens(
logits.unsqueeze(0),
torch.arange(mask_start, mask_end),
temperature=temperature
)
sampled_ids = sampled_ids.squeeze(0)
unfrozen = unfrozen.squeeze(0)
# Collect unfrozen tokens
new_tokens = sampled_ids[unfrozen].unsqueeze(1)
if new_tokens.shape[0] > 0:
input_ids = torch.cat([input_ids, new_tokens], dim=1)
seq_len = input_ids.shape[1]
else:
# Fallback: generate at least one token
next_token = torch.argmax(logits[0])
input_ids = torch.cat([
input_ids,
next_token.unsqueeze(0).unsqueeze(0)
], dim=1)
seq_len = input_ids.shape[1]
# Stop if we hit end token
if input_ids[0, -1].item() == tokenizer.eos_token_id:
break
# Decode
generated_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
elapsed = time.time() - start_time
# Estimate speedup: SBD forward passes vs. autoregressive forward passes
autoregressive_passes = seq_len - len(tokenizer.encode(prompt))
estimated_speedup = autoregressive_passes / forward_passes if forward_passes > 0 else 1.0
return generated_text, {
'forward_passes': forward_passes,
'total_tokens_generated': seq_len - len(tokenizer.encode(prompt)),
'elapsed_seconds': elapsed,
'estimated_speedup': estimated_speedup
}
# Example usage
# prompt = "The future of AI is"
# text, stats = generate_with_sbd(
# model,
# tokenizer,
# prompt,
# max_length=256,
# entropy_threshold=1.5,
# block_size=8
# )
# print(f"Generated: {text}")
# print(f"Speedup: {stats['estimated_speedup']:.1f}x")
| Parameter | Range | Impact | Recommendation |
|---|---|---|---|
entropy_threshold (γ) | 0.5-4.0 | Controls tokens per forward pass; lower = fewer tokens = slower but more accurate | Start at 1.5, tune based on latency SLA |
block_size | 2-16 | Training block size; wider distribution = more robust inference | Use 4-8 for balanced training; sample variably |
ntp_weight | 0.5-2.0 | Relative importance of autoregressive objective | Keep at 1.0 unless NTP performance degrades |
matp_weight | 0.5-2.0 | Relative importance of masked prediction objective | Keep at 1.0; both losses are critical |
temperature | 0.8-1.2 | Sampling temperature at inference | Use 1.0 for faithful generation; lower for determinism |
Ideal scenarios:
Strong results observed on:
Avoid SBD in these scenarios:
Insufficient training iterations: SBD requires ~34k iterations to close the performance gap with NTP-only training on 3B models. Scale linearly; 8B models need proportionally more steps. Monitor NTP and MATP losses separately—if either plateaus early, increase learning rate or training duration.
Ignoring loss balance: Removing either NTP or MATP loss causes 7-12% accuracy degradation. Both are essential. Use balanced weighting during training.
Entropy threshold too high: γ > 3.0 often unfreeze most/all masked tokens, negating speedup. Start at 1.5, measure wall-clock time, then adjust. Profile on target hardware (H100, A100) since overhead is hardware-dependent.
Block size mismatch at inference: Training with block_size sampled from 2-16 but inferring with fixed size 8 introduces distribution shift. Use variable block sizes at inference too, or validate that fixed inference block size was included in training distribution.
KV-cache not managed correctly: SBD maintains KV-cache across forward passes. If you don't increment cache positions after sampling, you'll recompute attention. Ensure your generation loop updates past_key_values indices correctly.
Model convergence before reaching target speedup: If after 50k training steps performance hasn't recovered, try:
Key Papers and Concepts:
Roofline Model Analysis:
For H100 GPUs, forward pass overhead is ~1-1.5% per pass. A 3.5x reduction in forward passes translates to ~3.4x wall-clock speedup on inference-bound workloads (batch size 1-4, sequence length 2K+). Memory-bound phases (loading weights, I/O) benefit less; compute-bound phases (attention, FFN on large batches) benefit most.
Reproducibility:
@article{gat2025setblockdecoding,
title={Set Block Decoding is a Language Model Inference Accelerator},
author={Gat, Itai and Ben-Hamu, Heli and Havasi, Marton and Haziza, Daniel and Reizenstein, Jeremy and Synnaeve, Gabriel and Lopez-Paz, David and Karrer, Brian and Lipman, Yaron},
journal={arXiv preprint arXiv:2509.04185},
year={2025}
}