(論文読み)Sakana AI 自己適応型LLM「Transformer²」


# はじめに

1月15日にフィードされた Sakana AI の「Transformer²」がすごい！

https://sakana.ai/transformer-squared/

けど、仕組みと「どこまですごい！」を捉えきれていない。
どうにか理解したい！ので、論文を読んでみようと思います。

https://arxiv.org/abs/2501.06252



# 自己適応型LLMの概要
* LLMの自己適応性は、リアルタイムでのタスク対応能力を高める重要な進展である
* 従来の一回限りのファインチューニングでは柔軟性と効率性に限界があった

* Transformer²は動的な重み調整により自己適応するLLMシステムを実現した
    * 自己適応モデルは必要に応じて専門モジュールを活用できる
    * 従来のLoRAよりも効率的で高性能な結果を示した


# Transformer²の基本設計

2つのフェーズで説明されています。

**①トレーニングフェーズ**
   - SVFを用いて重み行列の特異値を専門家ベクトルzとして学習 
   - 各タスク(数学、コーディングなど)に特化したzを獲得
   - RL(強化学習)によりタスクパフォーマンスを直接最適化

**②推論フェーズ**
- 2パス方式の推論メカニズム
   - 1stパス: タスク特性を分析し最適な専門家を選択/組合せ
   - 2ndパス: 適応された重みで最終的な応答を生成

## ①トレーニングフェーズ

### SVF：Singular Value Fine-tuning 
* 重み行列W∈ℝn×mの効率的な調整手法
 * SVD分解により3つの成分に分解: $W = UΣV⊺$
   - **半直交行列**：$U∈ℝm×r$、$V∈ℝn×r$
   - **特異値の対角行列(降順)**: $Σ∈ℝr×r$
   - **特異値の数**：$r = min(m,n)$
 * 線形変換を独立項の和として分解
 * 各特異成分は直交する独立の処理を実行
 * 特異値$σi$が各成分の寄与度を制御

### RL：強化学習による最適化
* REINFORCE アルゴリズムを基本フレームワークとして採用
 * $θz = {z1,...,zN×M}$の最適化
   - $N$: レイヤー数
   - $M$: レイヤーあたりの重み行列数
 * 生成された回答$yi$に対する単位報酬 $r∈{-1,1}$
 * KL正則化項により元モデルからの逸脱を抑制

---

うーん。むずかしい。
それぞれの中で何をやっているかまで理解できないですが、SVF と 強化学習を組み合わせていることは理解しました。こんな感じかな？

```mermaid

sequenceDiagram
    title: トレーニングフェーズ
    participant User as 運用担当
    participant SVF as SVF最適化器
    participant RL as RL制御器
    participant Base as ベースLLM
    participant Experts as 専門家ベクトル

    rect rgba(255, 240, 245, 0.5)
        note over SVF,Base: SVD分解による初期化
        User->>SVF: タスクデータセットで初期化
        activate SVF
        SVF->>Base: 重み行列を分解 (SVD)
        Base-->>SVF: U, Σ, V⊺を返却
    end

    rect rgba(240, 248, 255, 0.5)
        note over SVF,RL: 強化学習による最適化
        SVF->>RL: 学習ループ開始
        activate RL
        loop バッチ処理
            RL->>SVF: zベクトル生成
            SVF->>Base: W'を適用
            Base->>Base: 応答を生成
            Base-->>RL: 性能指標を返却
            RL->>RL: 報酬とKL距離を計算
            RL->>SVF: zベクトルを更新
        end
        deactivate RL
    end

    rect rgba(240, 255, 240, 0.5)
        note over SVF,Experts: 学習結果の保存
        SVF->>Experts: 最適化された専門家ベクトルを保存
        deactivate SVF
    end
```

## ②推論フェーズ

# 3つの適応戦略
### プロンプトベースの適応
 * "適応"プロンプトによりタスクを分類
 * 定義済みカテゴリから適切なzを選択
 * "others"カテゴリにより未知タスクにも対応
 * 実装が単純で追加学習不要

### 分類器ベースの適応
 * 専用の分類システムを構築
 * SVFで基本LLMを分類器として微調整
 * 学習済みの専門家タスクからデータセットを構築
 * プロンプト方式より高精度なタスク識別を実現

### Few-shot適応
 * テスト時の追加情報を活用した詳細な適応
 * K個の専門家ベクトルの線形補間: z' = Σk=1^K αk*zk
 * CEMによる最適な補間係数αkの探索
   - Few-shotプロンプトの性能で評価
   - 同点の場合は生成の対数尤度で判断
 * タスクごとに1回の適応で済む効率的な設計
 * 最も柔軟で高性能な適応を実現


こんな感じか...？

```mermaid

sequenceDiagram
    title: 推論フェーズ（2パス方式）
    participant User as ユーザー
    participant Dispatcher as タスク振分器
    participant Experts as 専門家ベクトル
    participant Base as ベースLLM
    participant CEM as CEM最適化器

    note over User,Base: 第1パス：タスク特性分析と専門家選択
    rect rgba(255, 240, 245, 0.5)
        User->>Dispatcher: 入力プロンプト
        activate Dispatcher
    
        alt プロンプトベース適応
            Dispatcher->>Base: 適応プロンプト送信
            Base-->>Dispatcher: タスク分類結果
        else 分類器ベース適応
            Dispatcher->>Base: 専用分類器使用
            Base-->>Dispatcher: タスク分類結果
        else Few-shot適応
            Dispatcher->>CEM: Few-shotサンプルで初期化
            activate CEM
            loop CEM最適化
                CEM->>Experts: α係数生成
                Experts-->>Base: 専門家ベクトル結合
                Base->>Base: 性能評価
                Base-->>CEM: 評価指標返却
            end
            CEM-->>Dispatcher: 最適な組み合わせ
            deactivate CEM
        end
    end

    note over User,Base: 第2パス：適応済み重みによる応答生成
    rect rgba(240, 248, 255, 0.5)
        Dispatcher->>Experts: 選択/結合された専門家ベクトル取得
        Experts-->>Base: 重みの適応を実行
        Base->>User: 最終応答を生成
        deactivate Dispatcher
    end
```
# 所感

- 前発表の[CycleQD](https://qiita.com/ymgc3/items/5e3b9343804576b6e3f4
)のSVDに、RL要素が加わったのが新規要素（トレーニング部分）
- その後ろのクエリプロセスの戦略がさらにセットになった（推論部分）
 
と理解しました。
クエリプロセスの新規性は理解しきれなかったので、次はソースコードを見てみようと思います。

https://github.com/SakanaAI/self-adaptive-llms

「自己適応」というので、リアルタイムでの学習を期待していた部分があったのですが、
トレーニング → 推論は一方向の関係で、推論結果をトレーニングに使用する機能の話ではなかったです。

運用内部的にはやっているとは思いますが、推論時の結果は新たな学習には活用される機構について気になっています。

システムの継続学習（lifelong learning）についても今後発表が出てくるのでしょうね。
今後もSakana AI の発表g楽しみです。

---

# 用語

| カテゴリ | 用語 | 説明 |
|----------|------|------|
| アーキテクチャ | Transformer² | 自己適応フレームワークとして設計されたLLMシステム。2パス推論メカニズムを特徴とする |
| アルゴリズム | CEM (Cross Entropy Method) | モンテカルロ法を用いた重要度サンプリングと最適化のための手法 |
| アルゴリズム | REINFORCE | 方策勾配法に基づく強化学習アルゴリズム |
| アルゴリズム | SVD (Singular Value Decomposition) | 行列を特異値と特異ベクトルに分解する手法 |
| アルゴリズム | SVF (Singular Value Fine-tuning) | 特異値の調整により効率的なモデル適応を実現する手法 |
| パラメータ | KL divergence | 2つの確率分布間の差異を測る指標 |
| パラメータ | α coefficient | Few-shot適応での専門家ベクトルの補間係数 |
| パラメータ | σ (Singular values) | 行列分解で得られる特異値 |
| モデル | Base LLM | 基礎となる大規模言語モデル |
| モデル | Expert vectors | タスク特化型の専門家ベクトル |
| モデル | LoRA (Low-Rank Adaptation) | 低ランク行列による効率的なモデル適応手法 |
| 学習手法 | Few-shot adaptation | 少数のサンプルを用いた適応学習 |
| 学習手法 | Policy gradient | 方策勾配に基づく強化学習手法 |
| 行列演算 | Diagonal matrix | 対角成分以外が0である行列 |
| 行列演算 | Semi-orthogonal matrix | 列ベクトルが互いに直交する行列 |
| 評価指標 | Log-likelihood | 生成されたトークンの対数尤度 |
| 評価指標 | Performance metrics | モデルの性能を評価する指標群 |