Hyenaの動作のイメージ。入力\(u\)から(\(v, x^1, \dots, x^N\))を生成し、それらの入力データに制御されたゲート(\(D_x^1, \dots, D_x^N\))と長い畳み込みフィルタ(Hyena Filter \(S_h^1, \dots, S_h^N\))による再帰的処理により、オリジナルのAttention並の性能を維持しつつ計算コスト削減に成功

入力長を増やした場合、Attentionを利用すると計算時間は飛躍的に増大するが、Hyenaの計算時間の増加は緩やかで、入力長が65,536に対しては100倍以上高速

こんにちは、グループ研究開発本部・AI研究室のT.I.です。またかと思われるかもですが、“Is Attention All You Need?”の続きです。“Is Attention All You Need?”にアクセスしますと、依然としてAttention機構は必要とされている現状がわかります。

前回のBlogでは、Transformerを利用しながらも200万トークンもの長い入力の処理が可能な“Recurrent Memory Transformer”を紹介しました。しかし、結局は一定長のセグメントに分割して処理を繰り返しているだけで、本質的な問題(入力長\(L\)に対して、計算コストが二乗 (\(\sim \mathcal{O}(L^2)\)))の解決には至っていません。また、前々回のBlogで紹介した状態空間モデル(SSM)を応用したS4は、計算コストを劇的に下げる可能性があるもののまだTransformerの性能には及びませんでした。しかし、SSMを利用する研究は着々と進められており、その発展系である今回紹介するHyenaは、ついに自然言語処理や画像処理などでTransformerに十分に匹敵するレベルの性能を発揮しました。

今回のBlogでは、このHyenaについて解説しますが、これはS4を発展させたH3(Hungry Hungry Hippos)を更に一般化したものにあたります。そのため、まずその前段階にあたるH3を解説します。

Hungry Hungry Hippos (arXiv:2212.14052)は、状態空間モデル(SSM)を利用したモデルで、前々回紹介したS4の改良版にあたります。奇妙な名前ではありますが、この元ネタはテーブルゲームの名前のようです。

このようなゲームみたいです

H3では、SSMの苦手とする長期の記憶を保持するために、H3 LayerというAttention機構に似た仕組みを利用しています。そして、H3 Layerを何層にも組み合わせることで、自然言語処理においてTransformerモデルと同等の性能を発揮することに成功しました。

まず、SSMについて簡単に復習します。入力 \(u(t) \in \mathbf{R}\)に対し出力\(y(t) \in \mathbf{R}\)とします。SSMでは、入力から出力への変換は行列 \(\mathbf{A} \in \mathbf{R}^{m\times m}\), \(\mathbf{B} \in \mathbf{R}^{m\times 1}\), \(\mathbf{C} \in \mathbf{R}^{1\times m}\), \(\mathbf{D} \in \mathbf{R}^{1\times 1}\)と状態変数 \(x(t) \in \mathbf{R}^m\)を含めた以下の方程式で定義されます。
$$
\begin{eqnarray}
\dot{x}(t) &=& \mathbf{A}x(t) + \mathbf{B}u(t) \\
y(t) &=& \mathbf{C}x(t) + \mathbf{D}u(t)
\end{eqnarray}
$$
これらの\(\mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D} \)はモデルの学習パラメータとなります。このSSMは入力(\(u_1, \dots, u_N\))に対して出力(\(y_1, \dots, y_N\))を計算するとき、フィルタ
$$
f = [\mathbf{CB}, \mathbf{CAB}, \mathbf{C} \mathbf{A}^2\mathbf{B}, \dots ,\mathbf{C}\mathbf{A}^{N-1}\mathbf{B}]
$$
を使った畳み込み(convolution)
$$
y_i = \mathbf{C}\mathbf{A}^i\mathbf{B}x_0 + (f \ast u)_i + \mathbf{D}u_i
$$
として処理できます。このように変形しておくことで、FFT(Fast Fourier Transformation)などを利用し非常に効率的に計算ができることをS4の紹介で解説しました。なお、以降では簡単のため
$$
y = \mathrm{SSM}_{\mathbf{A,B,C,D}}(u)
$$
と表記します。これらのフィルタ(行列\(\mathbf{A,B,C,D}\))は、モデルの性能を改善するために様々な条件が課されます。

さて、このSSMですが、実はある種の記憶保持タスクに対して性能が低いことが知られています。その例が以下の２つのタスクです。

1. Induction Head
2. Associative Recall

Induction
Headでは、入力で特別な記号\(\vdash\)の次の文字を記憶して出力するタスクです。具体的な例を見てみましょう。以下のアニメーションでは、入力の途中に\(\vdash\)が出現し、その後の単語\(f\)を記憶します。最後に出力を求める\(\vdash\)が現れるので、その次の文字\(f\)を出力します。

Induction Headの例(論文を元に作成)

次のAssociative RecallはInduction Headと似ていますが、文字のペアを記憶しておき最後に提示されるキーとペアとなる文字を出力するタスクです。具体的な例は以下の通りです。最後の文字\(a\)とペアとなっていた文字\(2\)を出力することが回答となります。

Associative Recall の例(論文を元に作成)

Synthetic language tasksの性能比較(論文を元に作成)

では、H3の性能の鍵は何でしょうか？H3では、TransformerのAttention機構にモチベートして設計されたH3 Layerを利用しています。これは、入力\(X\)をQuery (\(Q\))、Key (\(K\))、Value (\(V\))に射影し、Shift SSMとDiag. SSMという2つのSSMを利用します。数式ではこのように定義されます。
$$
\mathbf{Q} \odot \mathrm{SSM}_\mathrm{diag}(\mathrm{SSM}_\mathrm{shift}(\mathbf{K}) \odot \mathbf{V})
$$
ここで\(\odot\)は、要素毎の積(アダマール積)です。具体的なイメージを掴むために、H3 Layerの動作を以下のアニメーションで見てみましょう。

H3 layerでの処理のイメージ(図は論文を元に作成)

まず、入力(\(X\))から\(Q, K, V\)を生成します。それには入力シークエンス\(u \in \mathbf{R}^{N\times d}\)に対して、重み行列\(\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V \in \mathbf{R}^{d\times d}\)を用いて
$$
Q = u\mathbf{W}_Q, \quad K = u\mathbf{W}_K, \quad V = u\mathbf{W}_V
$$
で定義されます。そして、\(K\)をShift SSMに入力します。Shift SSMでは、\([a,b,c] \rightarrow [0, a, b]\)のような操作を含むSSMで、SSMのパラメータである行列\(\mathbf{A}\)の一部の要素のみが1(\(A_{i,j}=1\) for \(i -1 = j\))でそれ以外はゼロ(\(A_{i,j}=0\) for \(i – 1 \neq j\))となるような行列です。その出力結果と\(V\)の要素積を取り、Diag. SSMに入力します。Diag. SSMの行列\(\mathbf{A}\)は対角行列で、対角化HiPPO(S4D)を利用します。その出力結果と\(Q\)との要素積を取り、出力の射影行列\(\mathbf{W}_O \in \mathbf{R}^{d \times d}\)を掛け合わせて出力とします。このH3の計算コストとしては、\(\mathcal{O}(N\log N)\)であることがこの論文では示されています。これは通常のAttention機構である\(\mathcal{O}(N^2)\)よりも効率的です。なお、論文ではFlashConvというSSMの処理を高速化し学習を効果的にする手法も提案されています。詳細は論文を参照してください。

では、H3の自然言語処理での性能を紹介します。比較対象として、Transformer、S4D、Gated State Spacesを利用したモデルを用意しました。すべてのモデルは12層でパラメータ数は約125Mです。以下の図はこれらのモデルのOpenWebTextに対するPerplexityを比較したものです。この数値は小さいほど性能が高いことを示します。H3は、S4DとGated State Spacesよりも性能が高いことがわかります。ただし、ベースラインといえるTransformerには僅かに及びません。純粋なH3 Layerだけではなく、12層のうち2層をAttentionに差し替えたH3 Hybridでは、19.6とTransformerを1ポイント上回っています。

H3と他のモデルのOpenWebTextを用いたPerplexityの比較(図は論文を元に作成)

次に、SuperGLUEベンチマークについて、H3とGPT、OPTと比較した結果を紹介します。それぞれ同じ程度のパラメータ数のモデルを4段階(125M,355M,1.3B,2.7B)で比較しました。学習はThe PILEデータセットを用いて行いました。以下の図はそれぞれのスコアをレーダー・チャートと平均スコアをライン・チャートで示したものです。それぞれ、評価方法(Zero-shotか3-shot)とレーダー・チャートについてはパラメータ数(125M, 355M, 1.3B, 2.7B)を選択することで変更できます。タスク毎に揺らぎはありますが、zero-shotや3-shotでのスコアの平均を見るとHybrid H3では、Transformer系モデルであるOPT, GPTを上回っています。

Model Parameter: 2.7B
Model Parameter: 1.3B
Model Parameter: 355M
Model Parameter: 125M
Performance: zero-shot
Performance: 3-shot

SuperGLUEによる各種ベンチマーク結果のまとめ(論文を元に作成)(注) 実際には355Mでは、OPT-350M

最後に、Hybrid H3の推論スピードをTransformerと比較したものが以下の図です。評価はGPU A100でメモリーは80GBを利用します。入力長を3種類に変えた場合、どの長さのプロンプトに対してもHybrid H3の方が高速に推論できています。

推論のスピードの比較(図は論文を元に作成)

さて、H3がTransformerに匹敵する能力を持つことは分かりましたが、性能評価ではHybrid H3というAttentionを一部含むモデルであることが懸念です。Perplexityの比較にあるように、実はH3だけでは、Transformerの性能を越えることができていません。そのため、純粋にH3 Layerだけではなく、結局はAttentionの性能によるものではないかという疑問が残ります。この点を改良したモデルがHyena (https://arXiv/abs/2302.10866)になります。Hyenaのポイントは、Hyena Recurrenceという、Attentionに類似した行列の畳み込み処理とHyena Filterというデータの処理にあります。

Hyenaの論文の冒頭で、従来のTransformerを改良した二乗未満でスケールする手法ではオリジナルのTransformerほどの性能が果たせなかった点に関し、Attentionの高い性能の要因として以下の3つの点に注視しました。

1. Data control: データに制御された線形演算であり、1つのブロックで全ての線形関数を埋め込むことができる。Attention行列\(A\)を用いて、\(y = A(k,q)v \)と書き下せます。入力\(k, q\)で\(v\)についての線形関数を制御していることになります。
2. Sublinear parameter scaling: Attention layer のパラメータの数は入力長と独立している。そのため、Attention layer間のニューラルネットワークなどに更にパラメータを追加できる。
3. Unrestricted context: 入力のコンテキスト(順序など)に依存せずに任意の要素間の関係性を扱える。

これらの性質を指針としてHyenaは設計されております。

まず、入力長\(L\)のシグナル\(u_t \in \mathbf{R}^D\)のフィルタ\(h \)による畳み込み操作について考察します。
$$
y_t = (h\ast u)_t = \sum_{n=0}^{L-1} h_n u_{t-n}
$$
\(D \)は一般に入力の次元数ですが、簡単のために\(D=1 \)としたSISO(single input single output)の場合を考えますと、この計算はToeplitz kernel matrix \(S_h \in \mathbf{L \times L}\)との行列積
$$
(h \ast u) =
\begin{pmatrix}
h_0 & h_{-1} & \dots & h_{-L+1} \\
h_1 & h_0 & \dots & h_{-L+2} \\
\vdots & \vdots & \ddots & \vdots \\
h_{L-1} & h_{L-2} & \dots & h_0
\end{pmatrix}
\begin{pmatrix}
u_0 \\
u_1 \\
\vdots \\
u_{L-1}
\end{pmatrix}
$$
によって表現されます。Toeplitz行列とは、\(A_{i,j} = A_{i+1,j+1}\)と対角方向に同じ行列のことです。Convolutional Neural Network(CNN)などでは、このフィルタ\(h\)を直接に学習のパラメータとして最適化します。この場合、カーネルの大きさ\(M\)は入力長よりも一般的に小さくなります(\(M \ll L\))。

カーネルを大きくするほどより離れた要素間の関係性を扱えるようになりますが、その分パラメータが増えます。そこで、直接フィルタそのものを学習(explicit convolution)するのではなく、パラメータを用いてフィルタを間接的に表現してみます。
$$
h_t = \gamma_\theta(t)
$$
ここで、\(\theta\)は、関数\(\gamma_\theta\)のパラメータです。このような間接的な畳み込み(implicit convolution)を利用することで、入力長と同等の長い畳み込みフィルタを扱うことができます。その例が、まさにSSMを利用したフィルタです。
SSM
$$
\begin{eqnarray}
x_{t+1} &=& Ax_t + Bu_t \\
y_t &=& Cx_t + D u_t
\end{eqnarray}
$$
は、以下の畳み込み操作に帰着できます。
$$
y_t = \sum_{n=0}^t (CA^{t-n}B + D\delta_{t-n})u_n
$$
これをフィルタによる畳み込みと見ると、
$$
t \rightarrow h_t = \begin{cases}
0 \qquad &t < 0& \\ CA^t B + D \delta_t \quad &t \geq 0& \end{cases} $$ となります。

HyenaはH3を更に一般化させたものです。Order-\(N\) Hyena演算子は以下で定義されます。まず、入力\(u\)を\(N+1\)個のベクトル(\(v, x^1, \cdots, x^N\))に射影します。これは、AttentionやH3で入力をKey、Query、Valueと3つに分けたことの拡張になります。そして、入力長と同じ長さのフィルタ(\(h^1, \dots, h^N\))による畳み込みと次のベクトルとの要素積を再帰的に処理します。
$$
\begin{eqnarray}
z_t^1 &=& v_t \\
z_t^{n+1} &=& x_t^n(h^n \ast z^n)_t \quad n \in 1, \dots N \\
y_t &=& z_t^{N+1} \\
\end{eqnarray}
$$
このHyena recurrenceの計算コストは、\(\mathcal{O}(NL\log_2 L)\)になります。この計算をまとめると以下のようになります。
$$
y = x^N \cdot(h^N \ast(x^{N-1} \cdot (h^{N-1} \ast \cdots (x^1 \cdot (h^1 \ast v)))))
$$
Attentionでは、QueryとKeyの重みでValueを処理していました。Hyenaでは、それが\(N\)-stepも続くような機構となっています。これだけでは抽象的すぎて何のことだかわからないので、以下で具体的に各ステップでの演算と特徴を説明します。

SelfAttention matrixとHyena matrixの比較(論文を元に作成)

Hyenaの演算を具体的に行列で表現してみます。実はOrder-2 Hyenaは、H3と等価となります。簡単のためSISO(\(D=1\))のH3がどのようにHyenaのフレームワークで記述されるか見ていきます。ここで、\(D_q\)と\(D_k\)という、\(L \times L\)の対角行列とします。これはH3であったQuery、Keyの要素を対角成分に持つ行列で、入力データに対応し制御する箇所になります。H3の計算は以下の式にまとめられます。
$$
\begin{eqnarray}
A(q, k) &=& D_q S_\psi D_k S_\varphi \\
H3(q, k, v) &=& A(q, k)v
\end{eqnarray}
$$
ここでフィルタに相当する\(S_\psi, S_\varphi\)は、\(L \times L\)のToeplitz行列で学習パラメータです。具体的に書き下すと
$$
S_\psi =
\begin{pmatrix}
\psi_0 & 0 & \cdots & 0 \\
\psi_1 & \psi_0 & \cdots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
\psi_{L-1} & \psi_{L-2} & \cdots & \psi_0 \\
\end{pmatrix}, \quad
S_\varphi =
\begin{pmatrix}
\varphi_0 & 0 & \cdots & 0 \\
\varphi_1 & \varphi_0 & \cdots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
\varphi_{L-1} & \varphi_{L-2} & \cdots & \varphi_0 \\
\end{pmatrix}
$$
となります。因果性を持たせるために、\(S_\psi, S_\varphi\)の上三角の要素は0とします。これを一般化すると
$$
y = H(u)v = D_x^N S_h^N \cdots D_x^2 S_h^2 D_x^1 S_h^1v
$$
とHyena recurrenceをHyena行列の積として表現できます。ここで、\(D_x^n = \mathrm{diag}(x^n) \in \mathbf{R}^{L \times L}\)、\(S_h^n\)はフィルタ\(h^n\)の要素からなるToeplitz行列です。

Hyena演算子はunbounded contextを持つことが示されます。それにより局所的だけではなく長く離れた要素感も関係性を学習可能です。

畳み込みのフィルタとして、以下のような Feed-Forward Neural Network (FFN)とWindow関数、PositionalEncodingを組み合わせたものを利用します。
$$
h_t = \mathrm{Window}(t) \cdot (\mathrm{FFN} \circ \mathrm{PositionalEncoding})(t)
$$
ここでPositionalEncodingとは、位置情報を付与する処理です。Hyenaでは、\(\rho_k(t) = e^{i2\pi k t/L}\) (\(k \in 0, \dots, K-1\))として、以下のように定義します。
$$
\mathrm{PositionalEncoding}(t)
= [t, \mathrm{Re}[\rho_0](t), \mathrm{Im}[\rho_0](t), \dots, \mathrm{Re}[\rho_{K-1}](t),
\mathrm{Im}[\rho_{K-1}](t)]
$$
これにより、元々の入力長に加えて\(D_e = 2K+1\)の次元が追加されます。また、Window関数として、
$$
\mathrm{Window}(t) = \exp\{-\alpha t\}
$$
を利用します。

Hyena Filterの構成要素(論文を元に作成)

最後にHyenaの計算手順をまとめます。

1. Hyena Projection: 入力\(u \in \mathbf{R}^{L \times D}\)に対して、\(x^1, \dots, x^N, v \in \mathbf{R}^{D \times L}\)を計算する。
   $$x^1, \dots, x^N, v = \mathrm{Projection}(u)$$
   
2. Hyena Filter: \(h^1, \dots h^N \in \mathbf{R}^{D \times L}\) を計算する。
   $$h^1, \dots, h^N = \mathrm{HyenaFilter}(L, D_e)$$
   
3. Forward pass: Hyena FilterをFFTを利用した高速なアルゴリズムで順次畳み込み計算していく。
   $$
   v_t \leftarrow x_t^n \cdot \mathrm{FFTConv}(h^n, v)_t
   $$
   を順次計算(\(n = 1, \dots, N\))し、最後の\(v (= y)\)として出力。

最終的な計算コストは入力 \(u \in \mathbf{R}^{L \times D}\)に対して、order-N Hyena演算子は
$$
\mathcal{O}(NDL(\log_2 L + D))
$$
であり、入力長の二乗未満でスケールしております。

Hyena Recurrenceのイメージ(論文を元に作成)

まずは、Hyenaの長距離の畳み込みフィルタリングの性能を紹介します。直接フィルタのパラメータを最適化する手法(Conv1d, FNO)、SSMなどを介して間接的に最適化する手法(H3, TransferFunc, CKConv)と比較します。以下の図が、Associative Recallの結果です。間接的な手法の方が性能は良い傾向がありますが、ボキャブラリーのサイズが大きくなるほど、徐々にHyena以外のモデルの性能が低下していきます。一方でHyenaは十分長い入力長があれば、高い性能を維持しております。

Hyenaと他のフィルタリングの手法の比較(論文を元に作成)

では、更に長い入力長にした場合はどうなるでしょうか?以下の図は、Transformerなどの他のモデルに対して、長い入力長(最大13万)まで入力した場合の精度の比較です。Hyenaが他のモデルを圧倒して高い精度を維持しております。

Hyenaと他のモデルとの入力長を変えた場合のAssociative Recallの性能比較(図は論文を元に作成)

次に、言語モデリングの性能を紹介します。同程度のモデルサイズ(125Mパラメータ)で、TransformerなどとWikiText103とTHE PILEデータセットに対するPerplexityのまとめが以下です。PerformerなどのTransformerを効率化を測ったモデルでは、元のTransformerより性能が低い(高いPerplexity)ではありますが、一方のHyena-3は、Transformerと同程度の高い性能(低いPerplexity)を達成しました。

Hyenaと他のモデルのPerplexityの比較(図は論文を元に作成)

THE PILEデータセットの結果では、学習のトークン数を上げるにつれて、HyenaはGPTとほぼ同程度の性能改善をしておりますが、パラメータ数が355Mのものでは、GPTと比較して、学習の計算コスト(FLOPs)が20%程度少なくなっています。

Hyenaと他のモデルの学習時のPerplexityと計算コストの比較(論文を元に作成)

THE PILEデータセットでの学習結果として興味深い点が、Hyenaの性能にScaling Lawsらしきものが見られることです。これはLLMがパラメータ数やデータ数を増やすと、いくらでも性能が向上していくという性質で、現在までその限界は見えておらず、LLMを更に巨大化する大きな動機付けとなっています。

SuperGLUEによる各種ベンチマーク結果のまとめが以下の図です。比較対象として、RWKVとGPTNeoのモデルの結果を併記しました。Hyenaは137Bのトークンでの学習、他のモデルは332Bと300Bトークンでの学習とHyenaは半分ではありますが、ほぼ同程度の性能を発揮しております。

HyenaとGPTNeo、RWKVのSuperGLUEのパフォーマンスの比較(論文を元に作成)

次に計算時間についてです。HyenaとAttention、そしてFlashAttention layersの計算時間を比較したものが以下の図です。短い入力長に対してはHyenaは性能を発揮できませんが、入力長に応じてAttentionとFlashAttentionの実行時間がどんどん増加しますが、Hyenaはほとんど変わりません。65,536トークンの入力長では、HyenaはFlashAttentionの100倍以上高速に処理可能です。

HyenaとAttention、FlashAttentionの入力長を変えた場合の計算時間の比較(論文を元に作成)

最後に画像処理タスクについて紹介します。そのために、Vision Transformer(ViT)のAttentionをHyenaに置き換えたモデル(Hyena-ViT)を用意しました。ImageNet-1kの性能は以下の図の通りです。ViTと比較して、Hyena-ViTは同等の性能を維持しております。また、SSMを利用するモデルS4NDとCIFAR-10の性能比較した場合、やや上回る性能(1.7ポイント)でありながら、モデルのサイズは25%も小さくなっております。

Hyenaを利用した画像分類タスクの性能比較(図は論文を元に作成)

さて、3回目となったIs Attention All You Need?シリーズですが、今回は、状態空間モデルを利用したS4を更に進化させたHyenaというモデルを紹介しました。Hyenaでは、H3(Hungry Hungry Hippos)で導入されたAttention機構にモチベートされたH3 Layerを更に一般化したものを利用しています。その結果、計算コストがTransformerのように\(O(L^2)\)で増加せずに、自然言語処理や画像処理タスクに対して高い性能を発揮できました。

しかし、もともとAttentionは必要か?という観点から見ると、結局、「Attention的な処理は必要である。」という結論になりました。ただし、従来のTransformerで利用するような二乗でコストが掛からないことが大きな利点です。これにより従来のAttention並の性能を維持しつつ、更なる大規模計算が可能になることは大きなブレイクスルーになるのではないかと期待しております。

グループ研究開発本部 AI研究開発室では、データサイエンティスト/機械学習エンジニアを募集しています。ビッグデータの解析業務などAI研究開発室にご興味を持って頂ける方がいらっしゃいましたら、ぜひ募集職種一覧からご応募をお願いします。皆さんのご応募をお待ちしています。