Tree Cross Attention

要約

クロスアテンションは、予測を行うために文脈トークンの集合から情報を取り出す ための一般的な手法である。推論時に、各予測に対して、Cross Attentionは$mathcal{O}(N)$トークンの全セットを走査する。しかし実際には、良好な性能のために必要なトークンのサブセットはわずかであることが多い。Perceiver IOのような方法は、情報を潜在トークン$L < N$の小さいセットに抽出し、それに対して交差注意を適用するため、推論が安く、結果として$mathcal{O}(L)$の複雑さしか生じない。しかし実際には、入力トークンの数と抽出すべき情報量が増加するにつれて、必要とされる潜在トークンの数も大幅に増加する。本研究では、Cross Attentionに基づき、対数$mathcal{O}(Γlog(N))$個のトークンから推論に必要な情報だけを取り出すモジュールTree Cross Attention (TCA)を提案する。TCAはデータを木構造で整理し、推論時に木探索を行い、予測に関連するトークンを取り出す。TCAを活用して、トークン効率の良い推論のための柔軟なアーキテクチャであるReTreeverを紹介する。我々は、ツリークロスアテンション(TCA)が、様々な分類や不確定性回帰タスクにおいて、クロスアテンションと同等の性能を発揮する一方で、トークン効率が大幅に向上することを実証的に示す。さらに、ReTreeverをPerceiver IOと比較し、推論に同じトークン数を使用しながら、大幅な性能向上を示す。

要約(オリジナル)

Cross Attention is a popular method for retrieving information from a set of context tokens for making predictions. At inference time, for each prediction, Cross Attention scans the full set of $\mathcal{O}(N)$ tokens. In practice, however, often only a small subset of tokens are required for good performance. Methods such as Perceiver IO are cheap at inference as they distill the information to a smaller-sized set of latent tokens $L < N$ on which cross attention is then applied, resulting in only $\mathcal{O}(L)$ complexity. However, in practice, as the number of input tokens and the amount of information to distill increases, the number of latent tokens needed also increases significantly. In this work, we propose Tree Cross Attention (TCA) - a module based on Cross Attention that only retrieves information from a logarithmic $\mathcal{O}(\log(N))$ number of tokens for performing inference. TCA organizes the data in a tree structure and performs a tree search at inference time to retrieve the relevant tokens for prediction. Leveraging TCA, we introduce ReTreever, a flexible architecture for token-efficient inference. We show empirically that Tree Cross Attention (TCA) performs comparable to Cross Attention across various classification and uncertainty regression tasks while being significantly more token-efficient. Furthermore, we compare ReTreever against Perceiver IO, showing significant gains while using the same number of tokens for inference.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL