HDT: Hierarchical Document Transformer

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.

We, 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.

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.

Thanks 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.

Training/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.

Figure 1 (b): is this training loss or loss calculated on a hold-out dev set?

In 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.

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.

The 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.

References

[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.

[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.

[3] 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.

Long document classification

During 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.

Scaling to larger models with Bs of params

The 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.

Warm-starting (model-recycling)

This 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.

Pre-training with ArXiv data

Thanks 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.

Comparison between HDT, HAT and Longformer on ListOps

We 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.

Scores in Table 2

See caption.

The experiments do not compare to more recent approaches as Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention

Thanks 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.

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.

The 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.

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.

Thank 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.

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?

Thank 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.

Could you please add the missing description of notation used in equation 2-7?

We will make sure that all symbols and indices are described in the main text.