HDT: Hierarchical Document Transformer

Efficient Transformer Leverages Document Structures

Most documents are organized into structural constituents like sections, paragraphs, sentences, bulleted lists, figures, and footnotes. This structure is represented in the visual layout and conveys the author's semantic organization of the text. We exploit this document structure by (1) introducing auxiliary anchor tokens to represent each element in the hierarchy, and (2) developing an efficient sparse attention kernel that exchanges information only between tokens at the same level (siblings) as well as between the respective parent and child tokens. By stacking multiple HDT blocks, information from any token can reach any other token. At the same time, the attention pattern of HDT blocks is highly sparse, leading to gains in computational (and memory) efficiency. More specifically, for the document hierarchy introduced above, we prepend additional [SENT] anchor tokens to the beginning of every sentence, [SEC] anchor tokens to the start of each section, and a [DOC] anchor token to the beginning of the document as illustrated in the following figure:

Redesigned Attention Mechanism

The attention pattern is restructured into a multi-level hierarchy. Information is exchanged only between tokens at the same level (siblings) and between parent and child tokens. This method maintains sparsity, significantly enhancing computational and memory efficiency.

Sparse Attention Kernel

A key innovation of HDT is the development of a novel sparse attention kernel based on the Triton library. This kernel efficiently handles the sample-dependent hierarchical attention patterns unique to each document structure.

(Left) The Longformer sparse attention pattern is identical for all samples in a mini-batch.
(Right) Our proposed dynamic hierarchical attention pattern considers the document structure and therefore is different for each sample in a mini-batch.

Implementation

Following FlashAttention, we copy queries, keys and values block-wise to SRAM for fast attention computation using a fused kernel. To maximize the number of empty blocks that can be skipped, we leverage a simple heuristic which is illustrated with an example in the following figure.

Specifically, before copying keys and values to SRAM, we first sort them based on their hierarchy level while keeping the order of the queries unchanged. This ensures adjacency of the most related tokens and hence increases the probability of large empty blocks that can be skipped. Afterwards, we copy the queries, keys and values to SRAM and apply block attention. We process all non-empty blocks in parallel, skipping empty ones.

Attention Pattern (Black) and Processed SRAM Blocks (Grey)

We compare practical computation of Block-Sparse FlashAttention (Left) and our HDT attention kernel (Right) on the same hierarchical attention patterns. We show the attention mask of the first 1k tokens (~25% of the total document size) of two different documents (row 1+2) in black. The blue grids illustrate the 128 x 64 SRAM blocks which are processed in parallel using the fused kernel. All blocks highlighted in grey contain at least one non-zero attention entry and hence require processing. Due to the reordering of keys and values (columns) in HDT, anchor tokens are aggregated within adjacent blocks leading to a larger number of blocks that can be skipped compared to Block-Sparse FlashAttention.

Experiments

We first verify the utility of structure-aware attention patterns on a simple mathematical task -> ListOps. Next, we evaluate our pre-trained encoder-only model on SciRepEval proximity tasks. We also investigate the expressiveness of the anchor token representations using our pre-trained encoder-decoder model on the FacetSum summarization tasks and SCROLLS benchmark. Results on SCROLLS below demonstrate that our model can even be applied to long texts which are not explicitly structured. Please find the results of ListOps, SciRepEval, and FacetSum in our paper.

Efficiency Analysis

Denote the length (number of tokens) of the longest sentence in a document as \(s\) and the input length as \(n\), the theoretical complexity of HDT attention is \(O(n \times s)\). Therefore, longer inputs bring larger computational savings because \(s \ll n\). This advantage makes it possible for us to pre-train our model on 1 GPU in one day with input length 8,192 that is even longer than most of the LLMs pre-training setting.

Below in the left figure we compare GPU runtime and memory usage of different attention layers, including standard dense attention, block-sparse FlashAttention (using our hierarchical pattern), Longformer sparse windowed attention, and HDT hierarchical attention.

Model	Longformer	HAT	HDT (ours)
Complexity	\(O(n \times w)\)/td>	\(O(n \times k)\)	\(O(n \times s)\)
#Params	148.66 M	152.73 M	108.99 M
Time (ms)	178.82\(_{\pm 6.84}\)	77.84 \(_{\pm 2.30}\)	79.8 \(_{\pm 2.96}\)
TFLOPS	5.29 \(_{\pm 0.19}\)	8.95 \(_{\pm 0.26}\)	8.99 \(_{\pm 0.34}\)
Memory	11.25 GB	5.3 GB	5.85 GB

The right table above reports runtime and memory consumption of several long-document models with 12 Transformer-based layers. We report complexity, parameters, inference time, throughput, and memory usage using context length \(n=4096\) and mini-batch size 1. Here, \(w=512\) is the Longformer window size, \(k=128\) is the fixed HAT segment length. \(s\) is the length of the longest sentence in the document. Usually, \(s \ll k\).