StreamingLLM - アテンションシンクによる効率的なストリーミング言語モデルの実装
(github.com/mit-han-lab)- StreamingLLMは、効率性や性能を犠牲にすることなく、LLMを無限長入力のストリーミングアプリケーションに展開するためのフレームワーク
- 長い相互作用が必要なマルチラウンド対話では、過去トークンのKey/Value状態キャッシュが大きなメモリを消費し、一般的なLLMは学習シーケンス長より長いテキストへ一般化できないという問題を扱う
- 最近のKVのみをキャッシュするwindow attentionは、テキスト長がキャッシュサイズを超えると破綻するが、StreamingLLMは初期トークンのKVを保持するattention sinkによってwindow attentionの性能を大きく回復する
- 有限長のattention windowで学習されたLLMが、fine-tuningなしで無限シーケンス長へ一般化できるようにし、Llama-2、MPT、Falcon、Pythiaで最大400万トークン超まで安定かつ効率的な言語モデリングを実現
- ストリーミング設定では、sliding window recomputationベースラインと比べて最大22.2倍の高速化を示す
- コンテキストウィンドウを拡張せず、最新トークンとattention sinkのみを保持し、中間トークンは破棄する
- Llama-2が4096トークンのコンテキストウィンドウで事前学習されている場合、StreamingLLMにおけるLlama-2の最大キャッシュサイズも4096
- 長い本を入力した場合、モデルは最新トークンしか認識しないため、結論部分だけを要約できる
- 適したユースケースは、継続動作が必要で、過去データへの依存や大きなメモリ使用を避けたいストリーミングアプリケーションであり、例としてマルチラウンド対話やLLMベースの日常アシスタントが挙げられる
- 最近のコンテキスト拡張手法と直交的で統合可能であり、StreamingLLMの文脈におけるcontext extensionは、より大きなキャッシュサイズでより多くの最新トークンを保存できる可能性を意味する
- 実行例は
examples/run_streaming_llama.py --enable_streamingで、環境設定にはPython 3.8、torch、transformers==4.33.0、accelerate、datasets、evaluate、wandb、scikit-learn、scipy、sentencepieceが使われる - コアコードはLlama-2、MPT、Falcon、Pythiaを含めて公開済みで、perplexity評価コードとStreaming Llama Chatbotデモも公開済み、StreamEvalデータセットと評価コードは未公開
1件のコメント
Hacker News のコメント
この研究を完全な密なアテンションのようなものだと誤解しているように思える。
ここで言っているのは、遠く離れた内容を見られるようにするという話ではなく、スライディングウィンドウを毎回再計算して L² のコストを T 回払う代わりに、キャッシュを再利用してパープレキシティを維持するという効率改善に近い。
テストも
Q A Q A Q A Q A...のようにつなげて測ったもので、Q Q Q Q A A A A...のように、かなり後ろにある答えを探させたわけではない。パープレキシティの測定は「読めるテキスト」、つまり局所的にもっともらしい文を作れるという意味であって、アテンションが届かない大きな三角形状の空白から何かを「抽出」している証拠ではない。
本を渡して各段落の最初の単語を書かせたり、各章を一文で要約させたりしたら失敗しそうだ。
自分で試したところ、コンテキスト長の拡張をしてくれるようには見えず、実行は速いほうだった。
A100 のメモリを約 35GB 使用し、実行中ずっと使用量は固定されていた。
Project Gutenberg から本を持ってきて段落単位に分け、1つずつ入力して各段落ごとに「okay」と答えさせたあと、最後に質問したところ、答えは完全に幻覚だった。
ちなみに 10 分ほど触っている間、デフォルトモデルの
lmsys/vicuna-13b-v1.3に英語で答えさせることもあまりうまくいかなかった。https://gist.github.com/bluecoconut/9cae9e91fe3b1616ed650a96...
それでも興味深い研究で、核心は Figure 2 の発見にあるように見える。
最初の2つのレイヤーは直近のトークンにより注目する局所的なパターンを示すが、その下を通過すると、モデルはすべてのレイヤーとヘッドで初期トークンに強く注目する。
著者らはこれを「attention sinks」と呼んでおり、意味的に重要でなくても Softmax のためにアテンションスコアの合計が 1 にならなければならないので、余ったアテンションはどこかへ行く必要がある、と見ている。
自己回帰型言語モデルでは、初期トークンはほぼすべての後続トークンから見えるため、このようなシンクの役割を担うよう学習されやすい、という説明だ。
StreamingLLM は、LLM のアテンションウィンドウを切り出して使うときに生じるこの奇妙な挙動を補正する「ハック」に近く、Softmax 利用の亀裂が露呈した例なので、コンテキスト長に柔軟な LLM が欲しいなら、別の関数のほうがよいかもしれないと思わされる。
最初にざっと見ると、良すぎて本当か疑いたくなるほどだが、研究の質は悪くなさそうで、手法も驚くほど単純だ。
アイデアは、各レイヤーで最初のトークンとスライディングコンテキストウィンドウにだけアテンションを適用し、その間のトークンは無視するというもの。
これは、各レイヤーが関連情報をシーケンス後方へ少しずつ押し上げ、最上位レイヤー末尾のスライディングアテンションウィンドウがそれを見られるようにする、という意味に見える。
ただし、すべてのスライディングウィンドウがカバーする範囲がシーケンス全体をつなぐのに十分でなければ、重要な情報をすべて前方へ伝えられない可能性がある。
たとえば、すべてのウィンドウ長が同じときに
モデルの深さ × ウィンドウ長 < シーケンス長なら限界が生じる。これは、Softmax が合計を 1 にしなければならないという観察のおかげで可能になったように見える。
ざっと見たところ、モデルは以前のトークンに注目する必要がないとき、最初のトークンをプレースホルダーのように使う傾向がある。
この問題を初めて見たのは Evan Miller の HN 投稿で、アテンションヘッドに必ず以前のトークンへすべてのアテンションを配分させるのは間違いで、Softmax の分母に 1 を足して「注目しない」ことを許すべきだ、という内容だった。
再学習なしでこの観察を活用した点は良いし、Evan の提案に従っていたらモデルがどう変わっていたのかも気になる。
[2] https://news.ycombinator.com/item?id=36851494
すべての値が 0 の専用シンクトークンを置いてモデルを学習させたが、それでも他の初期トークンがシンクとして使われるようになるため、専用シンクトークンがあるほうがよい、という結論のように見える。
私の知る限り、実際の性能をより良くしたことはなかった。
学習後にアテンションウィンドウを操作する場合には Softmax + 1 のほうが合う可能性もあるが、大規模に試した人がいるかは分からない。
アテンションキャッシュメモリを追加するのは、この問題に対するとても興味深い解法だ。
数日前、Vision Transformer で関連する観察を行った論文も出ていた。
Transformer モデルはグローバル情報を保存するトークンを選んでいるように見え、一種の「考えるためのトークン」が必要に見える。
この目的の特定トークンを提供すると性能が少し良くなり、説明用の可視化もかなり興味深いものになる。
[0] https://arxiv.org/pdf/2309.16588.pdf
微調整では元のモデルパラメータを固定し、新しい「チューニング」キャッシュユニットに出入りするパラメータだけを調整できる。
そうすれば、異なるチューニングユニットのセットを入れ替えたり、一緒に使ったりできる。
俗語回避ユニット + 特定用語ユニット + 簡潔に書くユニットのように、一種のスーパープロンプトを混ぜて使う方式だ。
新しいパラメータ数が十分に少なければ、メモリは余計に使うが、高次最適化で高速かつ効果的なチューニングも可能かもしれない。
学習中にシーケンス長とユニット数を一緒に増やす方法も考えられる。
短いシーケンスには数個のユニットだけを使い、学習シーケンス長を伸ばしながらユニットを追加して継続学習する、という形だ。
任意のスケジュールではなく、性能や勾配分析によってキャッシュ拡張を制御することもできそうだ。
著者たちがFAQを公開しており、混乱をある程度整理してくれるかもしれない: https://github.com/mit-han-lab/streaming-llm/blob/main/READM...
「本のような長いテキストをStreamingLLMに入れて要約できるか？」については、長いテキストを入力することはできるが、モデルは最新のトークンだけを認識するため、本を入れても最後の数段落しか要約できず、あまり有用ではないかもしれない、と答えている
つまり、LLMのコンテキストウィンドウを拡張したり長期記憶を強化したりするものではなく、StreamingLLMの強みはキャッシュをリフレッシュせずに直近のトークンから流暢なテキストを生成することにある
自分が間違っている可能性もあるが、人々が考えているように、学習時の長さを超えた内容をLLMに参照させる技術ではないように思う
おそらく、長いテキストでモデル性能を維持する問題、正確には依然としてコンテキストウィンドウ内にある内容に対する性能に近いように見える
モデルがテキスト冒頭のトークンのアテンションにある種の荷物を載せるよう学習しており、それがウィンドウの外へ消えると壊れる、という説明だが、なぜそうなるのかはよく分からない
命令入力でないなら、中間のテキストも冒頭のテキストと同じくらい良いのではないかと思う
こうしたスライディングウィンドウ手法は、予期しない命令が最後にだけ出てくる場合をどう扱うのか気になる
たとえば本をモデルに入力し、最後の文に「以前の入力に含まれる文字 m の数を返せ」という命令があるとする。人間ならため息をついて読み直しながら数えるだろうが、LLMには入力をさかのぼって読み直す能力がない
この例では、文字を数えること自体のLLMの限界は無視するとしても、本当に解決するにはLLMが任意にループしたりジャンプしたりできる必要がありそうだ
もちろん、そうなるとまったく新しい問題が生じ、そもそも新しいアーキテクチャが必要になるかもしれない
質問を関連キーワードに変換し、調べ、内容をもう一度ざっと見て関連情報を探す、という方式だ
必要な前処理があるなら、「行って十分に調査してから答える」LLMは非常に強力になり得る
過去約1万年にわたり、私たちは個人の脳の容量と時間を超えられるよう知識管理技術を改善してきたのだから、言語モデルにも単なるBing検索ではなく、実際の研究手法と事前の消化を活用させるべきだ
短期記憶が、あるコード片が何をしたかを覚えている必要はなく、読むときにタグ付けし、拡張可能な共有タグ索引に依存すればよい
ただ、考えれば考えるほど一般的なLLMの事前学習に似ており、知識索引が巨大なLLMの重みの塊であるかのように感じられる
LLM自体の動作を変えるというより、その上に載せるレイヤーに近い
自分が誤解しているのかもしれないが、この場合は隠れ状態が検索問題を解決するのではないか？
答える前にはいずれにせよ入力全体を吸収しなければならないので、命令が先頭にあろうと最後にあろうと、アテンション以外には大きな影響はなさそうに見える
Claude 100Kはユーザーに命令を最後に置くよう求めている
あるいは高速なモデルで最後に命令があるか確認し、前の方へ持ってくればよい
現在のモデルが短い入力でもこれをできるのか、よく分からない
少し冗談めかして言えば、LLMはRNNを再発明しようと本当に懸命に努力していて、道具さえ与えれば最終的にそうするのだと思う
別の見方をすれば、Transformerモデルはリソース制約があるときに、RNNネットワークのどの部分を「保存する価値があるか」予測しようとしている
現在のTransformerは単純なヒューリスティックを使っており、この結果はそのヒューリスティックをより良くする
多くのNP完全問題のように、完全に正確ではなくても有用な近似があり得るし、Transformerはニューラルネットワークでもそれが可能であることを示している
オープンソースのリーダーボードでしばらく中位あたりにいたので、かなり正当なアプローチであり、単に流行していないだけだ
[1]: https://huggingface.co/blog/rwkv
TransformerがRNNに対して持つ主な利点は学習の並列化だ
RNNは学習中に勾配消失があり、全体の利用率を高めるのも難しく、大きなバッチが必要になるため扱いが難しい
RWKVのようなモデルの存在は、Transformerのように学習し、RNNのように推論する未来があり得ることを示している
関連して、MITのHan教授が公開のTinyML講義を進めている
https://news.ycombinator.com/item?id=37620507
https://efficientml.ai