今週のTOPIC

メインTOPIC

Transformer-Squared: Self-adaptive LLMs

Abstract

Introduction

• 従来のLLMのPost Trainingには限界がある。
• モデル全体を各タスク向けに再学習させるための計算資源と時間は莫大に必要。
• 一度タスク固有にFine Tuningしたモデルはその設定に静的に固定されてしまい、新たなタスクやドメインに直面すると再度別途微調整が必要になってしまう。
• 複数タスクを一つのモデルで扱おうとすると勾配の干渉による精度のトレードオフが発生したり、Catastrophic forgettingなどの問題がある。

• Post Trainingは高コストかつ静的で、多様なタスクに対応が難しかった。

• 自己適応型LLM
• エキスパートモジュールをオフラインで開発し、必要に応じて Base LLMに拡張することができる。
• 継続的な学習もサポート
• LoRAとMoEを使うのが一般的
• 複数のエキスパートモジュールを作るとトレーニングが必要なパラメータが大幅に増加する。

• 本論文は自己適応型LLMの実現を目指し、動的にタスクに応じた重み調整を行う手法を提案。

• Transformer-Squared（Transformer²）は推論時に二段階の適応を通じて柔軟に対応できるようになった。
• 特異値を用いたFine Tuning手法 SVF(Singular Value Finetuning)を提案
• 2パスの推論メカニズム

• このアプローチは、既存のLoRAやMoE手法よりもパラメータ効率が高く、実用的な適応性を示した。

先行研究

• 自己適応型LLMは以下の2パターンに分かれる
• 複数のLLMが協働するMacroview的な観点
• 単一のLLMが異なるタスクに特化するために内部適応するMicroview的な観点

Macroview

• ドメイン固有の専門知識を持つLLMsにクエリを指示し、エキスパートモデルからの出力を優先することで高い精度とタスク固有の最適化を実現

• メカニズム的には
1. 複数のLLMが異なる役割を果たし、共有された目標に向かって協調する
2. 相互の傾聴と討論を行う
3. 知識とskill planningを統合するために綿密なプロンプトを用いる

Microview

• モジュールやレイヤーのサブセットをタスクに応じて動的にルーティングする

• ルーティング方法の違い
• 従来のMoEは、トークンごとにエキスパートのサブセットをルーティングする
• Transformer²はサンプルごとにモジュールをルーティングする。

• エキスパートモジュールの構築
• 従来のMoEは、ゼロから学習されるか、Upcyclingなどの手法を用いたもの
• 参考：Upcycling → (おそらく)全パラメータを使うモデルを複数組み合わせる手法。SBIのなど。
• Transformer²はRLを用いてエキスパートベクトルを学習し、ドメイン固有知識を獲得する (後述)

• LoRAのような低ランク適応手法はさまざまな修正案が提案されている。

• Transformer²は低ランク行列に依存せず、代わりにフルランク空間にまたがる元のパラメータ行列の特異ベクトルをスケーリングさせる。

• SVDがPEFT (Parameter Efficient Fine Tuning)の帰納的バイアスとして使用されるようになっている。
• 重み行列を分解して、ノイズやロングテール情報に関連する特異成分を用いてLoRAの初期値を決めるなど。

• top-r個の特異ベクトルで元の重み行列を近似する先行研究などもある。
• top-r部分空間内の大きさと向きを調整するために、top-r部分の特異値行列の上に小さな学習可能な行列が導入される。
• この場合、特異値分布の歪みが少ない場合に、上位だけだと重要な情報が失われる可能性がある。

手法

1. 特異値分解に基づくファインチューニング (SVF)

2. Transformer²フレームワークにおける推論時の自己適応機構

Singular Value Fine-tuning, SVF

• 従来のFine Tuning手法は、重み行列を修正することで、事前に訓練されたモデルを新しい機能で補強することを目的としていることが多い。

• しかし、先行研究では既に大規模なLLMであれば多くの下流タスクを解くために必要な能力は、既に事前学習済みモデル内にすでに存在している。

• なので、効率的なFine Tuningを行うためには、これらの潜在能力をより表現しやすくすることに焦点を当てるべきである。

• これらの観点からTransformer²では、という単純なベクトルを学習させる。これを用いて、特異値の対角行列をのように調整する。これにより、LoRAなどと比べても桁違いに少ないパラメータで調整可能となる。
• 表現能力が少なくなりそうではあるが、逆にフルランクの情報にアクセスできるという利点がある。
• 特異値のみを修正するという操作なので、正則化的な効果もあり、Catastrophic forgettingも起きにくい。

End-to-end optimization with RL

• 最適化するパラメータはN Layer x M Matricesの重み行列に対する特異値ベクトル

• 単純にpromptを 、生成された回答を 、報酬を として強化学習する。その際に、元のモデルから逸脱しすぎないためのKLペナルティも追加する

• RLは一般的にNext Token Predictionより安定性が低いが、SVFは正則化特性が強いので、先行研究のような失敗を多く回避できることがわかった。

• LoRA fine tuningの場合は、Next Token Predictionをする必要があり、explaining textsが必要になるが、SVFではそれも必要ない。

Self-Adaptation

• SVFで訓練されたK個のエキスパートベクトル を組み合わせて、異なる種類の能力を提供する。
• 能力と訓練データセットの間のマッピングはデータセットのメタデータで行う。

• Transformer²はプロンプトが与えられると、一度目の推論を行い、そのtest-timeの挙動に合わせて、新しい ベクトルを生成する。このを2回目の推論で使用することで、新しく適応された重みで実際の応答を生成する。

• この適応プロセスを行うために以下の3つの手法を試した。
1. Prompt engineering
• Adaption promptを作り、LLMに直接入力プロンプトを分類してもらう。
• その分類結果に基づいて、事前学習に使用したドメイントピックのセットから1つのカテゴリを抽出し、対応するを直接抽出する。
• 「その他」カテゴリも用意することで、そのままの重みを使うことも選択させる。



2. Classification expert
• タスク識別を行うために特別なシステムを用意する。
• SVFを訓練したデータセット を用いて、元のベースモデルのLLMがタスクを分類できるようにFine-Tuningする。
• この際にもSVFの仕組みを使い、エキスパート を生成し、最初の推論パスで をそのまま利用する。
3. Few-shot adaptation
• Few-shotの例とCross-entropy method (CEM)を用いて のように線型結合させたエキスパートベクトルを生成する。
CEMの疑似コード



• どのような評価関数を使っているか明示的に論文内で書かれていないが、コード的にはこの辺り
• 実際推論するデータセットのラベルとかを使ってそうな雰囲気。そりゃ強くなるだろうという気はする。

実験

• SVFの効率と有効性の評価

• 提案された3つの適応戦略による自己適応性の実証

• 新しいフレームワークの特性の理解と解釈

実験設定

• 一般性のために異なるモデルファミリーとサイズのモデルを選定
• LLAMA3-8B-INSTRUCT
• MISTRAL-7B-INSTRUCT-V0.3
• LLAMA3-70B-INSTRUCT.

• 以下の3つのデータセットからベクトルを得る
• GSM8K：小学生レベルの数学文章題データセット。評価指標は厳密一致の正解率（Accuracy）
• MBPP-pro：Pythonプログラミング問題集（GoogleのMBPPを拡張したもの）。コードが正しく実行・解答できた割合で評価。指標はPass@1（1回の出力で正解できる率）
• ARC-Easy：科学常識問題集（AI2 Reasoning Challengeの簡単版）。4択問題の正解率で評価。

• VLMドメインへのSVFの適用性を評価するためにTextVQAの言語backboneとして を使う場合のzセットとしても利用する。

• Fine-tuningの学習曲線は以下のようになった。
• 点線はLLAMA3-8B-INSTRUCTの性能
• 赤い丸の部分でEarly stoppingして後段に使用している。
• SVFの学習効率が良いこともアピールポイント (ちゃんと過学習されている)

• 最後に4つの未知タスクに対して評価する。
• MATH：数学競技レベルの難易度の高い数学問題データセット。評価指標は正解率。
• Humeneval：オンラインジャッジ形式のコード生成タスク（OpenAI提供）。指標はPass@1。
• ARC-Challenge：ARC-Easyの難問版
• OKVQA： 常識的知識を要する画像質問応答タスク（Open Knowledge VQA）。画像内容＋外部知識による応答の正確さで評価（Accuracy）。

実験結果

• Fine-tuning段階のテストセットでは、一貫した精度向上が見られた。
• (Early Stoppingしてるのだからそれはそうでは？)

• 未知タスクにおける性能
• LoRAと公平に比較のために、検討したトレーニングタスクの全てのチェックポイントを使用して、テストタスクの最高性能のみを報告した。
• Transformer²はほとんどのタスクで精度向上が見られるが、LoRAの場合はMATHやHumanevalで性能が大幅に悪化する。
• これはLoRAが小さなデータセットであるGSM8KやMBPP-Proなどのデータセットで学習した場合のオーバーフィットに敏感である可能性を示唆している。
• 上記のFigure5のVQAの場合にも同様の傾向が見られる。

• 適応戦略の比較
• 明確な単調増加が観測される。特にFew-shot手法がどのタスクでも最も効果的である。

• 推論速度
• ARC-Challenge以外は、1段階目の推論はボトルネックにはならない。
• ARC-Challengeはinput tokenが長くなりがち

分析

• Classification系のAdaptation戦略の分類精度
• Reasoningの精度がClassification Expertの方が高い。
• Mistral-7BやLlama3-70Bはプロンプトだけでも精度が高い。
• ただし、分類精度が高いからといって後段のタスクでも精度が高いとは限らないことに注意。

• Few-shot Adaptationの係数分析
• HumenEvalやArc Challengeでは、どのモデルでも関連するタスクでのエキスパートの寄与度が高いが、MATHだけ寄与度が他のエキスパートの方が大きい。
• これは訓練で使ったGSM8Kが小学生の問題であり、MATHとは求められる能力が異なり、論理的推論などが必要になるからではないかと考察される。

• Ablation Studies
• Module Sensitivity：異なるモジュールにSVFを適用した場合の性能を比較 (#1~3)
• Objective Function: RL or Next Token Prediction (#2, #4)
• SVF vs LoRA: LoRAにおけるRLも検証。

• Cross-model compatibility
• 異なるLLMに対しても同じエキスパートベクトルを適用できるかを試した。
• Llama3-8B → Mistoral-7Bで実験
• 以外にも二つのタスクでは正の影響があった。二つのモデルのアーキテクチャが類似していることに起因している？

結論

• Transformer²は既存の自己適応型LLMに比べて以下のような利点がある。
• コスト削減
• 高い構成性
• オーバーフィットに対する正則化
• 優れた性能

• また、3つの適応戦略を開発し、Test時にアクセスするログを増やすことで単調に性能が増加することを明らかにした。

感想

• 意外と中身自体はシンプルな構成だった。

• MoEとかSelf Adaptationとかが流行ってはいるが、タスクの種類が実務的ではないので、実際実務で使うとしたらどのようにタスクを分解するべきかが難しい。実際のところ1つのエキスパートを選ぶ系の適法手法であれば、使う場所によって単純に使うLLMやプロンプトを変えるだけで良いのでは？という気がする。
• この辺りがどんどん省力化していくと、マルチテナントSaaSでテナント毎にエキスパートベクトルを用意する未来とかもあるのか