[Paper Review] 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 Modeling4 citations

TL;DR

This paper proposes ChunkBERT, a simple yet effective method to extend BERT for long-text classification by chunking input sequences into shorter segments, processing them with BERT, and aggregating representations via a TextCNN layer. The approach enables inference on arbitrarily long texts with only 6.25% of the original memory footprint and achieves state-of-the-art performance on long-text benchmarks, outperforming LongFormer by 5.7% on average and 18% on complex datasets.

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.

Motivation & Objective

• To address the 512-token limit in BERT that restricts its use in real-world long-text applications.
• To develop a simple, memory-efficient method that enables fine-tuning of pre-trained BERT models on arbitrarily long inputs without custom CUDA kernels.
• To evaluate the method on a unified long-text classification benchmark across multiple tasks, including binary, multi-class, and multi-label classification.
• To demonstrate that simple architectural extensions can outperform complex models on long-text tasks, challenging assumptions about model complexity.
• To provide a systematic alternative to random or heuristic truncation that preserves salient information in long documents.

Proposed method

• Input text is divided into fixed-size chunks (e.g., 128 tokens) to bypass BERT’s 512-token limit.
• Each chunk is encoded independently using a pre-trained BERT model to generate contextualized token representations.
• A TextCNN layer is applied to each chunk’s representation to extract local, hierarchical features and reduce dimensionality.
• The outputs from all chunks are concatenated and passed through a final classifier head for downstream classification.
• The method scales linearly with input length and supports arbitrary-length inputs during inference.
• Fine-tuning is performed on standard GPU hardware, with memory usage reduced by a factor of 16 (6.25% of original) through chunked processing.

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,

Experimental results

Research questions

• RQ1 Can a simple chunking and convolution-based extension to BERT outperform complex long-context models on long-text classification tasks?
• RQ2 Does the proposed method maintain high performance while drastically reducing memory usage compared to standard BERT fine-tuning?
• RQ3 How does the performance of ChunkBERT compare to strong baselines and state-of-the-art models like LongFormer on long-text benchmarks?
• RQ4 Does the method generalize across diverse long-text classification tasks, including binary, multi-class, and multi-label settings?
• RQ5 Can the method effectively capture salient information when key content is located beyond the standard 512-token context window?

Key findings

• ChunkBERT achieves an average 5.7% performance improvement over LongFormer across the long-text benchmark, with a 18% improvement on complex datasets.
• The method outperforms strong baselines, including key-sentence selection and random truncation, especially on long samples from the test set.
• Memory usage is reduced to only 6.25% of the original BERT memory footprint, enabling training and inference on standard GPUs.
• The model maintains consistent performance across diverse tasks, including binary, multi-class, and multi-label classification, demonstrating robustness.
• Performance degrades when increasing the chunk size during inference beyond the training size, suggesting caution in extrapolating chunk size.
• Inverted-Eurlex, where key information is at the end, shows improved performance with ChunkBERT, indicating effectiveness in capturing non-local salient content.

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.