Transformer²: Self-Adaptive LLMs

これは何

この記事では、最新のAI論文Transformer²: Self-adaptive LLMsをわかりやすく解説します。この論文は、AI研究の最前線であるTransformerモデルに関するもので、特に大規模言語モデル（LLM）がリアルタイムでタスクに適応する能力を強化する新しいフレームワークを提案しています。この研究はSakana AIという組織から発表されており、現在のLLMが抱える課題を解決することを目指しています。

LLM（大規模言語モデル）は、多くのタスクで優れた性能を発揮しますが、柔軟性や効率性の面でまだまだ改善の余地があります。例えば、新しいタスクに適応するには時間や計算リソースが必要だったりします。この論文では、それを克服するための手法として、モデルの内部をリアルタイムで調整する“自己適応”の仕組みが提案されています。これにより、計算コストを抑えつつ、さまざまなタスクに対応できる未来のLLMの可能性が広がります。

Transformer²とは

Transformer²（トランスフォーマースクエア）はLLMに自己適応機能を持たせるためのフレームワークです。LLMは新たなタスクを学習するのが人間よりもずっと苦手で無駄に多くのデータと計算が必要になります。Transformer²は“重み行列の特定の成分”（特異ベクトル）のみを動的に調整することで、モデル全体を再訓練することなく、新しいタスクに柔軟に適応できるように設計されています。

Transformer²は２つの革新的な技術で構成されます。

• エキスパートベクトルの適用: タスク固有のエキスパートベクトルを使ってモデルの動作をリアルタイムで微調整します。
• ディスパッチシステム: タスクの特定を分析して採用するエキスパートを決定します。

これによりTransformer²は計算コストを大幅に削減しながら従来のLoRAなどのアプローチよりも優れたタスク適応を実現します。

エキスパートベクトル

1. 少量データへの過学習を防ぐため
2. 計算コストを減らすため
3. 付け替え可能なエキスパートを学習するため

Transformer²では重み行列を特異値分解します。特異値は行列の代表的な性質を表しているので編集することで効率よくモデルのふるまいを変更できそうな感じがします。そこでタスクに応じて専用の特異値の列＝エキスパートベクトルを学習することを提案しています。
学習するパラメータが大幅に少なくなるので過学習への不安と計算コストがぐっと減りますね！

ディスパッチシステム

問題は個別に学習したエキスパートのどれを使うかです。この論文ではプロンプトに対して事前分析行うことでこれを実現します。といっても非常にシンプルでヒューリスティックな手法になっていて以下の３つのパターンが提案されています。

プロンプトエンジニアリング
直接LLMにプロンプトのカテゴリを尋ねることで利用するエキスパートを選択します。直感的。
選択肢には事前に学習したk個のエキスパートの他に「その他」も用意されていて、「その他」が選択された場合にはファインチューニングされていないベースモデルを利用します。

分類エキスパート
アイデアはプロンプトエンジニアリングと一緒でLLMに直接どのエキスパートを利用すればよいかを尋ねます。入力のプロンプトをk個のエキスパートに分類する専用のエキスパートベクトルを追加で学習します。幸いなことにk個のエキスパートの学習に利用したデータセットがあるので、これを分類タスクの学習データとすることができますね。

Few-Shot適応
ここでは単一のエキスパートを使うのではなくエキスパートの線形補間を使って全く新しいエキスパートベクトル$z′z'$を用意します。

$$z′=Σk=1Kαkzk z' = \Sigma^K_{k=1}{\alpha_k z_k}$$

この重み$αk\alpha_k$を得るためにはプロンプトと似たような問と答えの例（Few-Shot Examples)およびそれを評価することのできるシステムが必要です。CEM（Cross-Entropy Method）と呼ばれる手法を使いますが、要はランダムに生成したウェイトの中からFew-Shot Examples上で最も性能が良かった重みを利用します。

RLでの学習

エキスパートベクトルの学習には教師あり学習ではなく強化学習を使います。目的関数には1992年に発表された古典的なREINFORCEの目的関数にもとのモデルからの乖離を抑えるKL正則化を加えたものとしました。

$$J(θz)=E[log⁡(πθW′(yi^∣xi))r(yi^,yi)]−λdKL(πθW′∣∣πθW) J(\theta_z) = \mathbb{E}[\log(\pi_{\theta_{W'}}(\hat{y_i} | x_i))r(\hat{y_i}, y_i)] - \lambda d_{KL} (\pi_{\theta_{W'}} || \pi_{\theta_W})$$

第一項はごついような気もするのですがモデルの確率ベクトルと報酬の掛け算です。つまり報酬のいいアクションの確率は高くしろ！ということですね。直感的。

ここまでTransformer²のアルゴリズムを解説しました。思ったよりシンプルですよね。以下にはこの手法の性能評価結果を記述します。

実験

基本：実験にはLlama3-8B-Instruct, Mistral-7B-Instruct-v0.3およびLlama3-70B-Instructを利用します。またタスクとしてGSM8K, MBPP-pro, ARC-easyを考えます。
おまけ：マルチモーダルなタスクでの性能を図るためにLlama3-8B-Instructを使ってTextVQAタスクも解かせてみます。

①３つのタスクに特化したエキスパートの用意と②新たなタスクを適切なエキスパートにディスパッチできるかがポイントですね。またParameter Efficientな学習方法の筆頭であるLoRAもベースラインに加えてあります。

エキスパートの学習

ベンチマークデータセットを使ってモデルをファインチューニングしました。ほとんどの場合でLoRAと同等以上の性能を示せています。LoRAがいくつかのタスクでファインチューニング後に性能が低下しているのは過学習が発生しているからです。
特異値飲みを扱うSVFに比べて学習パラメータ数が多いLoRAは過学習に敏感ということが示されました。個人的にはLoRAのランクを大幅に下げて学習パラメータ数を減らした場合にどうなるのかは気になります。画像タスクにおいてもSVFはBase Model, LoRAを圧倒していますね。

ディスパッチの性能

Unseen Taskでの性能は予想された通り、一人だけFew-Shot ExamplesにアクセスしているTransformers²（Few-Shot）が圧倒する形になりました。

ちなみにタスクの分類についてはこの３つのタスクですともとのモデルやClassification Expertでかなり高い性能が出るようです。Prompt EngineeringとClassification Expertの性能差が軽微なのも納得ですね。

最後に

学習データがスケールしなくなった昨今では推論時の工夫で性能を向上させる試みに注目が集まっています。この論文の推論時に動的にエキスパートを付け替えるアイデアもとても興味深く、欧米のように大規模モデルを訓練するのが難しい日本の企業にとって注目の技術になるでしょう。それが日本に籍を置くSakana AIから出てきたのもうれしいですね！今後の研究活動も応援しています。