ルーティングTransformerを用いた効率的なコンテンツベースのスパースアテンション

近年、自己注意（self-attention）はさまざまな系列モデリング問題に広く採用されている。その有効性は認められているが、系列長に伴い計算量およびメモリ使用量が二次関数的に増加するという課題を抱えている。この複雑性を低減するための成功したアプローチは、局所的なスライディングウィンドウに注目するか、コンテンツに依存しない少数の位置に注目する手法に集約されている。本研究では、関心のあるクエリと関係のないコンテンツに計算資源やメモリを割り当てないよう、動的なスパースな注目パターンを学習する手法を提案する。本研究は2つの研究分野に根ざしており、コンテンツに基づくスパース注意のモデリングの柔軟性と、局所的かつ時系列的なスパース注意に基づく効率性の両方を統合している。提案モデルであるルーティングトランスフォーマー（Routing Transformer）は、オンラインk-meansに基づくスパースルーティングモジュールを自己注意に組み込み、系列長がn、隠れ次元がdの場合、従来の$O(n^2d)$から$O(n^{1.5}d)$へと全体の注意計算量を削減する。実験では、Wikitext-103における言語モデリング（困惑度：15.8 vs 18.3）およびImageNet-64における画像生成（3.43 vs 3.44 ビット/次元）において、同等のスパース注意モデルを上回る性能を達成するとともに、自己注意層の数を削減している。さらに、新しく公開されたPG-19データセットにおいても、長さ8192の系列で学習した22層のルーティングトランスフォーマーにより、33.2という新たな最良のテスト困惑度を達成し、新たな最先端性能を樹立した。