Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation

TL;DR

This paper introduces a semantic segmentation dataset and a block distillation training framework to enhance the generalization and efficiency of block attention in long-context NLP tasks.

Contribution

It presents a large dataset for automatic text segmentation and a novel distillation method that improves block attention performance without degrading accuracy.

Findings

The segmenter outperforms heuristic baselines in segmentation quality.

Block distillation achieves near-full-attention performance efficiently.

The methods improve long-context processing in multiple models.

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.