∞-former: Infinite Memory Transformer

https://arxiv.org/abs/2109.00301v1
2021

トランスフォーマーは長い文脈を扱う場合、文脈の長さに応じて計算量が増加するため、長期記憶を効果的にモデル化することができません。この問題を解決するためにいくつかのバリエーションが提案されているが、いずれも記憶容量に限界があり、古い情報を削除せざるを得ない。本論文では、バニラ変換器を拡張し、非限定的な長期記憶を持つ∞-formerを提案する。 ∞-formerは、長期記憶に注目するために連続空間の注目メカニズムを利用することで、注目の複雑さが文脈の長さに依存しなくなる。これにより、任意の長さのコンテキストをモデル化し、一定の計算量を維持しながら「スティッキーな記憶」を維持することが可能になります。合成ソートタスクでの実験により、∞-formerが長いシーケンスからの情報を保持する能力を実証した。また、言語モデリングの実験では、ゼロからモデルを学習したり、事前に学習した言語モデルを微調整したりすることで、拘束力のない長期記憶の利点を示しています。

1 はじめに

文書を読んだり書いたりするときには、以前に読んだり書いたりした情報を記憶しておくことが重要です。人間には、長期的な文脈を記憶する顕著な能力があり、関連する詳細を記憶に留めておくことができます（Carroll, 2007; Kuhbandner, 2020）。近年、変換器ベースの言語モデルは、コンテキストサイズを大きくすることで素晴らしい結果を得ている（Radford et al., 2018, 2019; Dai et al., 2019; Rae et al., 2019; Brown et al., 2020）。しかし、変換器の計算負荷はコンテキストの長さに応じて大きくなるため、これらのモデルには、メモリに収まる情報量に関する計算上の制限があります。例えば、通常の変換器は、入力シーケンスを処理するのに2次時間を必要とし、新しい単語を生成するたびに文脈に注意を払うのに1次時間を必要とする。この問題に対処するために、いくつかのバリエーションが提案されている（Tay et al.、2020b）。いくつかは、データに依存するパターン（Kitaevら、2020年、Vyasら、2020年、Tayら、2020a、Royら、2021年）またはデータに依存しないパターン（Childら、2019年、Beltagyら、2021年）のいずれかの疎な注意メカニズムを使用することを提案している。 2020; Zaheer et al., 2020）、自己注意の複雑さを軽減する（Katharopoulos et al., 2020; Choromanski et al., 2021; Peng et al., 2021; Jaegle et al., 2021）、過去の表現を記憶にキャッシュする（Dai et al., 2019; Rae et al., 2019）。これらのモデルは、注意の複雑さを軽減し、より長いコンテクストにスケールアップすることができます。しかし、それらの複雑さは依然としてコンテキストの長さに依存しているため、拘束されていないコンテキストを扱うことはできません。

本論文では、∞-former (infinite former; 図1)を提案する。これは、拘束力のない長期記憶(LTM)で拡張された変換モデルであり、このモデルは任意の長さの文脈に注目することができる。 LTMが拘束されないようにするために、我々は連続空間注意フレームワーク（Martins et al.をつかう. このフレームワークでは、入力シーケンスは、放射状基底関数（RBF）の線形結合として表現される連続信号として表されます。この表現には2つの大きな利点があります。 (i) 文脈のトークン数よりも小さな基底関数の数Nを用いて文脈を表現することができ、注意の複雑さを軽減することができる。 (ii) Nは固定することができ、§3.2と図2で説明するように、解像度を失う代償として、注意の複雑さを増加させることなく、O(L^2 + L × N)（ここで、Lは変換シーケンス長に相当する）を用いて、メモリ内の無限の文脈を表現することができる（詳細は§3.1.1）。

古い記憶の解像度が失われるという問題を軽減するために、LTMの新しい信号のより大きなスペースを、以前の記憶の信号の関連領域に帰属させることで、「スティッキー・メモリー」（§3.3）という概念を導入します。これにより、LTMに永続性の概念が生まれ、脳の長期増強と可塑性に類似した、関連する情報を失うことなく、長い文脈をよりよく捉えることができるようになります(Mills et al., 2014; Bamji, 2005)。要約すると、我々の貢献は以下の通りです。

我々は、連続的な長期記憶を持つ変換モデルを拡張した∞-formerを提案する（§3.1）。注意の計算量は文脈の長さに依存しないので、∞-formerは長い文脈をモデル化することができる。
我々は，∞-formerが非限定的な文脈をメモリに保持することを可能にする手続きを提案する(§3.2)．
また、重要な情報をLTMに保持するための手順であるsticky memoryを導入する(§3.3)。
合成タスク(§4.1)では、次第に長くなるシーケンスを考慮し、言語モデリングでは、ゼロからモデルを学習し(§4.2)、事前に学習した言語モデルを微調整する(§4.3)ことで、経験的な比較を行う。これらの実験は，束縛されないメモリを使うことの利点を示している．

2 背景

2.1 トランスフォーマー

トランスフォーマー（Vaswani et al., 2017）は、複数の層で構成されており、多頭の自己注意層（Bahdanau et al., 2015; Graves et al., 2014; Weston et al., 2014）に続くフィードフォワード層を、残差接続（He et al., 2016）と層の正規化（Ba et al., 2016）とともに包含している。入力シーケンスをX = [x1, ... , xL] ∈ R L×e とし、Lはシーケンス長、eは注目層の埋め込みサイズであるとする。自己注意計算に用いるクエリQ、鍵K、値Vは、前の層の入力すなわち出力Xを線形に射影して得られる。

ここで、WQ、WK、WV∈R e×eは学習可能な投影行列である。マルチヘッドセルフアテンションを行うために，Q，K，Vを{1, ... ... , H}内のhに対して，ヘッドQh，Kh，Vh∈R L×dに分割する。ここで、Hはヘッドの数であり、d＝e／Hである。そして、各アテンションヘッドhに対応するコンテキスト表現Zh∈R L×dは、次のようにして得られる。

ここで、dは各Khの隠れたサイズの次元であり、ソフトマックスは行ごとに行われる。頭部のコンテキスト表現を連結して、最終的なコンテキスト表現Z∈R L×e を得る

ここで、WR∈R e×eは、すべてのヘッドの表現を集約した別の投影行列である。

2.2 連続的な注目

任意の連続的な信号を扱うために、連続的な注目メカニズム（Martinsら、2020年）が提案されている。ここでは、単語に対する注目の確率質量関数を、信号に対する確率密度に置き換える。これにより、時間間隔やコンパクトなセグメントを選択することができる。連続的な注目を行うためには、まず、X∈R L×eで表される離散的なテキスト列を連続的な信号に変換する必要がある。これは、基底関数の線形結合として表現することによって行われる。そのために，各xi（i∈{1, ... ... , L}）は，まず，基底関数と関連付けられる。各xiは，例えばti = i/Lとすることで，まず位置ti∈[0, 1]と関連付けられる。そして，任意のt∈[0, 1]に対して，連続空間表現X¯(t)∈R^eを次のように得る．

ここでψ(t)∈R^Nは[0, 1]に位置するN個の1次元RBFである。式4の係数行列B∈R^N×eは、各i∈[L]に対してX¯(ti)≒xiとなるように、多変量リッジ回帰基準(Brown et al., 1980)を用いて得られ、閉形式(詳細はApp.A参照)が導かれる。

ここで、F = [ψ(t1), ... , ψ(tL)] ∈ R N×L は、L個のロケーションの基底ベクトルをパックしたものである。 G∈R L×NはFに依存するだけなので、オフラインで計算することができる。入力シーケンスを連続信号X¯(t)に変換した後、第2のステップはこの信号に出席することです。そのためには、標準的なアテンションメカニズムのように入力シーケンスに対する離散的な確率分布（式2のような）を持つ代わりに、確率密度pを持つことになります。これはガウスN（t: µ, σ2 ）である可能性があり、µとσ2は神経コンポーネントによって計算されます。最後に，pがあれば，文脈ベクトルcを次のように計算することができます。

Martinsら(2020)は、RNNsのための連続的な注目フレームワークを導入した。次のセクション(§3.1)では、それがどのようにトランスフォーマーマルチヘッドアテンションに使用できるかを説明します。

3 無限記憶トランスフォーマー

モデルが長距離のコンテキストにアクセスできるようにするために、我々はバニラトランスフォーマーを、前のステップの入力エンベッディングと隠れた状態を記憶する連続的なLTMで拡張することを提案する。また、transformer-XLの記憶と同様に、LTMと短期記憶（STM）の2つの記憶を持つ可能性も検討しています（Dai et al.、2019）。モデルの図を図1に示します。

3.1 長期記憶

簡単にするために、まず、長期記憶には、層に応じて、過去のテキスト列の入力埋め込みまたは隠れた状態からなる明示的な入力離散列Xが含まれていると仮定する（後に、§3.2で、この考えを非束縛記憶に拡張する）。まず、Xを連続的な近似値X¯(t)に変換する必要があります。そのために、X¯(t)を次のように計算します。

ここで、ψ(t)∈R Nは基底関数であり、係数B∈R N×eは式5のように計算され、B^T＝X^T Gとなります。そして、LTMの鍵K∈R N×e、値V∈R N×eを次のように計算します。

ここで，WK, WV∈R e×eは学習可能な投影行列である。 i in {1, ..., L}の各クエリq{h,i}に対して、注目度スコアを入力とするパラメータ化されたネットワークを用いて、μ{h,i}∈]0, 1[ and σ^2_{h,i}∈R>0を計算する。

次に、連続ソフトマックス変換（Martins et al., 2020）を用いて、確率密度p{h,i}をN (t; µ{h,i}, σ^2_{h,i})として取得する。最後に、value 関数V¯ h(t)をV¯ h(t) = V_h^T ψ(t)として与え、頭部固有の表現ベクトルを式6のように計算します。

これが、アフィン変換を経た行列Z_{LTM,h}∈R^{L×d}の行、ZLTM=[Z{LTM,1}, ... , Z_{LTM,H}]W^Oを形成します。長期表現であるZ_LTMは、変換器のコンテキスト・ベクトルZ_Tと合計され、最終的なコンテキスト表現Z∈R L×eが得られる。

これはトランスフォーマーのフィードフォワード層への入力となる

3.1.1 注意の複雑さ

∞-formerはcontinueousspace attention framework(Martins et al., 2020)を利用してLTM信号上に注意を払うため、そのキーマトリックスサイズK_LTM,h∈R N×dは基底関数の数Nにのみ依存し、注意を払うコンテキストの長さには依存しない。したがって、∞-formerの注意の複雑さも文脈の長さに依存しない。これは、短期記憶も利用する場合はO(L^2 + L × L_STM + L × N)、LTMのみを利用する場合はO(L^2 + L × N)に相当し、いずれも<< O(L×L_LTM)となり、長期記憶に参加するバニラ変換器の複雑さとなる。このため、∞-formerは、計算量を増やすことなく、任意の長いコンテキストに対応することができます。

3.2 束縛されないメモリ

メモリを離散的なシーケンスとして表現する場合、それを拡張するには、新しい隠れた状態をメモリに格納する必要があります。バニラ変換では、メモリの必要量が多いため、長いコンテクストでは実現不可能です。しかし、∞-formerでは、次に説明し、図2に示すように、連続的な注意を用いることで、メモリ要件を増やすことなく、拘束されていないコンテキストにアテンダンスすることができます。 unbounded representationを構築できるようにするために、まず[0, 1]内のM個の位置をサンプリングし、それらの位置でX¯(t)を評価します。これらの場所は、単純に直線的な間隔でも良いし、§3.3で説明したように、領域の重要度に応じてサンプリングしても良い。次に、対応するベクトルを、短期記憶からの新しいベクトルと連結します。そのためには，まず，この関数をτ∈[0, 1]の係数で縮めて，新しいベクトルのためのスペースを確保する必要があります．そのためには、次のように定義します。

そして、M個の位置0 ≤ t1, t2, ... ... , tM ≤ τにおいてX¯(t)を評価することができる。, tM ≤ τとして

そして，これらのベクトルを行とする行列X_past = [x1, x2, ... , xM] >∈RM×eを定義する．その後、この行列と新しいベクトルX_newを連結すると、次のようになります。

最後に，多変量リッジ回帰を行い，式5のようにB^T = X^T Gを介して，新しい係数行列B∈R N×eを計算する．そのためには、Xpastのベクトルを[0, τ ]の位置に、Xnewのベクトルを[τ, 1]の位置に関連付ける必要があります。ベクトルは直線的に配置されていると考えられます。

3.3 スティッキー・メモリー

LTMを拡張する際には、式14に示すように、[0, 1]のM個の位置で現在の信号を評価します。これらの位置は、直線的に配置することができます。しかし、信号の一部の領域は、他の領域よりも関連性が高い場合があり、その結果、次のステップのLTMの信号において、より大きな「メモリスペース」が与えられるべきである。このことを考慮して、各領域での信号の関連性に応じてM個の位置をサンプリングすることを提案する（App.Bの図7参照）。そのためには、前のステップで信号の各区間に与えられた注目度に基づいてヒストグラムを構築する。そのために、まず、信号をD個の線形間隔のビン{d1, ... ... , dD}に分割する。dD}に分割します．そして，各ビンに与えられた確率p(dj )をj∈{1, ... ... , D}について計算する．D}とします。

ここで、Hはアテンションヘッドの数、Lはシーケンスの長さです。なお，式16の積分は，erf関数を用いて効率的に評価できる。

そして、LTMの信号が評価されるM個の位置をp.3.4に従ってサンプリングします。

3.4 実装と学習の詳細

多変量リッジ回帰を適用してLTMの入力離散シーケンスXを連続信号に変換する前に、ゲートとして単純な畳み込み層（ストライド＝1、幅＝3）を使用してシーケンスを平滑化します。

モデルの学習には，クロスエントロピー損失を用いる．言語モデルは，長さLのテキストXの列を入力として，次の単語の確率分布p(xt+1 | xt , ... , xt-L)を出力する． T個のトークンからなるコーパスが与えられたとき、その負の対数尤度を最小化するようにモデルを学習します。

さらに、LTM上での一様な分布を避けるために、LTMに与えられた連続的な注意を正則化し、注意の確率密度N (μh, σh)とガウス事前分布N (μ0, σ0)との間のKL (Kullback-Leibler) divergence, D_{KL}を最小化しました。異なる頭部は異なる領域に注目するはずなので、µ0 = µhとし、注目分散のみを正則化して、次のようになります。

したがって、最小化される最終的な損失は次のようになります。

ここでλKLは、KL正則化の量を制御するハイパーパラメータです。

4 実験

∞-formerが長い文脈をモデル化できるかどうかを理解するために，まず，長い文脈の中でトークンを頻度でソートするという合成タスクの実験を行った(§4.1)。次に，言語モデルをゼロから学習する実験（§4.2）と，事前に学習した言語モデルを微調整する実験（§4.3）を行いました。

ベースライン

実験のベースラインとして，transformer-XL（Dai et al., 2019）とcompressive transformer（Rae et al., 2019）を検討する。 transformer-XLは、バニラ変換器（Vaswani et al., 2017）を、前のステップの隠れた状態で構成される短期記憶で拡張したものである。圧縮型トランスフォーマーは、トランスフォーマー-XLを拡張したもので、短期記憶のほかに、CNNを用いて圧縮された短期記憶の古いベクトルで構成される長期記憶も持っています。 transformer-XLもcompressive transformerも、相対的な位置の符号化が必要です。一方、我々のアプローチでは、メモリ・ベクトルがあらかじめ定義された連続空間の基底係数を表すため、メモリ内の位置符号化は必要ありません。

4.1 並べ替え

このタスクでは、入力は（システムには知られていない）トークン確率分布に従ってサンプリングされたトークンのシーケンスで構成されている。目的は、トークンの頻度が低い順にトークンを生成することである。以下のような例があります。

長期記憶が効果的に使われているかどうか、また、変換器が最新のトークンをモデル化してソートを行っているだけではないかどうかを理解するために、トークンの確率分布が時間とともに変化するように設計します。すなわち、p = αp0 + (1 - α)p1という2つの分布の混合物として設定します。ここで、混合係数α∈[0, 1]は、配列が生成されるにつれて0から1へと徐々に増加します。語彙は20トークンで，長さ4,000，8,000，16,000のシーケンスで実験を行った．すべてのモデルにおいて、3つの層と6つのアテンションヘッドを持つ変換器を使用し、長さ1,024、メモリサイズ2,048のシーケンスを検討しました。

圧縮型の変換器では、両方のメモリのサイズを1,024としました。 ∞-formerでは、サイズが1,024の短期記憶と、1,024の基底関数を持つLTMも考えます。さらなる詳細とハイパーパラメータについては、App.C.1に記載されています。 C.1. 結果は以下の通りです。

図3に見られるように，短い配列長（4,000）の場合，変換器XLは，圧縮変換器や∞-formerよりもわずかに高い精度を達成している．これは、transformer-XLが、ほぼすべてのシーケンスをメモリ内に保持できるためである。しかし，配列の長さが長くなると，その精度は急激に低下する．圧縮型の変換器と∞-formerも、配列長が長くなると精度が低下します。しかし、この減少は∞-formerではそれほど顕著ではなく、これは∞-formerが長いシーケンスのモデル化に優れていることを示している。

4.2 言語モデリング

また、1億300万トークンのトレーニングセットと、217,646トークンと245,569トークンの検証セットとテストセットを持つWikitext-103データセット（Merity et al. そのため、16層と10個のアテンションヘッドを持つトランスフォーマーで構成されるtransformer-XL（Dai et al., 2019）の標準的なアーキテクチャに従っています。 transformer-XLについては、サイズ150のメモリを使って実験する。圧縮変圧器については、両方のメモリのサイズが150で、圧縮率が4であることを考える。 ∞-formerについては、サイズ150の短期メモリ、150のガウスRBFを持つ連続的な長期メモリ、900トークンのメモリ閾値を考える。さらなる詳細とハイパーパラメータについては、App.C.2に記載されています。C.2.

結果

表1に見られるように、長期記憶を用いてモデルを拡張することで、圧縮変換器と∞-formerの両方で、より良いperplexityを得ることができました。さらに、∞-formerはcompressive transformerをわずかに上回っています。また、スティッキーメモリを使用すると、やや低いperplexityになることがわかります。これは、モデルが関連するメモリにフォーカスするのに役立つことを示しています。 ∞- formerがLTM内の古い記憶に注目しているのか、それとも最新の記憶に注目しているのかをよりよく理解するために、検証セットのトークンを予測する際に、長期記憶の各領域に与えられた注目度のヒストグラムをプロットしました。図4に見られるように、最初と中間の層では、∞-formerはより古い記憶に注目する傾向があり、最後の層では、注目パターンはより均一である。図5と図8（App.D）では、∞-formerがtransformer-XLよりも低いperplexityを持つ単語の例を、∞-formerが最後の層のLTMに与えた注意とともに示しています。予測される単語は長期記憶の中に何度も存在しており、∞-formerはそれらの領域に高い注意を払っていることがわかります。また、スティッキーメモリを使用した場合と使用しない場合で、LTMの信号空間が大きくなるかどうかを知るために、スティッキーメモリを使用した場合と使用しない場合で、最後の層の長期記憶に存在する各単語に与えられる記憶空間をプロットした。図6と図9（付録D）には、スティッキーメモリを使用したときに最も大きなスペースが与えられたフレーズや単語の例を示しています。これらの例では、この手順が実際に古い記憶に大きなスペースを与え、時間とともに定着する記憶を作り出すことがわかる。また、これらの記憶は、例示されている語句からもわかるように、関連性のあるものであることがわかります。

4.3 事前に訓練された言語モデル

LTMが事前に訓練された言語モデルを拡張するためにも使用できるかどうかを理解するために、Wikitext103 (Merity et al., 2017)と、訓練セットの最初の2,000冊（≒2億トークン）を含むPG-19 (Rae et al., 2019)のサブセット上で、GPT-2 small (Radford et al., 2019)を微調整する。そのために、サイズ512の配列、512個のガウス型RBFを持つ長期記憶、記憶閾値2,048トークンを考慮します。さらなる詳細とハイパーパラメータについては，App. C.3. 結果表2に報告されている結果は、GPT-2に長期記憶を追加して微調整するだけで、Wikitext-103とPG19の両方でパープレキシティを改善できることを示しています。これは、ゼロから学習することも、事前に学習したモデルを改良するために使用することもできるという、∞-formerの汎用性を示しています。

e4exp / paper_manager_abstract