Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation

• URL: http://arxiv.org/abs/2605.15913v3
• Date: Fri, 22 May 2026 02:43:55 GMT
• Title: Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation
• Abstract summary: Block attention can improve KV cache reuse in long-context scenarios such as Retrieval-Augmented Generation (RAG)<n>However, its broader application is hindered by two key challenges: the difficulty of segmenting input text into meaningful, self-contained blocks, and the inefficiency of existing block fine-tuning methods that risk degrading performance.<n>We propose block distillation, a training framework that is more efficient than block fine-tuning, which uses a frozen full-attention teacher model to guide the block-attention student.
• Score: 61.19473093799777
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Block attention, which processes the input as separate blocks that cannot attend to one another, offers significant potential to improve KV cache reuse in long-context scenarios such as Retrieval-Augmented Generation (RAG). However, its broader application is hindered by two key challenges: the difficulty of segmenting input text into meaningful, self-contained blocks, and the inefficiency of existing block fine-tuning methods that risk degrading performance. To address these, we first construct SemanticSeg, a large and diverse semantic segmentation dataset containing over 30k instances across 16 categories-including books, code, web text, and conversations with text lengths ranging from 2k to 32k. Using this dataset, we train a lightweight segmenter to automatically partition text into human-instinct-aligned blocks with controllable granularity. Second, we propose block distillation, a training framework that is more efficient than block fine-tuning, which uses a frozen full-attention teacher model to guide the block-attention student. This framework integrates three novel components: block sink tokens to mitigate information loss at block boundaries, block dropout to leverage training signals from all blocks, and token-level loss weighting to focus learning on block-attention-sensitive tokens. Experiments across multiple models and benchmarks demonstrate that our segmenter outperforms heuristic and statistical baselines, and block distillation achieves near-full-attention performance under block attention, establishing a practical and scalable pathway for deploying block attention.

Related papers

• BlockBatch: Multi-Scale Consensus Decoding for Efficient Diffusion Language Model Inference [15.424211656458453]
  Diffusion language models generate text by iteratively denoising multiple token positions in parallel.<n>Small blocks preserve local conditioning but require many denoising steps, whereas large blocks expose more parallelism but can make premature commitments and accumulate cache error.<n>We propose BlockBatch, a training-free online inference framework that executes multiple block-size branches for the same request inside a batched forward pass.
  arXiv  Detail & Related papers  (2026-05-28T01:48:29Z)
• Sparser Block-Sparse Attention via Token Permutation [46.22204775916057]
  We propose Permuted Block-Sparse Attention (textbfPBS-Attn), a plug-and-play method that leverages the permutation properties of attention to increase block-level sparsity.<n>Powered by our custom permuted-FlashAttention kernels, PBS-Attn achieves an end-to-end speedup of up to $2.75times$ in long-context prefilling.
  arXiv  Detail & Related papers  (2025-10-24T09:11:50Z)
• Blockwise SFT for Diffusion Language Models: Reconciling Bidirectional Attention and Autoregressive Decoding [60.06816407728172]
  Discrete diffusion language models have shown strong potential for text generation.<n>Standard supervised fine-tuning misaligns with semi-autoregressive inference.<n>We propose Blockwise SFT, which partitions responses into fixed-size blocks.
  arXiv  Detail & Related papers  (2025-08-27T02:49:33Z)
• DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation [11.910667302899638]
  DiffusionBlocks is a principled framework for transforming transformer-based networks into genuinely independent trainable blocks.<n>Our experiments on a range of transformer architectures demonstrate that DiffusionBlocks training matches the performance of end-to-end training.
  arXiv  Detail & Related papers  (2025-06-17T05:44:18Z)
• The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training [51.84624027213658]
  We propose Blockwise Learning Rate (LR), a strategy that tailors the LR to each block's sharpness.<n>We demonstrate this acceleration across GPT-2 and LLaMA, with model sizes ranging from 0.12B to 2B.<n>We incorporate Blockwise LR into Adam-mini, a recently proposed memory-efficient variant of Adam, achieving a combined $2times$ speedup and $2times$ memory saving.
  arXiv  Detail & Related papers  (2025-02-26T10:06:37Z)
• Block-Attention for Efficient Prefilling [6.916657784681611]
  We introduce Block-attention, an attention mechanism designed to address the increased inference latency and cost in Retrieval-Augmented Generation (RAG) scenarios.<n>By defining each passage as a block, Block-attention enables us to reuse the KV states of passages that have been seen before.<n>Block-attention significantly reduces the time to first token (TTFT) and floating point operations (FLOPs) to a very low level.
  arXiv  Detail & Related papers  (2024-09-14T02:34:26Z)
• Towards Universal Dense Blocking for Entity Resolution [49.06313308481536]
  We propose UniBlocker, a dense blocker that is pre-trained on a domain-independent, easily-obtainable corpus. By conducting domain-independent pre-training, UniBlocker can be adapted to various downstream blocking scenarios without requiring domain-specific fine-tuning. Our experiments show that the proposed UniBlocker, without any domain-specific learning, significantly outperforms previous self- and unsupervised dense blocking methods.
  arXiv  Detail & Related papers  (2024-04-23T08:39:29Z)
• BASS: Block-wise Adaptation for Speech Summarization [47.518484305407185]
  We develop a method that allows one to train summarization models on very long sequences in an incremental manner. Speech summarization is realized as a streaming process, where hypothesis summaries are updated every block. Experiments on the How2 dataset demonstrate that the proposed block-wise training method improves by 3 points absolute on ROUGE-L over a truncated input baseline.
  arXiv  Detail & Related papers  (2023-07-17T03:31:36Z)
• Landmark Attention: Random-Access Infinite Context Length for Transformers [45.69864961773124]
  We present a novel approach that allows access to the complete context while retaining random-access flexibility. Our method uses a landmark token to represent each block of the input and trains the attention to use it for selecting relevant blocks. We demonstrate that our method can obtain comparable performance with Transformer-XL while significantly reducing the number of retrieved tokens in each step.
  arXiv  Detail & Related papers  (2023-05-25T17:53:42Z)
• Structural block driven - enhanced convolutional neural representation for relation extraction [11.617819771034927]
  We propose a novel lightweight relation extraction approach of structural block driven - convolutional neural learning. We detect the essential sequential tokens associated with entities through dependency analysis, named as a structural block. We only encode the block on a block-wise and an inter-block-wise representation, utilizing multi-scale CNNs.
  arXiv  Detail & Related papers  (2021-03-21T10:23:44Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.