通过自动分割与块蒸馏实现块注意力的泛化

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