HDT: Hierarchical Document Transformer
Authors: Haoyu He
Markus Flicke
Jan Buchmann
Iryna Gurevych
Andreas Geiger
Venue: COLM (2024)
Keywords: compute & memory efficient Transformer, sparse attention, encoder-only, encoder-decoder, long-text Transformer
TL;DR: A compute efficient Transformer using hierarchical sparse attention.

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.

Published as a conference paper at COLM 2024
HDT: Hierarchical Document Transformer
Haoyu Hel.2; Markus Flicke Jan Buchmann Iryna Gurevych?, Andreas Geiger
1 University of Tübingen
Tübingen AT Center
3 Technical University of Darmstadt and Hessian Center for Al (hessian.Al) Thaoyu.he he, markus. flicke, a. geiger (@uni-tuebinne baan.buchaan iryna. www.eggueeeuuuuuooegee ttaddhrstadddddt.....
Abstract
In this paper, we propose the Hierarchical Document Transformer (HDT), a novel sparse Transformer architecture tailored for structured hierarchi- cal documents. Such documents are extremely important in numerous domains, including science, law medicine. However, most existing solutions are inefficient and fail to make use of the structure inherent to doc- unents. 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 compu- tational and memory efficiency while exploiting the document structure as an inductive bias. We address the technical challenge of implementing HDT's cample-depended
hierarchical attention pattern by developing a novel sparse attention kernel that considers the hierarchical structure of documents. As demonstrated by our experiments, utilizing structural in- formation present in documents leads to faster convergence, higher sample efficiency and better performance on downstream tasks.
Introduction
Many natural l 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 ha
architectures Vaswani et al. (2017) have established themselves as the dominant modeling paradigm, demonstrating state-of-the
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 consumer GPUs.
At the same time, most documents are naturally structured: Words form sentences, sentences form sections and sections form documents. Surprisingly, this structure 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 perfor- mance 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 consumer hardware.
To test this hypothesis, we propose the Hierarchical Document Transformer (HDT), 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. la. Second, redesign the
"Fqual contribution.

---

Published as a conference paper at COLM 2024
HIDT
10
Ours
Longformer
FFN
Ours (converged]
Longformer (converged)
Hierarchical Attention
MLM
Document
Sectoo
Anchor Taken
Regular Token
10
25
Hierarchical Positional Encoding
Training Time (hours)
(a) Model Architecture
(b) Convergence Speed
Figure 1: (a) We propose a sparse attention kernel that considers the hierarchical structure docu- ments. 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-traine treiniriinng (b) and better performance on downstream tasks. We use the held-out validation set in (b) to calculate the MLM loss.
attention pattern of Transformer block into 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 few HDT layers establishes communication between all tokens.
However, implementing our model required solving a key technical challenge: Existing sparse Transformer architectures like LongFormer Beltagy et al. (2020) assume fixed sparsity pattern for all inputs in a mini-batch which can be implemented using standard libraries. In contrast, HDT imposes different sparsity pattern for each sample within mini-batch as each document is structured differently. Towards this goal, we developed novel, flexible and efficient attention kernel based on the Triton (Tillet al., 2019) library. We pre-train HDT encoder models using Masked Language Modeling (MLM) and HDT encoder- decoder models using UL2 (Tay et al., 2023b) on three document datasets: arXiv (Saier et al., 2023), HUPD Patents (Suzgun et al., 2023) and Wikipedia. We demonstrate improved pre- training convergence rates (see Fig. 1b) as well as better downstream task performance on several proximity, summarization, QA and NLI tasks. Our code and data are available at https:/ github. wtsttuu..com/ee httmm//uuoooooossssiooo///ooe
2
Related Work
This study combines two lines of work that try improve the performance of bi-directional Transformers in handling long-document inputs: structural modeling and efficient attention.
Structural modeling approaches exploit the hierarchical organization of documents into sections, subsections, paragraphs, sentences, etc. They can be divided into two categories Buchmann et al. (2024): (1) In structure infusion, structural information is added to the Trans- former input, e.g., via special tokens (Aghajanyan et al., 2022; Liu et al., 2022; Buchmann et al., 2024), position embeddings (Bai al., 2021; Cao Wang, 2022), attention masks (Wang et al., 2019; Liu et al., 2021; Hong et al., 2022; Wu et al., 2021b) or fusing self-attention layers with GNN layers Sachan et al. (2021); Ahmed et al. (2019). While observing performance improvements, these works do not improve efficiency. (2) Hierarchical processing employs architectures that first contextualize tokens on a local (e.g., sentence) level, aggregate the local representations (e.g., in a "ICLS]" token) and then contextualize the aggregates one or several higher levels (e.g., paragraph or section, Yang et al. 2016; Chalkidis et al. 2019; Yang 2019; Ruan et al. 2022; Zhang et al. 2022; Dai et al. 2022). Local and higher-level
2

---

Published as a conference paper at COLM 2024
[DOC]
[8EC]
[SENT]
T1
[SEC]
[SEC]
+ Section Token
[SEC]
T2
[SENT]
[SENT]
[SENT]
Sentence Token
T3
T1
T2
T4
T5
T6
+ Regular Taken
T5
1o0 (BEO) BENTT
TI
61 IEaal
78
Figure 2: Hierarchical Document Decomposition. Left: Tree representation of document. Tokens within the same box attend to each other. Tokens that do not share a box attend to each other only indirectly (e.g, 11 and via the sentence and section tokens). Right: Sparse attention matrix.
contextualization 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 single layer using an efficient hierarchical attention pattern.
Due 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 work aims at speeding up attention computation by considering & 2 (Dao et al., 2022; Dao, 2023) propose
GPU hardware characteristics. PlashAttention tiling with block-wise attention to reduce memory IOs and compute attention blocks in parallel in Static Random Access Memory (SRAM). Matteo al. (2023) extends the idea of IlashAttention to key /query dropping and hashing-ban -base attention that was originally proposed by Kitaev et al. (2020). Our work builds upon both branches: We propose dynamic sample-dependent sparse attention pattern which exploits the structure of text documents. For efficiency, we follow Dao al. (2022) and implement this pattern as customized memory-award kernel.
3 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 encoder-decoder HDT architectures which are used in our experiments in Section
3.1 Standard Transformer
The original transformer encoder model by Vaswani et al. (2017) composed of Multi-Head Self-Attention (MHSA) layers that are interleaved with shallow feed-forward networks and residual connections. Because the attention pattern identical across multiple heads in an attention layer, for simplicity of notation, we represent queries, keys and values for single attention head as Q K, V Rxx,,, where is the sequence length and d, the head dimension. All MHSA layers in the model, as well as the embedding layers, produce outputs of dimension dmodel- We further denote the number of heads as it, hence dmodel de h. A

---

Published as a conference paper at COLM 2024
Tree Level
[DOc]
[SEC]
(SENT)
T1
[SEC]
(SENT)
T2
(SENT)
T3
T4
Senttence
8
1
1
1
2
2
2
3 Token
0 1
3
Figure 3: Hierarchical Positional Encoding. We represent the position of each token in the hierarchy with one linear index p' per hierarchy level yielding an index vector p (p' ....,p')' Above, we show an example with = 3 levels. Note that level 0 (document) does not require an index. Each index in is passed through sinusoidal encoding functions which are summed over all levels to form the final encoding vector according Eq. (2).
self-attention head computes its output O E RRux as follows:
QK
A
= softmax (A)
== SV
(1)
where the softmax operator is applied rowe The standard attention block has O(n2) time and memory complexity where и is the input length. It further does not explicitly take the structure of documents into account.
3.2 Hierarchical Document Transformer
Most documents organized into structural constituents like sections, paragraphs, sen- tences, bulleted lists, figures, and footnotes. This structure represented 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 document hierarchy with three levels: tokens, sentences and sections. More specifically, we split document with n tokens (E,, .., E Sections &, are split into sentences E; = S2.,, ( which are split into sequences regular tokens S = (ti 12, fsss))
D (4, 1,,,,,,,,, In) into 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) 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 the beginning of the document as illustrated in Fig.
Hierarchical Positional Encoding: We extend the sinusoidal position encoding to model hierarchy levels. To inform each token about its position within the hierarchy, we assign it one linear index p' per hierarchy level as illustrated in Fig. 3, yielding an index vector p= (p),... ...., p4)T Each index in p is passed through set of standard sinusoidal encoding functions which are summed over all levels to form the final hierarchical positional encoding (HPE) vector at p:
L [sin(axp))
if = 2k
HPE(p,1) Σ
where Wx
(2)
cos(uu)p) p') if 2k
20 4model
Hierarchical Attention: To model document structure, impose document-specific attention pattern by modifying the standard Transformer attention in Eq. (1) as follows
OK
A
S softmax(A O 1l)
O SV
(3)
Vdk
with (Aj ④ 1M;) = Aij if Mj = and - if Mij = 0. For clarity of notation, we use p/,
to represent the position index for hierarchy level I of the i'th token and the j'th token
4

---

Published as a conference paper at COLM 2024
High Bandwidth Memory (HBM)
Original Order
(Q.
Stor
Vo Kays and Values
Quertes
SRAM
Q
Skip
Skip
Skip
Copy to
SRAM
Skip
Skip
Skip
Figure 4: Hierarchial Attention Kernel. We copy queries, keys and values block-wise to SRAM for fast attention computation using fused kernel. To increase the number of empty blocks that can be skipped, we sort keys and values based on their hierarchy level. Larger examples are shown in Fig. 10.
f,. respectively. The attention mask M E {0, [0,1] 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 3-level document structure, we first compute the attention mask for each hierarchy level separately
Mpp [r? = 0] [p] 0]
(4)
MSDC [p? = 0] [p] 0] [p] = p/l
(5)
MSSNTT [p! = p/] Ir? = přl
(6)
where [-] denotes the Iverson bracket which evaluates if the argument is true and otherwise. Finally, we perform an OR operation to obtain the full attention mask M:
M
M DOC
SBC
① SENT
Note that M is highly sparse in practice (see Fig. 10 for an example) and hence reduces theoretical complexity from O(n²) to ( where is the length of the longest sentence in the document. Thus, computational savings are largest for long documents for which s < 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 be computed efficiently in SRAM. This 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 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 simple heuristic which is illustrated with an example Fig. 4. Specifically, before copying keys and values to SRAM, we first sort them based on their hierarchy level from / = 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 Qi keys Kj and values V, of block (i,f) to SRAM and apply block attention. We process non-empty blocks in parallel, skipping empty ones. Finally, we write the result O back to HBM using the online softmax

---

Published as a conference paper at COLM 2024
SM
SM (SUM-MOD)
0.5
MED
Vadidativn Accuracy
MED
MAX
0.4
MAX
BERT
HDT
0.-
Lonffomar
[SUM-MOD 3 IMED 2 3J 9 3]] Ground truth:
11AT
MED
MAX
50000
100000
150000
200000
Training Step
Figure 5: ListOps Sample
Figure 6: ListOps Attention
Figure 7: ListOps Convergence
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.
Hierarchical Encoder and Decoder Stacks: 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. la. The first sub-layer performs hierarchical multi-head self-attention as described above, and the second is a simple position-wise fully connected ford-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 stack of N identical standard Transformer decoder blocks, including sub-layer of causal multi-head attention that prevents tokens from attending to subsequent positions, sub-layer to perform multi-head cross attention wrt. the output of the encoder stack, and a standard feed-forwart 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 mathe- matical 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 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.
4.1 Mathematical Reasoning Tasks
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 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. shows the different types of attention patterns we use. The combination of redggneen+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 "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 increase
Model
Acc.
BERT
75,9
HDT, g lb
79.7
HDTg+b HDT
85.6
86.2
Table 1: ListOps Acc.

---

Published as a conference paper at COLM 2024
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 ListOps even without using any positional embedding the operators in the ListOps task are invariant to the positions of their operands.
4.2 Language Tasks
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 GPUs (~50k steps) using a batch size of 128. Compared 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 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 (Rehúřek & Sojka, 2010). In total, our corpus includes over 12M long documents with extracted structural information. Our word tokenizer utilizes vocabulary size of 32,768 tokens, trained with Byte Pair Encoding (BPE) (Sennrich et al., 2016). For splitting sentences, we use the NLTK sentence tokenizer'. We report results of HDT-E pre-trained with and without arXiv data 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.
Baselines: 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/ memorr usage for processing long text. HDT-E differs from HAT as follows: HD) 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 implementa- ba
tion. As the sparse attention encoder-decoder variant baseline, we choose Longformer (Beltagy et al., 2020) and its Senfffomer--Enoder---eeeeee
(LED) which are effective on various long-text tasks previous works (Tay et al., 2021; Dasigi 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.
Pre-training & Fine-tuning: We pre-train the encoder-only model HDT-E and the encoder- decoder model HDT-ED from scratch on our training corpus. All models are pre- pre-trained using an input length of 8,192 tokens and a mini-batch size of 128 (via gradient accumula- tion). HDT-E is pre-trained on the standard Masked Language Modeling (MLM) objective with mask ratio of 15%. We also pre-train Longformer from scratch 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 unified pre-training paradigm with range of denoising tasks. We fine-tune all models on the downstream tasks using the default settings, see Section A.2.2 for details.
Results on SciRepEval (Encoder Models): SciRepEval (Singh et al., 2023) is a scientific document representation benchmark containing various classification, regression, and proximity tasks. Unfortunately, the original SciRepEval dataset comprises only titles and abstracts as for many papers in the original benchmark the full text is not publicly available.
http:://www.ailkooog//////kkkkkk
7

---

Published as a conference paper at COLM 2024
Full Text
SciDocs
High.
Model
Feeds-M
Avg.
Cite
CoCite
CoView CoRead
Infl.
Pretrained Only
SciBERTbase
53.75
66.73
66.37
53.20
63.18
40.80
57.34
Longf ormer
56.64
71.92
71.51
61.57
63.66
43.82
61.52
HAT
60.14
75.55
73.62
67.65
65.02
45.81
64.63
HDT-E
62.50
78.51
75.69
72.12
65.19
43.60
66.27
HDT-E (-arXiv)
59.03
76.05
72.85
71.71
65.41
43.69
64.79
Pretrained Finetuned with Contrastive Learning
SciNCL @684k
64.77
81.67
78.55
77.48
70.22
48.66
70.23
SciNCL @19k
62.56
82.29
77.84
75.84
67.11
46.23
68.65
Longf ormer @19k
61.75
79.87
78.20
74.25
67.80
43,85
67.62
HAT @19k
63.46
81.24
79.43
75.76
69.31
47.37
69.42
HDT-E @19k
64.23
82.44
78.95
77.09
71.22
49.37
70.55
HDT-E @19k (-arXiv)
63.34
82.18
79.06
76.78
70.64
48.95
70.16
Table 2: Results SvilepEval Proximity Tasks. Top: Models pre-trained with MLM without fine- tuning. Bottom: Models pre-trained with MLM and fine-tuned using SciNCL's contrastive learning objective. Full text documents are available only for subset 19k training triplets. For reference, we also report the results of the original SciNCL model which is trained on all 684k title+abstract triplets. We also report HDT-E pre-trained without arXiv data (-arXiv) to study the impact of scientific documents pre-training data to model's performance the SciRepEval tasks which in the scientific domain. All numbers are mean average precision. SciBERT and SciNCL use only title and abstract as input.
To adapt SvilepEval for long document models, we hence consider 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 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 ("Pretrained Only") shows the mean average precision for models pre-trained MLM. Our results demonstrate that both, pre-training 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 al., 2020) and RoBERTa (Liu et al., 2019) (respectively) for best performance. This result underscores the sample efficiency of our structure-aware HDT model.
of candidate papers by their relatedness
Currently, the state-of-the-art performance on SciRepEval is reported by SciNCL (Ostendorff et al., 2022) which finetunes pre-trained SciBERT (Beltagy et al., 2019) model using contrastive learning objective. When comparing in this setting ("Pretrained + Finetuned with Contrastive Learning"), we observe that HDT-E again outperforms all baselines, despite with 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 scientific documents for pre-training, we exclude the entire arXiv corpus from our pre- training cltasett. We observed that our model's performance nearly matches HAT in the pre-training-on)) setting and surpasses HAT the contrastive learning setting, even when using only out-of-domain data for pre-training.
publicly available. In addition, to test our model's reliance
Results on FacetSum (Encoder-Decoder Models): 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, design
8

---

Published as a conference paper at COLM 2024
Model
Purpose
Method
Findings
Value
Purpose-ZS
Method-ZS
baSe
39.51
19.31
19.22
24.26
18.54
14.12
HDT-ED- SEC]
30.70
17.65
17.15
18.90
15.01
12.67
[SENT]
34.29
19.72
18.38
20.53
19.67
13.94
+ tokens
40.60
22.21
22.11
22.11
21.75
15.43
Table 3: Results on FacetSum Summarization Task. Following the original paper, we report ROUGE- as the metric here. For HDT-ED-(SEC) , the decoder cross-attends only to the section anchor token [SEC]. We observe that even when attending only to the anchor tokens (+ (+[SENT), our model is on par with LED, where the decoder attends to all tokens of the section, demonstrating the expressiveness of the learned intermediate representation of anchor tokens. When attending to additional regular tokens (+tokens), our model outperforms LED. We also report zero-shot (ZS) performance for "Purpose" and " "Method", training only on "Findings" and "Value".
RovRep ROUGE-1/2/L
SumSer ROUGE-1/2/L ROUGE-1/2/L
QMSum
Qspr
Nrtv
QALT
CNLI
Avg
Model
Fi
F1
EM-T/H
EM
Score
LEDpase
56.2/26.6/28.. 28.8 24.2/4.5/554 49.8/22.2/25.8
25.1/6.7/18.8
26.6
18.5
25.8/25.4
71.5
29.16
HDT-ED
00.8/7.1/88.6 28.3/6.7/18.7
33.1
14.2 29.4/26.4
81.4
31.41
Table 4: Results on the SCROLLS Summarization, QA and NLI Benchmark. We compare HDT- (LED) on the official SCROLLS benchmark without document strucfure. We choose LED as baseline as it has comparable number of parameters (162M) to HDT-ED (124M). We remark that neither model is competitive with state-of-the- art billion-parameter models such as CoLI5 XL (score 43.5) which are trained large GPU clusters.
ED (pre-trained for 12 GPU days) to Sengffomer-Encoder---r
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 LEDbase 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 LEDbase 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.
Results on SCROLLS (Encoder-Decoder Models): 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 compused of a fixed number of 32 sentences for hierarchical attention, while sentences are extracted usual. We fine-tune HDT-ED on the 7 tasks in SCROLLS separately and submit the test set predictions to the public leaderboard?. Table 4 show our results. Compared to LEDbase with 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 CoL15 XL which are trained for many epochs on large industrial GPU clusters.
Denoting the length (number of tokens) of the longest sentence in document as s, the theoretical complexity of HDT attention is O(n x s). Fig. 8 compares GPU runtime and memory usage of different attention layers, including standard dense attention, block-sparse Flash (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 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.
Efficiency Analysis:

---

Published as a conference paper at COLM 2024
Standard Antenton
Model
Longformer
HAT
HDT-E
60
Longlormer
Complex.
(n)
O(nk)
(n))
Block Spanse Maas--ittttiio
Time per bateh (ms)
50
HOT Attention
#Params
148.66 M
152.73 M
108.99 M
Time (ms)
178.82 6.8
77.84 2:30
79.8 2.6
TFLOPS
5.29 0.19
8.95 = 0.26
8.99 0.34
9636 MB
13146 MB
Memory 11.25 GB
5.3 GB
5.85 GB
20-
1988
Table 5: Runtime and Memory consumption of several long-document models with 12 layers. We report complex- ity, parameters, inference time, throughput, and mem- ory usage using context length и 409 and mini-batch size 1. Here,
10
8k
16k
32k
Ingut Length
512 is the Longformer window size, k 128 the fixed HAT segment length. is the length the longest sentence in the document.
Figure Runtime and Memory con- sumption of single attention layer.
Conclusion
We presented Hierarchical Document Transformer (HDT), 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 represen- tations 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 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 they do for LLMs, and if hierarchical language models hold similar potential for unlocking emergent abilities.
Acknowledgments
Andreas Geiger is member of the Machine Learning Cluster of Excellence, funded by the Deutsche Forschungssemeinschaat (DFG, German Research Foundation) under Germany's Excellence Strategy EXC number 2064/1 Project number 390727645. Jan Buchmann was funded by the European Union (ERC, InterText, 101054961). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council. Neither the European Union nor the granting authority can be held responsible for them.
10

---

Published as a conference paper at COLM 2024
References
Armen Aghajanyan, Dmytro Okhonko, Mike Lewis, Mandar Joshi, Hu Xu, Gargi Ghosh, and Luke Zettlemoyer. HTLM: Hyper-text pre-training and prompting of language models. In Proc. of the International Conf. on Learning Representations (ICLR), 2022. 2
Mahtab Ahmed, Muhammad Rifayat Samee, and Robert E. Mercer. You only need attention to traverse trees. In Anmual Meeting of the Association for Computational Linguistics (ACL), 2019. 2
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. ETC: encoding long and structured inputs in transformers. In Conference on Empirical Methods in Natural
Language Processing (EMNLP), 2020. 3
Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David C. Uthus, Mandy Guo, James Lee-Thory Thorp, Yi Tay, Yun-Hsuan Sung, and Sumit Sanghai. Colt5: Faster long-range transformers with conditional computation. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2023. 3
Cevdet Aykanat, Ali Pinar, and Ümit V. Çatalyürek. Permuting sparse rectangular matrices into block-diagon form. SIAM Journal on Scientific Computing (SICS), 2566) 2004. 6
He Bai, Peng Shi, Jimmy Lin, Yuqing Xie, Luchen Tan, Kun Xiong, Wen Gao, and Ming Li. Segatron: Segment-aware transformer for language modeling and understanding. In Proc. of the Conf. on Artificial Intelligence (AAAI), 2021. 2
Iz Beltagy, Kyle Lo, and Arman Cohan. Scibert: A pretrained language model for scientific text. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2019. 7, 8
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document trans- former. ar Xixioorr 2004.05150, 2020. 2, 2,3,6,7,11 6,7,15
Jan Buchmann, Max Eichler, Jan-Micha Bodensohn, Ilia Kuznetsov, and Iryna Gurevych. Document structure in long document transformers. In Conference of the European Chapter of the Association for Computational Linguistics (EACL), 2024. 2
Shuyang Cao and Lu Wang. HIBRIDS: attention with hierarchical biases for structure-aware long document summarization. In Annual Meeting of the Association for Computational Linguistics (ACL), 2022. 2
Ilias Chalkidis, lon Androutsopoulos, and Nikolaos Aletras. Neural legal judgment predic- tion in English. In Annual Meeting of the Association for Computational Linguistics (ACL), 2019. 2
Ilias Chalkidis, Manos Fergadiotis, Dimitrios Isarapataanss, Nikolaos Aletras, lon Androut- sopoulos, and Prodromos Malakasiotis. Paragraph-level rationale extraction through regularization: A case study on European court of human rights cases. In Annual Meeting of the Association for Computational Linguistics (ACL), pp. 226-241, 2021. 19
Ilias Chalkidis, Xiang Dai, Manos Fergadiotis, Prodromos Malakasiotis, and Desmond Elliott. An exploration of hierarchical attention transformers for efficient long document classification. ar arxioorr, 2210.05529, 2022. 2, 6,7, 19
Mingda Chen, Zewei Chu, Sam Wiseman, and Kevin Gimpel. Summscreen: A dataset for ab- stractive screenplay summarization. In Annual Meeting of the Association for Computational Linguistics (ACL), 2022. 19
Xiang Dai, Ilias Chalkidis, Sune Darkner, and Desmond Elliott. Revisiting transformer-based models for long document classification. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2022. 2
11

---

Published as a conference paper at COLM 2024
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. ar arXiv.org, 2307.08691, 2023. 3,5
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022. 3, 5, 16, 17
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. A dataset of information-seekinn questions and answers anchored in research papers. In North American Chapter of the Association for Computational Linguistics (NAACL), 2021. 7, 19, 20
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In North American Chapter of the Association for Computational Linguistics (NAACL), 2019.
Michael C. Ferris and Jeffrey D. Horn. Partitioning mathematical programs for parallel solution. Mathematical Programming (MP), 80:35-61, 1998. 6
Jonas Geiping and Tom Goldstein. Cramming: Training language model on single GPU in one day. In Proc. of the International Conf. on Machine learning (ICML), 2023. 7
John T. Guthrie, Tracy Britten, and K. Georgene Barker. Roles of Document Structure, Cog- nitive Strategy, and Awareness in Searching for Information. Reading Research Quarterly, 26(3), 1991.
Ivan Habernal, Daniel Faber, Nicola Recchia, Sebastian Bretthauer, Iryna Gurevych, Indra Spiecker genannt Döhmann, and Christoph Burchard. Mining legal arguments in court decisions. Artificial Intelligence and Law, 2023. 19
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proc. JEEE Conf. on Computer Vision and Pattern Recognition (CVPR), 2016.
Bruce Hendrickson and Tamara G. Kolda. Partitioning rectangular and structurally unsym- metric sparse matrices for parallel processing. SIAM Journal on Scientific Computing (SICS), 21(6):2048-200 2000.
Giwon Hong, Jeonghwan Kim, Junmo Kang, and Sung-Hyon Myaeng. Graph-induced transformers for efficient multi-hop question answering. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2022. 2
Luyang Huang, Shuyang Cao, Nikolaus Nova Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. In North American Chapter of the Association for Computational Linguistics (NAACL), 2021. 19, 20
Alistair E. W. Johnson, Tom J. Pollard, Lu Shen, Li-wei H. Lehman, Mengling Feng, Moham- mad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G. Mark. Mimic-iii, a freely accessible critical care database. Nature, 3, 2016. 19
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: Proc. of the International Conf. on Learning Representations (ICLR), 2020. 3
The efficient transformer. In
Tomás Kociský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativega reading comprehension challenge. Transactions of the Association for Computational Linguistics (TACL), 6:317-328, 2018. 20
Yuta Koreeda and Christopher D. Manning. Contractnli: A dataset for lecument-level natural language inference for contracts. In Conference on Empirical Methods in Natural
Language Processing (EMNLP), 2021. 20
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: denoising sequence-to- sequence pre-training for natural language generation, translation, and comprehension. In Annual Meeting of the Association for Compulational Linguistics (ACI.), 2020. 8
12

---

Published as a conference paper at COLM 2024
Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. ar arXio.org, 2307.03172, 2023. 1
Yang Liu, Jiaxiang Liu, Li Chen, Yuxiang Lu, Shikun Feng, Zhida Feng, Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. ERNIE-SPARSE: learning hierarchical efficient transformer through regularized self-attention. ARXIV, abs 2ttp//220....2222 2022.
Ye Liu, Jian-Guo Zhang, Yao Wan, Congying Xia, Lifang He, and Philip S. Yu. HETFORMER: heterogeneous transformer with sparse attention for long-tex, extractive summarization. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2021. 2
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. ARXIV, (bs//197.11692, 2019. 8
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proc. of the International Conf. on Learning Representations (ICLR), 2019. 19
Pagliardini Matteo, Paliotta Daniele, Jaggi Martin, and Fleuret François. Faster causal attention over large sequences through sparse flash attention. arXiv.org, 2306.01160, 2023.
Rui Meng, Khushboo Thaker, Lei Zhang, Yue Dong, Xingdi Yuan, Tong Wang, and Daqing He. Bringing structure into summaries: a faceted summarization dataset for long scientific documents. In Annual Meeting of the Association for Computational Linguistics (ACL), 2021. 7,9,15 19
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXivoorg 1805.02867, 2018.
Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher Manning. Grokking of hierarchical structure in vanilla transformers. In Annual Meeting of the Association for Computational Linguistics (ACL), 2023.
Nangia Nikita and Bowman Samuel R. ListOps: diagnostic dataset for latent tree learning. In North American Chapter of the Association for Compulational Linguistics (NAACL), 2018. 18
Malte Ostendorff, Nils Rethmeier, Isabelle Augenstein, Bela Gipp, and Georg Rehm. Neigh- borhood contrastive learning for scientific document representations with citation embed- dings. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2022.
Richard Yuanzhe Pang, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma, Jana Thompson, He He, and Samuel R. Bowman. Quality: Question answering with long input texts, yes! In North American Chapter of the Association for Computational Linguistics (NAACL), 2022. 20
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research (JMLR), 11 2020. 6
Radim Rehūřek and Petr Sojka. Software framework for topic modelling with large corpora. In International Conference on Language Resources and Evaluation (LREC) Workshops, 2010. 7
Child Rewon, Gray Scott, Radford Alec, and Sutskever Ilya. Generating long sequences with sparse transformers. ar arXiv.org 2019. 3
Qian Ruan, Malte Ostendorff, and Georg Rehm. Histruct+: Improving extractive text summarization with hierarchical structure for Computational Linguistics (ACL., 2022. 2
struuccur information. In Annual Meeting of the Association
13

---

Published as a conference paper at COLM 2024
Devendra Sachan, Yuhao Zhang, Peng Qi, and William L. Hamilton. Do syntax trees help pre-trained transformers extract information? In Conference of the European Chapter of the Association for Computational Linguistics (EACL), 2021. 2
Tarek Saier, Johan Krause, and Michael Färber. unarxive 2022: All arxiv publications pre- processed for nlp, including structured full-text and citation network. In ACM/IEEE Joint Conference on Digital Libraries, JCDL (JCDL), 2023. 2, 7, 8, 19
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Annual Meeting of the Association for Computational Linguistics (ACL), 2016. 7
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: standardized comparison over long language sequences. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2022. 7
Amanpreet Singh, Mike D' Arcy, Arman Cohan, Doug Downey, and Sergey Feldman. Scirepe- val: A multi-format benchmark for scientific document representations. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2023. 7, 19
Mirac Suzgun, Luke Melas-Kyriazi, Suproteem K Sarkar, Scott Kominers, and Stuart Shieber. The harvard USPTO patent dataset: A large-scale, well-structured, and multi-purpose corpus of patent applications. In Proc. of the Neural Information Processing Systems (NeurIPS) Track on Datasets and Benchmarks, 2023. 2,7
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena A benchmark for efficient transformers. In Proc. of the International Conf. on Learning Representations (ICLR), 2021. 7
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, ACM Computing Surveys, 15:1:129) 2023a.
Yi Tay, Mostafa Dehghani, Vinh Q. Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Huaixiu Steven Zheng, Denny Zhou, Neil Houlsby, and Donald Metzler. UL2: unifying language learning paradigms. In Proc. of the International Conf. on Learning Representations (ICLR), 2023b. 2,7
Barbara M. Taylor and Richard W. Beach. The Effects of Text Structure Instruction on Middle-Grade Students' Comprehension and Production of Expository Text. Reading Research Quarterly, 19(2):134-146, 1984. 4
Philippe Tillet, Saiang-Tsung Kung, and David D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (MAPL), 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NIPS), pp. 5998-6008, 2017. 1,3
Yaushian Wang, Hung-Yi Lee, and Yun-Nung Chen. Tree transformer: Integrating tree structures into self-attention. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2019.
Chuhan Wu, Fangzhao Wu, Tao Qi, and Yongfeng Huang. Hi-transformer: Hierarchical interactive transformer for efficient and effective long document modeling. In International Joint Conference on Natural Language Processing (IJCNLP), 2021a.
Hongqiu Wu, Hai Zhao, and Min Zhang. Code summarization with structure-induced transformer. Annual Meeting of the Association for Computational Linguistics (ACL), 2021b. 2
14

---

Published as a conference paper at COLM 2024
Liu Yang.
Fine-tune BERT for extractive summarization. 2019.
arXiv.org,
Zichao Yang, Diyi Yang, Chris Dyer, Xiaodong He, Alex Smola, and Eduard Hovy. Hier- archical attention networks for document classification.. In North American Chapter of the Association for Computational Linguistics (NAACL), 2016.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Haopeng Zhang, Xiao Liu, and Jiawei Zhang. HEGEL: : hypergraph transformer for long doc- ument summarization. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2022.
Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadallah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, and Dragomir R. Radev. Qmsum: A new benchmark for query-based multi-domain meeting summarization. In North American Chapter of the Association for Computational Linguistics (NAACL), 2021. 19
A Appendix
A.1 Hierarchical Attention Kernel
poc p0r1y
1000
p001
71
1200
1o0
13
091
JENT
poo pequenre pecqpeent (a9T)
(a) Longformer's fixed-window attention
(b) Our dynamic hierarchical attention
Figure 9: (a) The Longformer (Beltagy al., 2020) sparse attention pattern is identical for all samples mini-batch. (b) contrast, the proposed dynamic hierarchical attention pattern considers the document structure and hence is different for each sample in mini-batch.
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 im- plementation. 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 sample- dependent attention patterns. Fig. 9 provides visual illustration comparing our dynamic attention pattern to the fixed attention pattern of Longformer (Beltagy et al., 2020), assuming mini-batch size of for clarity. Notably, within mini-batch containing multiple samples, the attention pattern applied 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 different hierarchies and attending to one another, the attention devoted to these anchor tokens can be notably fine-grained and sparse within
15

---

Published as a conference paper at COLM 2024
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 FlashAttention Dao et al. (2022). Our algorithm for forward pass of the attention layer provided in Algorithm 1. This algorithm is executed in parallel for all elements of the Cartesian product B x H, where represents the batch size and H denotes the number of dz).
and extends
heads, assuming input tensors of shape (B H
Our primary adjustment to the FlashAttention algorithm is in line 2 of Algorithml, where we implement 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 grounded the observation that the computational ineffi- ciency 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 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 E R Rx query block size M, key-value block size N, softmax statistics vectors m E R R and E R output tensor Ouxd,
1: Partition Q into T, I
M blocks Qi...., Qr,
2: Sort K, V according to the hierarchy level from t I..
3: Partition K, V into T. I blocks K,,..., KT,, and Vi,..., VT.
4: for 1 < < T, do
5: Load Q; O; m; from HBM to SRAM
6: for S < Tr do
7:
Load K, Vj from HBM SRAM
8:
On chip, compute mask M according Equation (7)
9:
if Non-zero values in M then
S,j K,Ky @ Mij
11:
mi = (Swmax(Si))
12:
jj exp(S,) (S m,)) (pointwise)
l rowsum (P₁;)
meew (ax)) R new eh mecw
Mij-nn T1, E R
15:
Write O, diag(! (I(ee) 1 (ddiag())ee):-"""")
O Mij m
V,) to HBM
16:
Write I; wne m; m.ew to HBM
17:
end if
18:
end for
19:
end for
16

---

Published as a conference paper at COLM 2024
(a) Block-Sparse FlashAttention (Dao et al., 2022)
(b) HDT Attention
(c) Block-Sparse FlashAttention (Dao et al., 2022)
(d) HDT Attention
Figure Attention Pattern (Black) and Processed SRAM Blocks (Grey). Comparison between practical computation of Block-Sparse FlashAttention (Dao et al., 2022) and our HDT attention kernel on the same hierarchical attention patterns. We show the attention mask of the first 1k tokens (~255% of the total document size) of two different documents (row 1+2) in black. The blue grid illustrates the 128 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 the reordering of keys and values (columns) in HDT, anchor tokens are aggregated within adjacent blocks leading to larger number of blocks that can be skipped compared to SSprre FlashAttention (Dao et al., 2022).
17

---

Published as a conference paper at COLM 2024
A.2
Experiments
In this section, we provide additional details about the model settings and datasets.
A.2.1
Mathematical Reasoning Tasks
The models we train on ListOPs 20d have 12 transformer encoder layers with 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 520, with 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 patterns also hold on ListOps data generated with 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.
ListOps:
conclusions regarding the benefits of sparse structure-aware attention
15
Correct
Cornet
14
Longformer
HAT
16
Incomect
Incerrect
14
14
12
12
LagiCunn
20
J02
300
4:00
500
120
23c
300
40F
500
(rput Langhh
I1p4 Lenghh
Figure 11: Stacked histograms of predictions by HDT, BERT, Longformer and HAT ListOPs 20d.
A.2.2 Language Tasks
Model Settings: We choose typical setting for our model architecture, with d 768, d: = 64 with 12 heads, and an intermediate hidden size 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 larger learning rate of 1e-3 which empirically find better suited for all models in
18

---

Published as a conference paper at COLM 2024
SRH
Search
Feedb-1
Peer Review Score
Tweet Mentions Kendall's /
hindex Kendall's 7
Citation Count Kendall's /
Publication Year
Average
nIXG
MAP
Kendall's 7
Kendiall's 7
BETRPAY
56.79
69.2
17.39
TD:19
16.34
15,43
13.37
28.44
Longforeer
56.11
66.46
17.34
12.68
16.65
26.16
31.54 54
BeAMCL
56.05
71.85
HAT
37.01
70.01
14.02
13.18
13.64
19.75
19.8
29.72
HDT-E E (tble) alsttrcc)
56.62
7126
13.79
13.62
15.65
26.34
24.66
31.76
Table Results on remaining SciRepEval tasks. Although our model is not the best across all tasks, the performance is comparably stable different tasks than other models.
Model
MIMIC
ECIHR-LJP
ECIHR-ARG
F1
F1
acc
Longformer
78.7
78.6
66.7
НAT
78.9
79.8
82.6
HDT-E
79.1
79.5
82.9
Table 7: Results on document classification task ECUHR-LJP, ECHHR-ARG, MIMIC We follow the experimental setting of HAT and cover three long document classification tasks and report the test scores.
our experiments. We use AdamW (Loshchilov & Hutter, 2019) as the optimizer for both pre-training and fine-tuning.
The Scill benchmark we use in this work 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 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.
SvilepEval:
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 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, 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 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 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 comparable 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, 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.
FacetSum:
The SCROLLS benchmark contains 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), question answering dataset for scientific papers;
SCROLLS:
19

---

Published as a conference paper at COLM 2024
GovRep
SumScr
QMSum
Qspr
Nrtv
QALT
CNLI
Avg
Model
ROUGE-I/2/L
ROUGE-1/2/L
ROUGE-1/2/L
F1
F1
EM-T/H
EM
CoLTSX,
61.3/32.2/33.8 36A//101/21.7
36.2/12.9/2.22
53.9
31.1
48.1/43.8
88.4
43.51
LongT5X1
61.1/32.3/33.7
35.8/9.6/2111
34.9/11.82355
53.1
29.3
46.0/42.1
88.2
42.53
CoLT5iarge
60.7/31.3/32.9
36.7/10.6/22.0
34.9/11.5/2.11
49.8
27.7
39.9/36.8
88.7
41.04
LED HDT-ED
56.2/26.6/28.8
24.2/45/15.4
25.1/6.7/18.8
26.6
18.5
25.8/25.4
71.5
29.16
49.8/22.2/25.8
30.8/7.1/18.6
28.3/6.7/18.7
33.1
14.2
29.4/26.4
81.4
31.41
Table 8: Results on the SCROLLS benchmark compared to SotA. We evaluate several models on the SCROLLS benchmark. t indicates results reported on the public leaderboard.
GovRep ROUGE-1/2/L
QASPER
Model
F1
Extractive
Abstractive
Yes/No
Unanswer.
Overall
LED
56.2/26.6/28.8
30.96
15.76
70.33
26.21
32.80
HDT-ED 49.89/21.5/22.266
30.57
11.42
67.14
46.45
33.14
struct.
49.42/21.22/24..97
33.12
13.02
64.19
43.17
34.02
Table 9: Effect of using document structure on GovRep and QASPER, Note that the models are evaluated with the original GovRep (Huang et al., 2021) and QASPER (Dasigi et al., 2021) datasets, therefore results in this table might be slightly different from the results in Table 4 for the same task as the authors have cleaned the datasets in SCROLLS.
NarrativeQA (Kociský et al., 2018), a question answering dataset over entire books from Project Gutenberg', QuALITY (Pang et al., 2022), a multiple choice question answering dataset over stories and articles sourced from Project Gutenberg; and Contract NLI (Koreeda & Manning, 2021) as natural language inference dataset in the legal domain. Table 8 compares our model with SotA iillion-arrametr encoder-decoder models which have been trained on large industrial compute clusters.
Effect of document structure: Our models are initially pre-trained on structured documents, and since most downstream tasks also maintain document structure, we aim to investigate the impact of such structures. To accomplish this, we selected two tasks, GovReport (Huang et al., 2021) and QASPER (Dasigi et al., 2021), where document structure preservation is integral. We proceeded by flattening the documents to create flattened version of the dataset. Subsequently, we trained both LED and HDT-E on these datasets, employing the pseudo-section setting introduced in Section 4.2 for the flattened data. A comparative analysis between models trained with and without real document structure for the two datasets is presented in Table 9, While we observe that modeling document structure is not having a very large effect on downstream task performance in this setting, it is yet very important to utilize document structure during pre-training in experiments. In fact, models pre-trained on pseudo sections and tested on the structured downstream tasks did not deliver any reasonable results. Our interpretation is that during pre-training, HDT learns to represent hierarchical information in different granularity via its anchor tokens, which allows to adapt pseudo-section data via fine-tuning.
Phttps:/wwww.gguteeeerggg/
20


Reviews:
[{"review_id": "ShrH9i3rzK", "reviewer": "Reviewer_9hSs", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**The experiments do not compare to more recent approaches as Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention**\n\nThanks for the suggestion. We first want to remark that the suggested paper considers auto-regressive decoder-only models and hence cannot be compared to encoder models directly. However, for this rebuttal, we adapted their model to an encoder-only variant by dropping the causal mask and using MLM for pre-training. However, we found this model difficult to train and could not get sensible results within the rebuttal period. We will continue experimenting with this model after the rebuttal and report our results in the paper, if successful.\n\n**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.**\n\nThe superscript $l=1,2,3$ consistently denotes the level of hierarchy and is explained in Figure 3. Both indices $i, j$ consistently denote positions. We will make this more clear at the start of the methodology section and make sure that all symbols and indices are named in the main text as well.\n\n**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.**\n\nThank you. The examples we use actually are real documents from arXiv (we will clarify this further) and show the first 1k tokens of each document (~25%) as described in the caption. The actual text is indeed shown at the axes but unfortunately not legible due to the figure resolution to keep the pdf size reasonable. We will provide high-resolution versions of these figures on our project page.\n\n**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?**\n\nThank you for pointing this out. This was indeed a typo, which will be corrected. Our method is slightly (2 ms) slower than HAT, which is due to our dynamic sample-specific attention pattern while HAT uses the same attention pattern for every sample. \n\n**Could you please add the missing description of notation used in equation 2-7?**\n\nWe will make sure that all symbols and indices are described in the main text.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The usage of position information for different hierarchy is interesting.\n- The results on SciRepEval Proximity tasks, and it leads to comparable or better results on SCROLLS summarization, QA, and NLE benchmark", "reasons_to_reject": "-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\n-The writing up of methodology could be improved: \nThe 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_to_authors": "- 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?\n- Could you please add the missing description of notation used in equation 2-7?", "extra_scores": {}}, {"review_id": "6aG2P2s4sf", "reviewer": "Reviewer_P6rv", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Long document classification**\n\nDuring the rebuttal period, we did additional experiments on the ECtHR-LJP dataset, a multi-label document classification task used in the HAT paper. We trained both models using our implementation and evaluated them with the micro-F1 metric. Our model shows marginal improvements compared to HAT (79.1 -> 79.2). Note that results for CONTRACT-NLI (CNLI) and QuALITY (QALT) can be found in Table 4, evaluated as a text-to-text generation task instead of a classification task. We have also applied for access to the MIMIC dataset and will add results once we have been approved.\n\n**Scaling to larger models with Bs of params**\n\nThe main focus of our paper is on processing structured datasets which are currently smaller than unstructured datasets used for training LLMs. However, we agree that scaling the model and datasets (eg, using additional structured data or combining structured data with unstructured data) is an interesting avenue which we want to explore in future work.\n\n**Warm-starting (model-recycling)**\n\nThis is indeed interesting, and we did some preliminary experiments on this. We found that in contrast to other models like Longformer and LongT5, our model can be stably trained from scratch which we consider an advantage of our model. One difficulty in warm-starting HDT from existing pre-trained models is that we use additional anchor tokens which must be trained as well. We plan to investigate this in the future.\n\n**Pre-training with ArXiv data**\n\nThanks for pointing this out. The reason that we choose SciRepEval as benchmark and Arxiv as a part of our pre-training data is that both include structured data. However, note that Arxiv documents make up for only 15% of our pre-training data. For this rebuttal, we conducted additional experiments and pre-train our model from scratch without ArXiv documents, following the same setting as before (1 GPU day), and evaluating it on SciRepEval. We find that with contrastive fine-tuning, our model maintains its performance (70.55 -> 70.33), while the pretrained-only model performance drops slightly (66.27 -> 65.16), yet still outperforms all baselines.\n\n**Comparison between HDT, HAT and Longformer on ListOps**\n\nWe have conducted new experiments comparing to Longformer and HAT on sequences up to 512 tokens. HDT (0.794) significantly outperforms BERT (0.599), Longformer (0.622) and HAT (0.584) with 5 operands and a maximum tree depth of 20.  \n\n**Scores in Table 2**\n\nSee caption.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The newly proposed method, HDT, is a great intuitive step forward for Hierarchical Transformers, similar to HAT, that try to impose document structure in the modeling. \n* The newly proposed method, HDT, seems to outperform several notable baselines in different benchmarks.\n* The paper is nicely written, organized, and easy to follow with great figures that can help readers understand the technical issue of processing long documents and the newly proposed architecture.", "reasons_to_reject": "* 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. \n* 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.\n* 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_to_authors": "* 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.\n* 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.\n* What are the scores for \"Feeds-M\" and \"High. Infl.\" refer to in Table 2?", "extra_scores": {}}, {"review_id": "GB8rwRO195", "reviewer": "Reviewer_XCB9", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**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.**\n\nThe main goal of this work was to design and implement an efficient Transformer for long documents using sparse attention. In this process, the choice of the attention pattern is a central question. As has been shown in studies with both humans [1,2] and language models [3], the document structure can be a useful signal when solving tasks on long documents, document structure offers a natural sparse attention pattern to process long documents more efficiently. To fully exploit document structure in sparse attention, a major challenge is the implementation of an attention kernel that handles arbitrary attention patterns efficiently. We solve this challenge and present results that show promising performance. As stated in the introduction, we will release our code and models.\n\n**References**\n\n[1] Barbara M. Taylor and Richard W. Beach. The Effects of Text Structure Instruction on Middle-Grade Students’ Comprehension and Production of Expository Text. Reading Research Quarterly, 19(2):134–146, 1984.\n\n[2] John T. Guthrie, Tracy Britten, and K. Georgene Barker. Roles of Document Structure, Cognitive Strategy, and Awareness in Searching for Information. Reading Research Quarterly, 26(3), 1991.\n\n[3] Jan Buchmann, Max Eichler, Jan-Micha Bodensohn, Ilia Kuznetsov, and Iryna Gurevych.\nDocument structure in long document transformers. In Conference of the European Chapter\nof the Association for Computational Linguistics (EACL), 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Clarity, well written. Proposes a novel, well executed and detailed approach.", "reasons_to_reject": "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.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8dwa2lSSz6", "reviewer": "Reviewer_vjsH", "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).\nThe 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). \nAdditionally, 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.\nTo 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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**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.**\n\nWe, too, are very interested in scaling our ideas to larger model sizes and task sets, such as code which is also structured data. We consider our work as a proof-of-concept to demonstrate that incorporating natural hierarchical structure via sparse attention is possible. \n\n**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 exciting.**\n\nThanks to our hierarchical attention design and the implementation of our kernel, we are able to process 8192 tokens as input during training and inference using standard (academic) compute hardware. While recently LLMs have been extended to larger contexts, they do not consider the input structure and require large distributed compute (eg, RingAttention) for training or alternative non-structured memory mechanisms. \n\nTraining/Fine-tuning our pre-trained structured model with longer inputs is possible both theoretically and computationally. Figure 8 shows how our hierarchical model scales to larger inputs in terms of runtime. In preliminary experiments, we increased the input length to 30k for NarrativeQA (Nrtv) in Table 4, resulting in slight improvements from 14.2 -> 15.3. However, most datasets that we evaluate on do not require input lengths larger than 8k.\n\n**Figure 1 (b): is this training loss or loss calculated on a hold-out dev set?**\n\nIn Figure 1(b) we report the training loss as also done in prior work (eg, CRAMMING) for evaluating pre-training. However, the validation loss behaves similarly, and we will also report it in the final paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The proposed architecture sounds promising: the hierarchical transformers have been investigated a lot for different NLP tasks; the paper furthers this direction by developing a new sparse attention kernel which can be used as a drop-in replacement of existing standard attention.\n* The authors evaluate the proposed transformer variant on various tasks to assess the effectiveness of the learned representations. The reported results surpass those of both HAT and Longformer. \n* The authors categorize the literature on long-document transformers into two categories and provide a nice overview.", "reasons_to_reject": "* 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.\n* 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_to_authors": "* Page 9, Efficiency Analysis `12-layer models models` -> `12-layer models`\n* 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.\n* Figure 1 (b): is this training loss or loss calculated on a hold-out dev set?\n* Not sure whether Figure 1 (a) and Figure 2 are duplicated, it seems Figure 2 is a more detailed version or version with example?\n* Suggest adding full name of SRAM when you first mention it.", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept