Gated Delta Net

概要

Gated Delta Netは系列データを扱うモデルであり、Attentionをベースにしている。Attentionは過去系列の情報を現在の入力に応じて重み付けして出力を計算するため、key-valueペアを記憶する必要がある。Attentionの指数関数を分解できる形で近似したLinear Attentionでは、このkey-valueペアを行列 に圧縮できる。正規化して分母を落とせば、更新は単純な漸化式ひとつで済む。ただしこの漸化式は足すことしかできず、大きさの決まった に情報を足し続けるので中身が飽和してしまう。Gated Delta Netでは、キーの方向だけを書き換える更新と、状態全体を薄れさせる減衰の2つを持たせることで、ピンポイントな上書きと不要な記憶の消去を両立させている。

長い系列を扱うほど、過去をすべて保持するAttentionに対して有利になる。実際、言語モデリングや長文の理解ではTransformerと同等以上の性能が報告されている。一方で、過去の任意の位置をそのまま引いてくるような検索では及ばない。固定サイズの状態に圧縮する以上、そこに入り切らない情報は失われるためである。そのため実際には、層の大部分をGated Delta Netにして一部にAttentionを混ぜた、ハイブリッドの形で使われることが多い。

Linear Attention

一般的なAttentionでは系列の1要素 から線形変換でクエリ・キー・バリューを作った。

ここで形状は 、 とする。すると、現在時刻 までの情報のみをみるCausal Attentionでは、出力 は

で計算できる。ここで、非線形な 関数を次のような因数分解可能な関数 で近似してみる。

すると、これを用いたAttentionの出力は次のように書き換えられる。

これは一般的にLinear Attentionと呼ばれる。分子について考えると、

となる。ここで、

と定義すると

と表せる。また、これらは定義より以下の漸化式で再帰的に計算できる。

Delta Rule

Gated Delta Netでは上記の計算において分母を無視する。これは正規化のために分母を計算する必要があるが、別の部分で正規化を行うことにする。つまり以下のように があらかじめ正規化されているとする。

また、その後の出力も正規化したものとすると以降は簡単のため

と書くことにする。すると、Linear Attentionは次のように書き換えられる。

ここで、 はkeyのベクトル空間をvalueのベクトル空間に写す線形変換である。なぜなら、 が一つのkey-valueペアを記憶していると仮定する。つまり、 とすると、正規化されたベクトルの元では

が成立する。逆に のとき

となる。つまり、 はkeyのベクトル空間をvalueのベクトル空間に写す線形変換である。さて、この前提において を計算してみると、

つまり、 は真に を に写すような変換にはなっていない。 が残ってしまう。そこで、 を となるような変換になるためにオンライン学習される（Gated Delta Netを学習する段階ではなく、推論時にある意味で学習される）ものだと捉える。損失関数は単純に以下のように定義する。

すると、その勾配は

とある。これを用いて、 をいわゆる勾配法で更新して とする。つまり、学習率 を用いて

となる。これが Delta Rule である。このDelta Ruleがちゃんと新たな を復元できることを確認する。 のとき、更新後の は

となる。つまり、 は に向かって更新されることがわかる。このままではKey の方向に沿った更新しかできないので全体に減衰をかけるため、係数 を用いて への更新時、 を とする。これを計算すると、

となり、Delta Ruleは以下のようにまとめられる。

マルチヘッド化とGDNのブロック構造

シーケンスの1入力 からDelta Ruleで用いるための を作り、Delta Ruleから出力を得るまでを一つのブロックと考えることができる。それには色々な方法が考えられ、採用されるモデルによって多少のバリエーションがある。ここでは、元論文のもので説明する。

まず、基本的にはマルチヘッドを採用する。Delta Ruleの計算を同じ入力に対して複数回行う。ヘッド に対して、 を作って

を計算する。さて、各パラメータの計算方法を説明する。Attentionと同様に、クエリ・キー・バリューは線形変換で得られる。

つまり、入力行列 から

しかし、これらは正規化されていないので、正規化を行う。また、同時に時間方向に畳み込みを行う。この畳み込みはShortConvと呼ばれるものだが、Mambaのものと同様で過去のいくつかの時刻の情報をDepthwiseで畳み込むものである。 重み を用いて以下のように計算される。

なお、 は畳み込みのカーネルサイズである。また、正規化はL2正規化を行う。つまり、 に対して

と計算する。ただ、行列に対しては

のように、各行ベクトルごとに正規化を行う。これらを用いて、 は以下のように計算される。

は今までの議論から正規化を行う必要はないが、出力で正規化を行うことになる。 は入力 とパラメータとの内積をシグモイド関数で正規化することで得られる。概念的には1ニューロンへのバイアスなしの全結合である。

対して、 は以下のように計算する。

はスカラー、 はベクトルであり、これらは学習されるパラメータである。softplus関数は以下のように定義される。

これは は単純にkey方向に沿った更新の大きさを決めるため単純に係数を作るだけでいいが、 は減衰の大きさを決めるため指数減衰の形にする必要があるためである。指数減衰には様々な都合の良い性質があることが知られている。

さて、Delta Ruleによって得られた各ヘッドの出力を とすると、RMSNorm正規化を行う。RMSNorm正規化とは以下のようなものである。

RMSとはRoot Mean Squareの略であり、ベクトルの各要素の二乗平均平方根で正規化することを意味する。それだけでなく、正規化後に学習されるパラメータ を要素ごとに掛けることで、正規化後のスケールを調整することができる。RMSNormはLayerNormと同様の効果があるが、平均を引く操作がないため計算コストが低いという利点がある。これらはヘッドごとに独立に計算される。さらに、それぞれに対して入力との gating を行う。

入力ブランチ側は線形変換と非線形関数 を用いた変形を行う。

これらを結合して線形変換を行い、最終的な出力 を得る。結合の方法は単純に連結するだけである。