用自动分块与蒸馏训练,让模型在长文本中高效用块注意力。
Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation
- 用3万+样本的语义分割数据集训练轻量分块器,按语义自动切块。
- 块蒸馏框架使块注意力性能接近全注意力,且训练效率更高。
- 适合需要长文本处理的生成模型优化,如RAG系统部署。
块注意力将输入分割为互不交互的块,能显著提升长上下文场景(如检索增强生成)中的键值缓存复用率。但其广泛应用受限于两个问题:如何将文本切分为有意义、自包含的块,以及现有块微调方法效率低且易降性能。为此,我们构建了包含超过3万实例、覆盖16类文本(如书籍、代码、网页、对话)的大型语义分割数据集SemanticSeg,文本长度2k至32k。基于此,训练了一个轻量级分块器,可按人类直觉实现可控粒度的自动分块。其次,提出块蒸馏框架,以冻结的全注意力教师模型指导块注意力学生模型,包含三个新组件:块汇点令牌缓解块边界信息丢失,块丢弃利用所有块的训练信号,词级损失加权聚焦于对块注意力敏感的词。多模型、多基准实验表明,该分块器优于启发式与统计基线,块蒸馏在块注意力下实现接近全注意力的性能,为块注意力的实用化与规模化部署提供了可行路径。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。