Nyströmformer: Nyström 法による線形時間とメモリでの自己注意の近似

TL;DR

Nyströmformer は、標準的な自己注意機構を近似し、線形時間・メモリ計算量 $O(n)$ を実現する効率的な Transformer バリアントです。従来の二次計算量 $O(n^2)$ に比べ、はるかに長い入力シーケンスを扱うことができ、NLP とコンピュータビジョンのタスクにおいて競合できる性能を維持します。

行列近似のための Nyström 法

Nyström 法は、大きな行列 $P^{n \times n}$ を $m$ 行と $m$ 列のサブセットをサンプリングすることで近似します。これらのサンプルをブロック行列に配置することで、元の行列は 4 つの部分行列に分割されます：$A_{P}$ ($m \times m$)、$B_{P}$ ($m \times (n - m)$)、$F_{P}$ ($(n - m) \times m$)、および $C_{P}$ ($(n - m) \times (n - m)$)。

Because the entries of $A_{P}$, $B_{P}$, and $F_{P}$ are known through sampling, the unknown submatrix $C_{P}$ can be estimated as:

$$C_{P} = F_{P} A_{P}^{+} B_{P}$$

ここで $+$ は Moore-Penrose 擬似逆行列を表します。得られた Nyström 近似 $\hat{P}$ は 3 つの行列の積として表され、全行列を計算する負荷が軽減されます。

Nyström を自己注意に適用する

標準的な自己注意はソフトマックス行列 $S = \text{softmax}(\frac{QK^T}{\sqrt{d}})$ に依存します。$S$ に Nyström 法を直接適用することは不可能です。なぜなら、ソフトマックス行列の 1 列を計算するには、行方向のソフトマックス分母を求めるために他のすべての列の情報が必要になるからです。

この問題を解決するために、Nyströmformer は結果のソフトマックス行列からではなく、クエリ ($Q$) とキー ($K$) から「ランドマーク」（Nyström 点）をサンプリングします。モデルはクエリランドマーク $\tilde{Q}$ とキーランドマーク $\tilde{K}$ に基づいて 3 つの行列を定義します：

  • $\tilde{F} = \text{softmax}(\frac{Q\tilde{K}^T}{\sqrt{d}})$
  • $\tilde{A} = \text{softmax}(\frac{\tilde{Q}\tilde{K}^T}{\sqrt{d}})^{+}$
  • $\tilde{B} = \text{softmax}(\frac{\tilde{Q}K^T}{\sqrt{d}})$

近似されたソフトマックス行列 $\hat{S}$ は $\hat{S} = \tilde{F} \tilde{A} \tilde{B}$ と計算されます。この近似を値 ($V$) と掛け合わせることで、完全な $QK^T$ 積を計算せずに線形時間で自己注意を近似できます。

ランドマークの選択と実装

Nyströmformer はセグメント平均を用いて $\tilde{Q}$ と $\tilde{K}$ を構築します。$n$ 個のトークンを $m$ セグメントに分割し、各セグメントの平均をランドマークとして使用します。研究によれば、32 または 64 個程度のランドマークさえあれば、シーケンス長 $n = 4096$ や $8192$ の場合でも競合できる性能が得られます。

実装の詳細

  • Depthwise Convolution: アーキテクチャは、値に加えるスキップ接続として 1D depthwise convolution (DConv) を含みます。
  • Complexity: この手法は $O(n^2)$ の計算量を回避し、長いシーケンスでも効率的に動作します。

Hugging Face での利用可能性と使用方法

Nyströmformer は Hugging Face Transformers ライブラリを通じて Masked Language Modeling (MLM) に利用可能です。シーケンス長に応じて 4 つのチェックポイントが提供されます：nystromformer-512、nystromformer-1024、nystromformer-2048、nystromformer-4096。

ユーザーは NystromformerConfig の num_landmarks パラメータでランドマーク数 $m$ を制御できます。モデルは NystromformerForMaskedLM クラスまたは高レベルの pipeline API を使用して、fill-mask などのタスクにデプロイできます。

Sources