Tree Cross Attention

要約

クロス アテンションは、予測を行うために一連のコンテキスト トークンから情報を取得する一般的な方法です。
推論時に、各予測について、クロス アテンションは $\mathcal{O}(N)$ トークンの完全なセットをスキャンします。
ただし、実際には、多くの場合、良好なパフォーマンスを得るには、トークンの小さなサブセットのみが必要です。
Perceiver IO などのメソッドは、情報をより小さいサイズの潜在トークンのセット $L < N$ に抽出し、その後クロス アテンションが適用されるため、推論のコストが低くなり、複雑さは $\mathcal{O}(L)$ のみになります。 。 ただし、実際には、入力トークンの数と抽出する情報の量が増加するにつれて、必要な潜在トークンの数も大幅に増加します。 この研究では、ツリー クロス アテンション (TCA) を提案します。これは、推論を実行するために対数 $\mathcal{O}(\log(N))$ 個のトークンから情報のみを取得するクロス アテンションに基づくモジュールです。 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, Google