[论文解读] Breaking the Token Barrier: Chunking and Convolution for Efficient Long Text Classification with BERT

Aman Jaiswal, Evangelos Milios|arXiv (Cornell University)|Oct 31, 2023

Topic Modeling被引用 4

一句话总结

本文提出ChunkBERT，一种简单而有效的方法，通过将输入序列分块为较短的片段，利用BERT处理各片段，并通过TextCNN层聚合表示，从而扩展BERT用于长文本分类。该方法仅需原始内存占用的6.25%即可实现对任意长文本的推理，在长文本基准测试中达到最先进性能，平均比LongFormer高出5.7%，在复杂数据集上高出18%。

ABSTRACT

Transformer-based models, specifically BERT, have propelled research in various NLP tasks. However, these models are limited to a maximum token limit of 512 tokens. Consequently, this makes it non-trivial to apply it in a practical setting with long input. Various complex methods have claimed to overcome this limit, but recent research questions the efficacy of these models across different classification tasks. These complex architectures evaluated on carefully curated long datasets perform at par or worse than simple baselines. In this work, we propose a relatively simple extension to vanilla BERT architecture called ChunkBERT that allows finetuning of any pretrained models to perform inference on arbitrarily long text. The proposed method is based on chunking token representations and CNN layers, making it compatible with any pre-trained BERT. We evaluate chunkBERT exclusively on a benchmark for comparing long-text classification models across a variety of tasks (including binary classification, multi-class classification, and multi-label classification). A BERT model finetuned using the ChunkBERT method performs consistently across long samples in the benchmark while utilizing only a fraction (6.25\%) of the original memory footprint. These findings suggest that efficient finetuning and inference can be achieved through simple modifications to pre-trained BERT models.

研究动机与目标

• 解决BERT中512个标记的限制，该限制限制了其在真实世界长文本应用中的使用。
• 开发一种简单且内存高效的方案，使预训练BERT模型能够在无需自定义CUDA内核的情况下微调任意长度的输入。
• 在涵盖二分类、多分类和多标签分类的多个任务上，于统一的长文本分类基准上评估该方法。
• 证明简单的架构扩展可在长文本任务上超越复杂模型，挑战关于模型复杂度的既有假设。
• 提供一种系统性替代方案，以取代随机或启发式截断，从而保留长文档中的关键信息。

提出的方法

• 将输入文本划分为固定大小的块（例如128个标记），以绕过BERT的512个标记限制。
• 使用预训练的BERT模型独立编码每个块，生成上下文相关的标记表示。
• 对每个块的表示应用TextCNN层，以提取局部、分层特征并降低维度。
• 将所有块的输出拼接后输入最终的分类头，用于下游分类任务。
• 该方法在输入长度上呈线性扩展，推理时支持任意长度的输入。
• 微调在标准GPU硬件上进行，通过分块处理将内存使用量降低为原始值的1/16（即6.25%）。

Figure 1: Overview of BERT Input chunking. First, the complete input $(T)$ is tokenized using the BERT tokenizer and split into chunks of size $(C)$ e.g, 128 tokens each. The input chunks are independently processed by BERT and produce token embeddings for each chunk of size $(C\times 768)$ . Here,

实验结果

研究问题

• RQ1 对BERT进行简单的分块与卷积扩展，是否能在长文本分类任务上超越复杂的长上下文模型？
• RQ2 与标准BERT微调相比，该方法是否在大幅降低内存使用量的同时仍保持高性能？
• RQ3 在长文本基准上，ChunkBERT与强基线模型及SOTA模型（如LongFormer）相比表现如何？
• RQ4 该方法是否能在涵盖二分类、多分类和多标签设置的多样化长文本分类任务上泛化？
• RQ5 当关键内容位于标准512个标记上下文窗口之外时，该方法是否能有效捕捉显著信息？

主要发现

• 在长文本基准上，ChunkBERT平均比LongFormer高出5.7%，在复杂数据集上高出18%。
• 该方法优于强基线模型，包括关键句选择和随机截断，尤其在测试集中长样本上的表现更优。
• 内存使用量降低至原始BERT内存占用的6.25%，使标准GPU上可实现训练与推理。
• 模型在二分类、多分类和多标签分类等多种任务上均保持一致的高性能，展现出良好的鲁棒性。
• 当推理时的块大小超过训练时的大小时，性能会下降，提示在扩展块大小时需谨慎。
• 在Inverted-Eurlex数据集中，关键信息位于文档末尾，ChunkBERT表现出性能提升，表明其能有效捕捉非局部显著内容。

Figure 2: Chunk Attention Mask: illustration of chunk attention for the complete input $T$ . The space complexity becomes $O(nC^{2})$ , where $n$ is number of chunks and $C$ is the chunk-size.