Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation

📄 arXiv: 2605.15913v1 📥 PDF

作者: Shuaiyi Li, Zhisong Zhang, Yan Wang, Lei Zhu, Dongyang Ma, Chenlong Deng, Yang Deng, Wai Lam

分类: cs.CL, cs.AI

发布日期: 2026-05-15

备注: 16 pages, 2 figures

💡 一句话要点

提出基于自动分段和块蒸馏的块注意力泛化方法，提升长文本处理效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 块注意力 语义分割 知识蒸馏 长文本处理 检索增强生成

📋 核心要点

  1. 长文本处理中，块注意力机制受限于文本分割质量和微调效率，难以充分发挥其在KV缓存重用方面的优势。
  2. 论文提出自动分段和块蒸馏方法，利用语义分割数据集训练分割器，并采用蒸馏框架提升块注意力模型的性能。
  3. 实验结果表明，该方法在多个模型和基准测试中表现优异，分割器优于传统方法，蒸馏框架接近全注意力性能。

📝 摘要（中文）

块注意力通过将输入分割成独立的块来处理，在检索增强生成（RAG）等长上下文场景中具有提高KV缓存重用率的巨大潜力。然而，其更广泛的应用受到两个关键挑战的阻碍：将输入文本分割成有意义的、自包含的块的难度，以及现有块微调方法效率低下且可能降低性能。为了解决这些问题，我们首先构建了SemanticSeg，一个大型且多样化的语义分割数据集，包含超过3万个实例，涵盖书籍、代码、网页文本和对话等16个类别，文本长度从2k到32k不等。使用该数据集，我们训练了一个轻量级的分割器，以自动将文本分割成符合人类直觉的块，并具有可控的粒度。其次，我们提出了一种比块微调更有效的训练框架——块蒸馏，它使用冻结的全注意力教师模型来指导块注意力学生模型。该框架集成了三个新颖的组件：块sink tokens以减轻块边界的信息损失，块dropout以利用来自所有块的训练信号，以及token级别的损失加权以将学习重点放在块注意力敏感的token上。跨多个模型和基准的实验表明，我们的分割器优于启发式和统计基线，并且块蒸馏在块注意力下实现了接近全注意力的性能，为部署块注意力建立了一条实用且可扩展的途径。

🔬 方法详解

问题定义：现有块注意力机制在处理长文本时，面临两个主要问题。一是如何将长文本分割成有意义且自包含的块，传统的启发式或统计方法效果不佳。二是现有块微调方法效率低，且容易导致性能下降，无法充分发挥块注意力的潜力。

核心思路：论文的核心思路是利用语义分割技术自动将文本分割成语义相关的块，并采用知识蒸馏的方法，将全注意力模型的知识迁移到块注意力模型，从而提高块注意力模型的性能和泛化能力。通过这种方式，可以在保证性能的同时，提升长文本处理的效率。

技术框架：整体框架包含两个主要模块：1) 语义分割模块：使用SemanticSeg数据集训练一个轻量级的文本分割器，该分割器能够自动将输入文本分割成多个语义相关的块。2) 块蒸馏模块：使用一个预训练的、冻结的全注意力教师模型，指导块注意力学生模型的训练。学生模型通过学习教师模型的输出，从而获得更好的性能。

关键创新：论文的关键创新在于：1) 构建了一个大规模的语义分割数据集SemanticSeg，用于训练文本分割器。2) 提出了块蒸馏训练框架，该框架比传统的块微调方法更有效，且能够避免性能下降。3) 引入了块sink tokens、块dropout和token级别的损失加权等技术，进一步提升了块注意力模型的性能。

关键设计：在块蒸馏框架中，块sink tokens被添加到每个块的末尾，以缓解块边界的信息损失。块dropout随机丢弃一些块，以利用来自所有块的训练信号。Token级别的损失加权则根据token对块注意力的敏感程度，调整损失函数的权重，从而使模型更加关注重要的token。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，论文提出的分割器在文本分割任务中优于启发式和统计基线。块蒸馏方法在块注意力机制下实现了接近全注意力机制的性能，显著提升了块注意力模型的性能。具体而言，该方法在多个基准测试中取得了SOTA或接近SOTA的结果，证明了其有效性和泛化能力。

🎯 应用场景

该研究成果可广泛应用于需要处理长文本的场景，例如检索增强生成（RAG）、文档摘要、代码生成、对话系统等。通过提高长文本处理的效率和性能，可以提升这些应用的实用性和用户体验，并降低计算成本。未来，该方法有望应用于更多领域，例如长视频理解、长音频处理等。

📄 摘要（原文）

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.