An Exploration of Hierarchical Attention Transformers for Efficient Long Document Classification

TL;DR

Long document classification demands both effective understanding and computational efficiency. The authors propose fully pre-trained Hierarchical Attention Transformers (SWE+CSE) and systematically compare architectures, segmentation strategies, and pre-training regimes against Longformer baselines. They find that cross-segment contextualization throughout the model (especially in interleaved layouts) yields the best accuracy-efficiency balance, with end-to-end pre-training substantially outperforming ad-hoc variants. Across downstream tasks, HATs match or exceed Longformer while using less memory and achieving faster throughput, and the work provides code and pre-trained models to facilitate reproducibility.

Abstract

Non-hierarchical sparse attention Transformer-based models, such as Longformer and Big Bird, are popular approaches to working with long documents. There are clear benefits to these approaches compared to the original Transformer in terms of efficiency, but Hierarchical Attention Transformer (HAT) models are a vastly understudied alternative. We develop and release fully pre-trained HAT models that use segment-wise followed by cross-segment encoders and compare them with Longformer models and partially pre-trained HATs. In several long document downstream classification tasks, our best HAT model outperforms equally-sized Longformer models while using 10-20% less GPU memory and processing documents 40-45% faster. In a series of ablation studies, we find that HATs perform best with cross-segment contextualization throughout the model than alternative configurations that implement either early or late cross-segment contextualization. Our code is on GitHub: https://github.com/coastalcph/hierarchical-transformers.

Figures (5)

• Figure 1: Performance - Efficiency trade-off for HAT and Longformer on downstream tasks.
• Figure 2: Attention patterns for the examined architectures: Hierarchical (Segment-wise followed by cross-segment attention) and Sparse (Combination of windowed and global attention) Attention Transformers.
• Figure 3: Top: The two main modules (building blocks) of Hierarchical Attention Transformers (HAT): the Segment-wise (SWE), and the Cross-segment (CSE) encoders. Bottom: The four examined HAT variants.
• Figure 4: Example of a Hierarchical Attention Network with $N\times$ interleaved blocks.
• Figure 5: Text Segmentation Strategies (Greedy, Sentence-wise, Dynamic). In the presented example we have a text which comprises 4 sentences, each one with a different number of tokens. Greedy segmentation leads to split sentences across segments, e.g., the last token of $S_2$ and $S_4$ has been placed in a different segment compared to the rest of the tokens. Sentence-wise segmentation leads to excessive padding and document truncation, e.g., the last sentence ($S_4$) does not fit in the models since the model can encode up to 3 segments. Dynamic segmentation avoids spliting sentences and balances padding and truncation.