HDT: Hierarchical Document Transformer

Summary

This paper describes a sparse Transformer architecture tailored for structured hierarchical documents (document->section->sentence->token). The new architecture introduces auxiliary anchor tokens (indicating the boundary of document/section/sentence) and redesigns the attention mechanism into a sparse multi-level hierarchy (ensuring each token only attending to its parent, siblings and children). The authors point out that there are two main benefits of the hierarchical document transformer: exploiting the document structure as an inductive bias (better performance on downstream tasks); and, sparse attention, leading to gains in computational and memory efficiency (faster pre-training). Additionally, the authors have created an efficient attention kernel using the Triton library, which can seamlessly replace existing Transformers. The implementation aims to minimize empty block computations by employing a simple heuristic: sorting keys and values according to their hierarchy level first. To demonstrate the effectiveness of proposed structure-aware attention, the authors conduct a series of experiments on different datasets, including ListOps (mathematical reasoning), SciRepEval, FacetSum, SCROLLS. The authors show that their method outperforms Hierarchical Attention Transformer, which segments long text without sticking to the document structure, and Longformer, which uses sparse attention.

Rating

7

Confidence

4

Ethics flag

1

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. * 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. * 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. * 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` * Figure 6 and the corresponding text description is Section 4.1 is difficult to understand. Also it is not very clear how these experiments relate to the paper's contribution. * Figure 1 (b): is this training loss or loss calculated on a hold-out dev set? * Not sure whether Figure 1 (a) and Figure 2 are duplicated, it seems Figure 2 is a more detailed version or version with example? * Suggest adding full name of SRAM when you first mention it.

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.

Rating

7

Confidence

3

Ethics flag

2

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.

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.

Rating

7

Confidence

5

Ethics flag

1

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. * The newly proposed method, HDT, seems to outperform several notable baselines in different benchmarks. * 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. * The evaluation is mainly focused on the size regime of medium-sized models (~100M params.), which leaves open the question of how HDT would benefit from scaling to larger models with Bs of params, where it could also have a more fair comparison to LongT5 and CoLT5 models. * The authors did not consider the idea of warm-starting (model-recycling) used by Longformer, HAT, and LongT5, which is another dimension of efficiency concerning saving compute time and resources.

Questions 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. * I would like to see a comparison between HDT, HAT, and Longformer on long-form ListOps with Ks of numbers and ops. I think it would be a more interesting comparison instead of BERT which completely lacks any form of anchor tokens or explicit hierarchical attention. * What are the scores for "Feeds-M" and "High. Infl." refer to in Table 2?

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.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

- The usage of position information for different hierarchy is interesting. - 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 -The writing up of methodology could be improved: The description of HPE and hierarchical attention are not very clear. The same index $i$ is used for position and mask, superscript $1$,$2$,$3$ used in Equation 4,5 & 6 are not explained. Fig 10 is referenced for a mask that is sparse in practice, but it’s not clear what is the document structure for that example? Providing the text next to it could help to support this claim as currently it looks like a toy example.

Questions 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? - Could you please add the missing description of notation used in equation 2-7?