LongContextへの対応に向けたマルチスケール状態空間モデル

現在の大規模言語モデル（LargeLanguageModel;LLM）の多くに組み込まれている注意機構[1]には，計算量が入力の長さの2乗に比例して大きくなるという問題がある．実際には，スライディングウィンドウを用いた注意機構による計算量の削減や，並列処理による生成に必要な所要時間の短縮が行われているため，計算量の問題は重要視されにくい．しかしながら依然としてLLMの計算量は大きいため，性能を維持したまま計算量を小さくする研究が盛んに行われている．LLMの性能を維持したまま計算量を小さくするためのアプローチの一つとして，状態空間モデル（StateSpaceModel;SSM）がある．注意機構では各トークンに対してそれ以前のトークンを利用した計算をしていたが，SSMでは決められたサイズの行列を入力されたトークンの数だけ更新していく形で計算を進める．そのため，SSMの計算量は入力の長さに比例しており，長い入力の処理する場合注意機構よりも計算量を抑えやすい．特にMamba[2,3]というSSMを基にしたアーキテクチャは，言語モデルPythia-410mをPhonebookベンチマークで評価したときの正解率．として用いた場合，一部の評価ベンチマークにおいてTransformer[1]をはじめとする注意機構を含む言語モデルに匹敵する性能を示している[2,3]．一方で，数千トークン以上の入力を必要とする長文処理ベンチマークでは，Mambaの性能が著しく損なわれ，Transformerに対する計算量の優位性を活かせないことが問題となっている．これはSSMの原理として，隠れ状態と呼ばれる固定の大きさの行列に入力の情報を保持するように更新していくが，入力が長すぎると隠れ状態の情報の保持に問題が発生してしまうためだと考えられている．実際に，Phonebook[4,5]という入力に含まれる電話帳から特定の人物名を探して隣に書かれた電話番号を出力するといったベンチマークにおいて，電話帳の人物名と電話番号のペア数を変えながら正解率を測定したところ，図1のような結果となった．電話帳のペア数が100を超えると正解率はほぼ0になっていることが分かる．Mambaの長い入力に対する性能を向上させるために，既存研究のLongMamba[6]では入力情報の圧縮というアプローチが行われている．これは，Mambaが長い入力に対して隠れ状態の崩壊を招くため，長い入力を避けて短くしてから処理すればよいという考えに基づいている．実際，圧縮による情Thisworkispublishedwithoutpeerreviewandislicensedbytheauthor(s)underCCBY4.0(https://creativecommons.org/licenses/by/4.0/).org/licenses/by/4.0/).