Standard inference and training with transformer based architectures scale
quadratically with input sequence length. This is prohibitively large for a
variety of applications especially in web-page translation, query-answering
etc. Consequently, several approaches have been developed recently to speedup
attention computation by enforcing different attention structures such as
sparsity, low-rank, approximating attention using kernels. In this work, we
view attention computation as that of nearest neighbor retrieval, and use
decision tree based hierarchical navigation to reduce the retrieval cost per
query token from linear in sequence length to nearly logarithmic. Based on such
hierarchical navigation, we design Treeformer which can use one of two
efficient attention layers -- TF-Attention and TC-Attention. TF-Attention
computes the attention in a fine-grained style, while TC-Attention is a coarse
attention layer which also ensures that the gradients are "dense". To optimize
such challenging discrete layers, we propose a two-level bootstrapped training
method. Using extensive experiments on standard NLP benchmarks, especially for
long-sequences, we demonstrate that our Treeformer architecture can be almost
as accurate as baseline Transformer while using 30x lesser FLOPs in the
attention layer. Compared to Linformer, the accuracy can be as much as 12%
higher while using similar FLOPs in the attention layer.

本文提出了一种名为 Treeformer 的新型 transformer 模型，通过采用基于决策树的层次化导航方法，将 attention 计算的检索成本从线性降低为几乎对数级，同时使用 TF-Attention 和 TC-Attention 两种 attention 层实现了性能优于基准 Transformer 30 倍 FLOPs 的结果。