HDT: Hierarchical Document Transformer

Authors: haoyu  he
markus  flicke
jan  buchmann
iryna  gurevych
andreas  geiger
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-21 10:29:31

Abstract: In this paper, we propose the Hierarchical Document Transformer (HDT), a novel sparse Transformer architecture tailored for structured hierarchical documents. Such documents are extremely important in numerous domains, including science, law or medicine. However, most existing solutions are inefficient and fail to make use of the structure inherent to documents. HDT exploits document structure by introducing auxiliary anchor tokens and redesigning the attention mechanism into a sparse multi-level hierarchy. This approach facilitates information exchange between tokens at different levels while maintaining sparsity, thereby enhancing computational and memory efficiency while exploiting the document structure as an inductive bias. We address the technical challenge of implementing HDT's sample-dependent hierarchical attention pattern by developing a novel sparse attention kernel that considers the hierarchical structure of documents. As demonstrated by our experiments, utilizing structural information present in documents leads to faster convergence, higher sample efficiency and better performance on downstream tasks.

Main paper: 
 introduction: 
 Many natural language processing tasks including summarization and question answering require language models to encode long structured documents such as scientific papers or Wikipedia articles into meaningful representations. Since 2017, attention-based Transformer architectures Vaswani et al. (2017) have established themselves as the dominant modeling paradigm, demonstrating state-of-the-art performance on numerous NLP tasks. While dense attention yields context-rich representations, its computational complexity grows quadratically with the input. This is problematic when processing large documents, in particular on resource constrained hardware such as consumer GPUs.At the same time, most documents are naturally structured: Words form sentences, sentences form sections and sections form documents. Surprisingly, this structure is largely ignored by many existing language models that typically consider their context as a "flat" sequence of tokens. While transformer models can in principle learn to generalize hierarchically, structural grokking requires extremely long training (Murty et al., 2023) . Moreover, the performance of flat long context models depends on the position of relevant information (Liu et al., 2023) . We hypothesize that exploiting the structure of documents explicitly yields two major benefits: (1) Imposing the structure of documents as an inductive bias improves sample efficiency and generalization. (2) Adapting the attention pattern to the document structure leads to sparse representations which reduce computational and memory complexity, and enable processing of long documents even on consumer hardware.To test this hypothesis, we propose the Hierarchical Document Transformer (HDT), a novel sparse Transformer architecture for processing hierarchically structured documents. More specifically, we first introduce auxiliary anchor tokens for all structural elements such as sentences, sections and documents as illustrated in Fig. 1a . Second, we redesign the Here, regular tokens are illustrated in green, and auxiliary anchor tokens in yellow (document), red (section) and blue (sentence). Each token attends to its parent, siblings and children. Cross-level attention is illustrated using color gradients in the attention matrix. Utilizing structural information present in documents leads to faster pre-training (b) and better performance on downstream tasks. We use the held-out validation set in (b) to calculate the MLM loss.attention pattern of a Transformer block into a multi-level hierarchy where information is exchanged 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. In contrast to previous hierarchical models Wu et al. (2021a) ; Chalkidis et al. (2022) , our HDT block supports deeper hierarchies and communicates at all hierarchy levels simultaneously such that stacking of only a few HDT layers establishes communication between all tokens. contextualization is performed iteratively over several layers. While hierarchical processing improves efficiency, separation of contextualization levels into different layers results in two problems: (1) Contextualization is not conducted simultaneously across all levels, hence more layers are required potentially.(2) Because different layers assume different functions, they can't serve as drop-in replacements for common Transformer architectures. HDT resolves these problems by simultaneous contextualization over all hierarchy levels in a single layer using an efficient hierarchical attention pattern.Due to the quadratic memory and time complexity of the self-attention block of Transformer models (Vaswani et al., 2017) , major efforts have been devoted to reducing complexity, in particular when applying attention to longer sequences. One branch of work focuses on using sparse attention patterns to reduce computation while maintaining expressiveness including Sparse Transformer (Rewon et al., 2019) , ETC (Ainslie et al., 2020 ), BigBird (Zaheer et al., 2020) , Longformer (Beltagy et al., 2020) and CoLT5 (Ainslie et al., 2023) among others. Another branch of work aims at speeding up attention computation by considering GPU hardware characteristics. FlashAttention 1 & 2 (Dao et al., 2022; Dao, 2023) propose tiling with block-wise attention to reduce memory IOs and compute attention blocks in parallel in Static Random Access Memory (SRAM). Matteo et al. (2023) extends the idea of FlashAttention to key/query dropping and hashing-based attention that was originally proposed by Kitaev et al. (2020) . Our work builds upon both branches: We propose a dynamic sample-dependent sparse attention pattern which exploits the structure of text documents. For efficiency, we follow Dao et al. (2022) and implement this pattern as a customized memory-aware kernel. 
methodology: 
 We now introduce the proposed Hierarchical Document Transformer (HDT) for efficient long-document modeling. First, we briefly recap the standard transformer model. Next, we introduce the proposed HDT block. Finally, we describe the design of the encoder-only as well as encoder-decoder HDT architectures which are used in our experiments in Section 4.The original transformer encoder model by Vaswani et al. (2017) is composed of Multi-Head Self-Attention (MHSA) layers that are interleaved with shallow feed-forward networks and residual connections. Because the attention pattern is identical across multiple heads in an attention layer, for simplicity of notation, we represent queries, keys and values for a single attention head as Q, K, V ∈ R n×d k , where n is the sequence length and d k the head dimension. All MHSA layers in the model, as well as the embedding layers, produce outputs of dimension d model . We further denote the number of heads as h, henced model = d k × h. A [DOC] [SEC] [SENT] T1 T5 [SEC] [SENT] T2 [SENT] T3 T4 Tree Level 0 0 0 1 0 0 1 0 1 2 0 0 1 1 0 1 1 2 2 2 0 1 1 1 2 2 2 2 2 2 1 -Section 2 -Sentence 3 -Token = = = T6 2 2 2 2 3 4Figure 3 : Hierarchical Positional Encoding. We represent the position of each token in the hierarchy with one linear index p l per hierarchy level l yielding an index vector p = (p 1 , . . . , p L ) T . Above, we show an example with L = 3 levels. Note that level 0 (document) does not require an index. Each index in p is passed through sinusoidal encoding functions which are summed over all levels to form the final encoding vector according to Eq. (2).self-attention head computes its output O ∈ R n×d k as follows:A = QK T √ d k S = softmax(A) O = SV (1)where the softmax operator is applied row-wise. The standard attention block has O(n 2 ) time and memory complexity where n is the input length. It further does not explicitly take the structure of documents into account.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 (Taylor & Beach, 1984; Guthrie et al., 1991) . While our model is general and can handle arbitrary hierarchical structures, for simplicity we will focus our exposition on a document hierarchy with three levels: tokens, sentences and sections. 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 Fig. 2 .We extend the sinusoidal position encoding to model L hierarchy levels. To inform each token about its position within the hierarchy, we assign it one linear index p l per hierarchy level as illustrated in Fig. 3 , yielding an index vector p = (p 1 , . . . , p L ) T . Each index in p is passed through a set of standard sinusoidal encoding functions which are summed over all levels to form the final hierarchical positional encoding (HPE) vector at p:HPE(p, i) = L ∑ l=1 sin(ω k p l ) if i = 2k cos(ω k p l ) if i = 2k + 1 where ω k = 1 10000 2k/d model (2)Hierarchical Attention: To model document structure, we impose a document-specific attention pattern by modifying the standard Transformer attention in Eq. (1) as follows t j , respectively. The attention mask M ∈ {0, 1} n×n is defined such that information is only exchanged directly between tokens at the same level as well as between the respective parent and child tokens. For a 3-level document structure, we first compute the attention mask for each hierarchy level separatelyA = QK T √ d k S = softmax(A ⊙ 1 M ) O = SV (3) with (A ij ⊙ 1 M ij ) = A ij if M ij = 1 and -∞ if M ij = 0.EQUATIONEQUATIONEQUATIONwhere [•] denotes the Iverson bracket which evaluates to 1 if the argument is true and 0 otherwise. Finally, we perform an OR operation to obtain the full attention mask M:M = M DOC ⊕ M SEC ⊕ M SENT (7)Note that M is highly sparse in practice (see Fig. 10 for an example) and hence reduces theoretical complexity from O(n 2 ) to O(n s) where s is the length of the longest sentence in the document. Thus, computational savings are largest for long documents for which s ≪ n. However, when using parallel hardware (e.g., GPUs) as customary in deep learning, special care has to be taken regarding the kernel design to translate these theoretical savings into actual wall-clock time reduction. We hence develop a custom hierarchical attention kernel.More specifically, we build upon the recent tiling-based ideas of FlashAttention (Dao et al., 2022; Dao, 2023) . FlashAttention partitions attention computation into small 128 × 64 token blocks that can be computed efficiently in SRAM. This is in contrast to classical attention implementations that materialize intermediate outputs to slow High Bandwidth Memory (HBM). However, naïvely partitioning the attention matrix into regular blocks is suboptimal given the uniformly distributed entries of M. Furthermore, each document has a different structure and hence leads to a different sparsity pattern in the attention matrix.To maximize the number of empty blocks that can be skipped, we leverage a simple heuristic which is illustrated with an example in Fig. 4 . Specifically, before copying keys and values to SRAM, we first sort them based on their hierarchy level from l = 0 to L 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 as illustrated in Appendix Fig. 10 . Afterwards, we copy the queries Q i , keys K j and values V j of block (i, j) to SRAM and apply block attention. We process all non-empty blocks in parallel, skipping empty ones. Finally, we write the result O i back to HBM using the online softmaxMED MAX 1 2 3 3 3 9 SM (SUM-MOD) [SUM-MOD 3 [MED 1 2 3][MAX 9 3]]Ground truth: 4 algorithm (Milakov & Gimelshein, 2018) . The complete algorithm for a forward pass is provided in Appendix Algorithm 1. While we found this simple sorting heuristic to work well in practice (see Fig. 8 for a runtime comparison to Block-Sparse FlashAttention and Appendix Fig. 10 for qualitative examples), more advanced permutation algorithms for sparse matrices (Aykanat et al., 2004; Ferris & Horn, 1998; Hendrickson & Kolda, 2000) could possibly lead to larger computational savings at the cost of additional inference time per data sample. We will leave the investigation of such trade-offs for future work.We realize the encoder as a stack of N identical HDT blocks where each block is composed of two sub-layers as illustrated in Fig. 1a . The first sub-layer performs hierarchical multi-head self-attention as described above, and the second is a simple position-wise fully connected feed-forward network (FFN). We use layer normalization at the beginning of each sub-layer as well as residual connections (He et al., 2016) for each layer. We apply the hierarchical positional encoding to the input tokens and add them to the respective token embeddings as input to the encoder stacks. The decoder is composed of a stack of N identical standard Transformer decoder blocks, including a sub-layer of causal multi-head attention that prevents tokens from attending to subsequent positions, a sub-layer to perform multi-head cross attention wrt. the output of the encoder stack, and a standard feed-forward network (FFN). Following T5 (Raffel et al., 2020) , we use relative positional encoding for the decoder. 
experiments: 
 We first demonstrate the utility of structure-aware attention patterns on a simple mathematical task. Next, we show the effectiveness of our encoder-only model on SciRepEval proximity tasks. We also investigate the expressiveness of the anchor token representations using our encoder-decoder model on the FacetSum summarization tasks. Experiments on SCROLLS demonstrate that our model can even be applied to long texts which are not explicitly structured. Finally, we provide a detailed efficiency analysis. As proof of concept, we first compare our encoder model (HDT-E) to BERT (Devlin et al., 2019) , Longformer (Beltagy et al., 2020) and HAT (Chalkidis et al., 2022) on the ListOps mathematical reasoning task (Nikita & R., 2018) . The ListOps dataset is composed of simple list operations (Modular Sum, Minimium, Maximum, Median) with 2-5 operands and maximum tree depth 20. An example is shown in Fig. 5 . Fig. 6 shows the different types of attention patterns we use. The combination of red+green+blue entries corresponds to the attention pattern defined in Section 3.2. However, the special structure of the problem admits of further increasing sparsity by using a "causal" mask where operands do not attend to their operators (green+blue) and where operands do not attend to each other (blue only). Our ablations in Table 1 show that as we increase the inductive sparsity bias, test set accuracy improves. As illustrated in Fig. 7 , whilst BERT, Longformer and HAT use a learned positional embedding, we find that HDT outperforms all baselines on ListOps even without using any positional embedding as the operators in the ListOps task are invariant to the positions of their operands.We now investigate the utility of HDT for long-document processing tasks. As we are interested in efficiency, we conduct all experiments using a constrained compute budget. Following Geiping & Goldstein (2023) , we pre-train our encoder models (HDT-E) for 24 hours on 1 GPU (∼10k steps) and our encoder-decoder models (HDT-ED) for 72 hours on 4 GPUs (∼50k steps) using a batch size of 128. Compared to the baselines, our models require 5-20 times fewer pre-training steps to reach better or comparable downstream task performance. We leave scaling HDT to larger sizes and more data for future work.Data: To train our proposed structure-aware models, we build a large full-text document corpus from unarXive (Saier et al., 2023) (1.9M arXiv papers), HUPD (Suzgun et al., 2023) (utility patent applications) and the latest Wikipedia dump processed with Gensim ( Řeh ůřek & Sojka, 2010) . In total, our corpus includes over 12M long documents with extracted structural information. Our word tokenizer utilizes a vocabulary size of 32,768 tokens, trained with Byte Pair Encoding (BPE) (Sennrich et al., 2016) . For splitting sentences, we use the NLTK sentence tokenizer1 . We report results of HDT-E pre-trained with and without arXiv data on a full-document version of SciRepEval (Singh et al., 2023) . We evaluate HDT-ED on the FacetSum (Meng et al., 2021) and SCROLLS (Shaham et al., 2022) benchmarks which contain tasks for long text reasoning requiring context modeling. Benchmarks details are provided in Appendix Section A.2.2.We compare HDT-E to Hierarchical Attention Transformer (HAT) (Chalkidis et al., 2022) which is the closest work to us that applies attention hierarchically to reduce computation/memory usage for processing long text. HDT-E differs from HAT as follows:(1) HDT-E takes the natural structure of text to build hierarchies while HAT cuts long text into fixed-size segmentations that do not adheare to the structure of natural language, (2)HDT-E enables simultaneous information passing across all levels in a single layer while HAT requires multiple blocks to contextualize information between hierarchies. Moreover, HAT only considers two hierarchy levels. Despite the flexibility of our model, it achieves comparable latency/throughput supported due to our highly efficient kernel implementation. As the sparse attention baseline, we choose Longformer (Beltagy et al., 2020) and its encoder-decoder variant Longformer-Encoder-Decoder (LED) which are effective on various long-text tasks in previous works (Tay et al., 2021; Dasigi et al., 2021; Tay et al., 2023a) . We also include the current SotA methods SciBERT (Beltagy et al., 2019) and SciNCL (Ostendorff et al., 2022) as additional SciRepEval baselines. Both baselines use dense attention with an input length of 512 capturing only title and abstract. Unless stated otherwise, we use the officially released code and models by the original authors.We pre-train the encoder-only model HDT-E and the encoderdecoder model HDT-ED from scratch on our training corpus. All models are pre-trained using an input length of 8,192 tokens and a mini-batch size of 128 (via gradient accumulation). HDT-E is pre-trained on the standard Masked Language Modeling (MLM) objective with a mask ratio of 15%. We also pre-train Longformer from scratch in the same setting as ours to study the effectiveness of the attention pattern we propose. Fig. 1b shows that our model converges faster than Longformer. HDT-ED is pre-trained on UL2 (Tay et al., 2023b) , which is a unified pre-training paradigm with a range of denoising tasks. We fine-tune all models on the downstream tasks using the default settings, see Section A.2.2 for details. To adapt SciRepEval for long document models, we hence consider a subset of SciRepEval for which full-text articles are available in unarXive (Saier et al., 2023) . Table 2 shows the tasks with the largest number of samples, which are proximity tasks involving ranking a set of candidate papers by their relatedness to a query paper. Results on other SciRepEval tasks can be found in Appendix Table 6 . Without any fine-tuning, we take the hidden representation of the first token (either [CLS] for baseline models or [DOC] for HDT-E) and sort papers by cosine similarity between paper representations. The first set of results in Table 2 ("Pretrained Only") shows the mean average precision for models pre-trained on MLM. Our results demonstrate that both, pre-training on full-text documents and considering the structure of documents, leads to substantial performance gains in this setting. Moreover, our model pre-trains significantly faster (10k iterations) than Longformer (64k) and HAT (50k) and can be trained from scratch while Longformer and HAT require parameter initialization from BART (Lewis et al., 2020) and RoBERTa (Liu et al., 2019) (respectively) for best performance. This result underscores the sample efficiency of our structure-aware HDT model.Currently, the state-of-the-art performance on SciRepEval is reported by SciNCL (Ostendorff et al., 2022) which finetunes a pre-trained SciBERT (Beltagy et al., 2019) model using a contrastive learning objective. When comparing in this setting ("Pretrained + Finetuned with Contrastive Learning"), we observe that HDT-E again outperforms all baselines, despite with a smaller margin. Notably, HDT-E outperforms SciNCL even though SciNCL has access to 684k triplets for finetuning while HDT-E is trained only on those 19k triplets for which full-text documents are publicly available. In addition, to test our model's reliance on scientific documents for pre-training, we exclude the entire arXiv corpus from our pretraining dataset. We observed that our model's performance nearly matches HAT in the pre-training-only setting and surpasses HAT in the contrastive learning setting, even when using only out-of-domain data for pre-training.A major advantage of our model is its flexibility. As tokens are hierarchically grouped via anchor tokens, the document's content is compressed at different levels of granularity. We can query this information by letting the decoder of HDT-ED attends to only a subset of the tokens or anchor tokens (e.g., [SEC] ).To better understand how much information is retained where in our model, we design a 4 : Results on the SCROLLS Summarization, QA and NLI Benchmark. We compare HDT-ED (pre-trained for 12 GPU days) to Longformer-Encoder-Decoder (LED) on the official SCROLLS benchmark without document structure. We choose LED as baseline as it has a comparable number of parameters (162M) to HDT-ED (124M). We remark that neither model is competitive with state-of-theart billion-parameter models such as CoLT5 XL (score 43.5) which are trained on large GPU clusters.summarization task based on FacetSum (Meng et al., 2021) which provides section-level summaries for Emerald journal articles. Table 3 shows our results. We observe performance on par with LED base even when attending only to the anchor tokens, which demonstrates the expressiveness of the learned intermediate representations. When additionally attending to all regular tokens (as LED base does), performance increases further. The right part of the table shows generalization performance when training on "Findings" and "Value", but evaluating "Purpose" and "Method" summaries.While our model benefits from the structure of documents as present in our SciRepEval and FacetSum evaluations, we are also interested in the applicability of our model to less structured "flat" long-text tasks. Towards this goal, we evaluate our model on the SCROLLS benchmark, which contains 7 tasks spanning multiple domains that require reasoning over long texts for which document structure is not available. To apply our model to such inputs, we use pseudo sections composed of a fixed number of 32 sentences for hierarchical attention, while sentences are extracted as usual. We fine-tune HDT-ED on the 7 tasks in SCROLLS separately and submit the test set predictions to the public leaderboard2 . Table 4 show our results. Compared to LED base with a similar number of parameters, HDT-ED improves by over 2 score points. This demonstrates that our model is also effective on "flat" long-text tasks. However, unsurprisingly, our results are not competitive with state-of-the-art billion-parameter models such as CoLT5 XL which are trained for many epochs on large industrial GPU clusters.Efficiency Analysis: Denoting the length (number of tokens) of the longest sentence in a document as s, the theoretical complexity of HDT attention is O(n × s). Fig. 8 compares GPU runtime and memory usage of different attention layers, including standard dense attention, block-sparse FlashAttention (using our pattern), Longformer sparse windowed attention, and HDT attention. All kernels are fed with real data, i.e., documents, which are further transformed to multi-head queries, keys, and values using 12 heads with head dimension 64. We report the runtime for one batch with 4 samples on an NVIDIA A100 GPU, and plot peak GPU memory consumption at lengths 16k and 32k for each kernel, except standard attention which exceeds 40 GB at 16k tokens. Besides, we also compare runtime and memory consumption for three complete 12-layer models using different attention mechanisms in Table 5 . Despite its flexibility, HDT-E achieves runtime and memory consumption on par with HAT which uses 2 layers and fixed segment length. long-document models with 12 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. 
conclusion: 
 We presented Hierarchical Document Transformer (HDT), a novel approach for encoding long documents efficiently. By explicitly incorporating document structure into the attention mechanism, we achieve sparse representations, reducing computational complexity while improving sample efficiency and generalization. We believe that hierarchical text representations offer many exciting opportunities in the future: Extending the hierarchical structure down to byte-level could enable token-free language models. Hierarchical ideas might also inspire novel decoder architectures that generate language in a structured hierarchical fashion. Hierarchical models may also be advantageously combined with other models such as state space models, RNNs or ConvNets. Finally, it still remains unclear if scaling laws also hold for hierarchical language models as they do for LLMs, and if hierarchical language models hold similar potential for unlocking emergent abilities. In this section, we compare the proposed dynamic hierarchical attention pattern to the fixed sparse attention pattern of Longformer and provide algorithmic details of its implementation. Subsequently, we illustrate how our customized attention kernel reduces computational overhead through qualitative examples.Given that our attention pattern is rooted in document structure, it is different for each sample. Hence, we require an attention mechanism that is able to cope with sampledependent attention patterns. Fig. 9 provides a visual illustration comparing our dynamic attention pattern to the fixed attention pattern of Longformer (Beltagy et al., 2020) , assuming a mini-batch size of 2 for clarity. Notably, within a mini-batch containing multiple samples, the attention pattern applied to each sample varies, thereby complicating the realization of our attention pattern using existing attention kernels. Moreover, due to the presence of anchor tokens positioned at the onset of different hierarchies and attending to one another, the attention devoted to these anchor tokens can be notably fine-grained and sparse within the attention mask. This phenomenon is demonstrated in Fig. 10 , depicting hierarchical attention masks derived from real documents. Note that due to space limitations only the first 1k tokens (∼25% of the total document size) are visualized. To optimize the number of empty blocks that can be skipped, we employ the sorting heuristic outlined in Section 3.2. The core mechanism of our hierarchical sparse attention kernel is based on and extends FlashAttention Dao et al. (2022) . Our algorithm for a forward pass of the attention layer is provided in Algorithm 1. This algorithm is executed in parallel for all elements of the Cartesian product B × H, where B represents the batch size and H denotes the number of heads, assuming input tensors of shape (B × H × n × d k ).Our primary adjustment to the FlashAttention algorithm is in line 2 of Algorithm1, where we implement a sorting mechanism for the key and value based on their hierarchy levels, specifically prioritizing [DOC] tokens followed by [SEC] and [SENT] tokens before normal tokens. This modification is grounded in the observation that the computational inefficiency in block-wise hierarchical attention predominantly stems from attention interactions between anchor tokens situated distantly from each other, resulting in excessively sparse blocks. Fig. 10 offers a visual comparison between the computation patterns of our kernel and block-sparse FlashAttention, demonstrating the efficacy of our approach.Algorithm 1 Forward Pass of the HDT Hierarchical Attention Kernel Require: Query, key, value Q, K, V ∈ R n×d k , query block size M, key-value block size N, softmax statistics vectors m ∈ R n×n and l ∈ R n×n , output tensorO n×d k . 1: Partition Q into T r = n M blocks Q 1 , . . . , Q T r 2: Sort K, V according to the hierarchy level from l = 1 to L. 3: Partition K, V into T c = n N blocks K 1 , . . . , K T c , and V 1 , . . . , V T c 4: for 1 ≤ i ≤ T r do 5: Load Q i , O i , m i , l i from HBM to SRAM 6: for 1 ≤ j ≤ T c do 7:Load K j , V j from HBM to SRAM 8:On chip, compute mask M according to Equation (7) 9:if Non-zero values in M then 10: S ij = Q i K j ⊙ M ij 11: mij = rowmax(S ij ) 12: Pij = exp(S ij -mij ) (pointwise) 13: lij = rowsum( Pij ) 14: m new i = max(m i , mij ) ∈ R M , l new i = e m i -m new i l i + e mij -m new i li,j ∈ R M 15: Write O i ← diag(l new i ) -1 (diag(l i )e m i -m new i O i + e mij -m new i Pij V j ) to HBM 16: Write l i ← l new i , m i ← m new 
 
 Appendix: In this section, we provide additional details about the model settings and datasets.ListOps: The models we train on ListOPs 20d have 12 transformer encoder layers with a feature dimension of 128, an intermediate size of 512, a learning rate of 3e-4 and a batch size of 200. We split the original training set into 85k samples for training and 5k samples for validation. The test set accuracy is computed on 10k samples. The ListOps code (Nikita & R., 2018) generates many short samples and a few very long ones, occasionally exceeding the 512 token input size of the BERT model. The largest token length that can be generated by ListOps 20d is 5 20 , with a maximum tree depth of 20 and the maximum number of operands per operator being 5. We exclude such rare samples from the dataset for our experiments. We observe that our conclusions regarding the benefits of sparse structure-aware attention patterns also hold on ListOps data generated with a smaller maximum tree depth of 10d and 5d. As shown in Fig. 11 , HDT is capable of correctly predicting long ListOPs sequences, whereas BERT, Longformer and HAT exhibit lower performance for long samples. HDT's performance on ListOPs is stable across varying input lengths. Model Settings: We choose a typical setting for our model architecture, with d model = 768, d k = 64 with 12 heads, and an intermediate hidden size of 3,072. For HDT-E, we use 12 encoder stacks, and for HDT-ED we use 6 encoder stacks and 6 decoder stacks, resulting in 109M and 124M parameters, respectively. For fine-tuning, we use a constant learning rate of 5e-5, batch size 32 (by accumulating gradients of 8 mini-batches with mini-batch size 4), and a dropout rate of 0.1 for all tasks, except for summarization tasks for which we use a larger learning rate of 1e-3 which we empirically find better suited for all models in our experiments. We use AdamW (Loshchilov & Hutter, 2019) as the optimizer for both pre-training and fine-tuning.The SciRepEval benchmark we use in this work is a subset of the original dataset (Singh et al., 2023) with additional full-text data from unarXive (Saier et al., 2023) . We remove samples that are not in unarXive and drop the tasks that have less than 100 samples left after filtering. This leads to a benchmark spanning 13 tasks, including regression, proximity, and search tasks for scientific document representation evaluation. We show our results on proximity tasks in Table 2 and results on the remaining tasks in Table 6 .Long Document Classification Tasks: Following HAT (Chalkidis et al., 2022) , we validate our model on three additional long document classification tasks. The first task is MIMIC-III (Johnson et al., 2016) , which comprises nearly 50,000 discharge summaries from US hospitals. Each summary is annotated with one or more ICD-9 codes (labels). The model takes a discharge summary as input and outputs the relevant set of first-level ICD-9 codes (19 in total). The second task is ECtHR-LJP (Chalkidis et al., 2021) , containing approximately 11,000 cases from the European Court of Human Rights (ECtHR) public database. For each case, the dataset provides a list of factual paragraphs (facts) from the case description. Each case is mapped to articles of the ECHR that were allegedly violated. The model takes the list of facts as input and outputs the set of allegedly violated articles. Lastly, we include the ECtHR-ARG task (Habernal et al., 2023) , which includes approximately 300 cases from the European Court of Human Rights (ECtHR). For each case, the dataset provides a list of argumentative paragraphs from the case analysis, with spans in each paragraph labeled with one or more of 13 argument types. We follow HAT in re-formulating this task as a sequential paragraph classification task, where each paragraph is labeled with one or more argument types. The model takes the list of paragraphs of a case as input and outputs the set of relevant argument types for each paragraph. Our results in Table 7 show that our model is comparable to HAT on this task.The FacetSum dataset provides fixed 4-class summaries, "Purpose", "Method", "Findings", and "Value", for each document. To align classes to the corresponding sections, we follow Meng et al. (2021) to classify sections into "Introduction", "Method", "Result" and "Conclusion" first by keyword matching and then match the four classes of sections to the four classes of summaries, respectively.The SCROLLS benchmark contains a suite of tasks that require reasoning over long texts. The tasks cover GovReport (Huang et al., 2021) and SummScreenFD (Chen et al., 2022) that are summarization tasks in the domain of government reports and TV shows; QMSum (Zhong et al., 2021) , a query-based summarization task for meeting transcripts; QASPER (Dasigi et al., 2021) , a question answering dataset for scientific papers;https://www.nltk.org/api/nltk.tokenize.htmlhttps://www.scrolls-benchmark.com/leaderboardhttps://www.gutenberg.org/ 


Decision: Paper provides a good demonstration of a technique to improve efficiency (sample, and computational) to the increasingly important problem of long-context modeling, based on document structured and operationalized via sparse attention. Experimental work verifies the approach on realistic workloads. Although very long contexts are not used in the paper, this is a comprehensive demonstration that structured attention based on document properties holds promise.

Paper is relatively clearly written, and reviewers and rebuttal comments indicated it will be of interest to many readers.

Reviews: review: paper_summary: paper_summary: This paper studies hierarchical document representation using different position representations for the hierarchy of the document as section, paragraph and token; the attention mask is adjusted accordingly with a more sparse representation of the whole document. The proposed transformer architecture is compared with existing transformer architectures and some sparse attention transformers approaches as only encoder and encoder-decoder transformers  improves over mathematical reasoning and language tasks.
strength_weakness: strength_weakness: -The experiments do not compare to more recent approaches as Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention, https://openreview.net/pdf?id=UINHuKeWUa
-The writing up of methodology could be improved: 
The description of HPE and hierarchical attention are not very clear. The same index $i$ is used for position and mask, superscript $1$,$2$,$3$ used in Equation 4,5 & 6 are not explained. Fig 10 is referenced for a mask that is sparse in practice, but it’s not clear what is the document structure for that example? Providing the text next to it could help to support this claim as currently it looks like a toy example.
questions: questions: - In table 5, the bold values do not always correspond to the best value, e.g. time 77.84 for HAT is smaller than 79.8 for HDT-E. What is the rationale behind using bold values?
- Could you please add the missing description of notation used in equation 2-7?
paper_summary: paper_summary: The authors present a novel Transformer architecture, namely Hierarchical Document Transformer (HDT), targeting long document processing, which explicitly imposes document structure as an inductive bias in the attention mechanism using special anchor tokens for different document levels (document, sections, and sentences) in the hierarchy. The authors also designed a custom hierarchical attention kernel, influenced by FlashAttention (Dao et al., 2022-2023) to improve efficiency (clock time and memory). They experiment with both Encoder-Only and Encoder-Decoder models comparing to notable baselines, Longformer and HAT for Encoder-Only, and LED for Encoder-Decoder models. The authors showcase that the newly proposed architecture outperforms the baselines on average considering results on SciRepeval proximity (similarity) tasks for Encoder-Only models, and FacetSum and SCROLLs for Encoder-Decoder models. In their efficiency analysis, the authors find their models to be faster in terms of clock time and TFLOPs and more memory-efficient compared to Longformer, on par with HAT.
strength_weakness: strength_weakness: * The benchmarking is mostly limited to document proximity (similarity) tasks and summarization. Long document classification with datasets such as ECtHR, MIMIC, CONTRACT-NLI, or QuALITY,  is not part of the evaluation. 
* The evaluation is mainly focused on the size regime of medium-sized models (~100M params.), which leaves open the question of how HDT would benefit from scaling to larger models with Bs of params, where it could also have a more fair comparison to LongT5 and CoLT5 models.
* The authors did not consider the idea of warm-starting (model-recycling) used by Longformer, HAT, and LongT5, which is another dimension of efficiency concerning saving compute time and resources.
questions: questions: * Why you did not consider expanding the pretraining corpora beyond Arxiv and Patents? From a practical perspective, i.e., people reusing your model for real-life scenarios, this limits the model to fewer domains, while from a benchmarking perspective, it may give your model an extra edge (Arxiv -> SciRepeval) compared to HAT, which was not pre-trained from scratch. I would consider expanding the pertaining corpus with long documents, like those in legal corpora, e.g., the English part of the MultiLegalPile, or general documents >1K by filtering ThePile or C4. Wikipedia documents are also relatively short, close to 200 tokens on average.
* I would like to see a comparison between HDT, HAT, and Longformer on long-form ListOps with Ks of numbers and ops. I think it would be a more interesting comparison instead of BERT which completely lacks any form of anchor tokens or explicit hierarchical attention.
* What are the scores for "Feeds-M" and "High. Infl." refer to in Table 2?
paper_summary: paper_summary: This paper presents an approach to explicitly encode document structure inductive biases in the Transformer architecture's attention mechanism. The approach is interesting and well executed however the motivation for introducing the proposed approach is not clear besides "more efficient" and "sparse". A variety of techniques have been proposed to induce sparsity in the attention mechanism albeit tailored to specific tasks. This paper claims to improve general "long document tasks". To satisfy this claim more experiments should be included on a variety of tasks.
strength_weakness: strength_weakness: The motivation for encoding structure explicitly is not very clear and a variety of other techniques have been proposed for inducing sparsity. The code should be included.
paper_summary: paper_summary: This paper describes a sparse Transformer architecture tailored for structured hierarchical documents (document->section->sentence->token). The new architecture introduces auxiliary anchor tokens (indicating the boundary of document/section/sentence) and redesigns the attention mechanism into a sparse multi-level hierarchy (ensuring each token only attending to its parent, siblings and children).
The authors point out that there are two main benefits of the hierarchical document transformer: exploiting the document structure as an inductive bias (better performance on downstream tasks); and, sparse attention, leading to gains in computational and memory efficiency (faster pre-training). 
Additionally, the authors have created an efficient attention kernel using the Triton library, which can seamlessly replace existing Transformers. The implementation aims to minimize empty block computations by employing a simple heuristic: sorting keys and values according to their hierarchy level first.
To demonstrate the effectiveness of proposed structure-aware attention, the authors conduct a series of experiments on different datasets, including ListOps (mathematical reasoning), SciRepEval, FacetSum, SCROLLS. The authors show that their method outperforms Hierarchical Attention Transformer, which segments long text without sticking to the document structure, and Longformer, which uses sparse attention.
strength_weakness: strength_weakness: * Although the proposed architecture outperforms baselines (i.e., LED base) of similar scale, there is actually a big gap between the reported results and SOTA results (even some of them are really LLMs, e.g., BART-large SLED) on Scrolls, probably on other datasets as well.
* Considering there are Long-context LLMs (e.g., million tokens) available, the maximum sequence length (8192 for pre-training) considered in this paper may not be very exiting.
questions: questions: * Page 9, Efficiency Analysis `12-layer models models` -> `12-layer models`
* Figure 6 and the corresponding text description is Section 4.1 is difficult to understand. Also it is not very clear how these experiments relate to the paper's contribution.
* Figure 1 (b): is this training loss or loss calculated on a hold-out dev set?
* Not sure whether Figure 1 (a) and Figure 2 are duplicated, it seems Figure 2 is a more detailed version or version with example?
* Suggest adding full name of SRAM when you first mention it.

Comments: comment: Thanks for the response. I believe with the changes will improve the final. I keep my original rating.
Thanks for your thorough response. I've read the response. I think it will be great to include all additional experiments in the final version of your article, including results with MIMIC, if possible. I keep my original rating.
I've read the response, I keep my original rating.
Just acknowledging I have read the response; I keep my original rating.