Evolution Strategies Fine-Tuning: 大規模言語モデルの直接パラメータ最適化

Outcome

人口ベースの直接パラメータサーチを通じて大規模言語モデルを微調整し、多様なアーキテクチャ全体で堅牢なモデル改善を実現します。勾配ベースのRLメソッドと比べて15.5倍低い訓練分散を達成し、明示的なペナルティなしで報酬ハッキングに対する耐性を提供します。

Problem Context

現在のLLM微調整は、勾配ベースの強化学習(PPO、GRPO)を通じた逆伝播に依存していますが、以下の課題があります:

• 疎で長期の報酬: 推論タスクでは中間的な教師信号が利用不可であることが多く、長いシーケンスを通じた勾配が不安定になる
• 報酬ハッキング: 勾配ベースの最適化は明示的なKL制約なしでループホール(短いが無意味な出力)を利用する
• クロスモデル脆弱性: 微調整の成功は基本モデルアーキテクチャ全体で劇的に異なり、GRPOは特定のモデルで完全に失敗する
• 訓練不安定性: 実行間での高い分散(ESより15.5倍高い)は、大規模デプロイメント向けの高価な微調整を信頼できなくします
• 計算オーバーヘッド: 逆伝播とKLペナルティ計算は相当なメモリと計算負荷を追加する

Evolution Strategiesは代替手段を提供します:報酬信号のみを使用した直接パラメータ空間サーチであり、勾配は不要です。

Core Concept

Evolution Strategiesはモデルパラメータを進化圧にさらされるゲノムとして扱います。アルゴリズムは繰り返し:

1. 正規分布からパラメータ摂動をサンプリング
2. 摂動されたモデルを目的のタスクで評価して報酬を取得
3. 高報酬摂動の方向でパラメータを更新(自然勾配)

重要な知見: ESは報酬値のみを必要とし、勾配は不要です。これにより応答レベルの教師信号(モデルが問題を解決したか?)が可能になり、損失勾配ではなく、最適化をモデルアーキテクチャから切り離し、疎報酬状況での効果的なサーチを実現します。

十億パラメータ規模では、7つのエンジニアリング最適化がESを実行可能にします:ランダムシードを通じたノイズ再現可能性、並列GPU評価、インプレース摂動、報酬正規化、貪欲デコーディング、分解された更新、および簡略化された学習率。

Architecture Overview

人口ベースのサーチ

• 小さな固定人口(30メンバー vs. 先行研究での10,000以上)が摂動を並列で評価
• 各メンバー: 基本重みにシードからサンプリングされたスケーリングされたガウスノイズを追加
• GPU全体での並列評価;Hugging Face Accelerateを通じて単一マシンまたは分散クラスタ

報酬駆動パラメータ更新

• 各人口メンバーから報酬信号(スカラー、遅延OK)を収集
• 報酬を零平均単位分散に正規化
• ユーティリティ加重摂動の平均を計算: Δθ ∝ Σ(utility_i × noise_i)
• 学習率を適用: θ_new = θ_old + α × Δθ

メモリと計算効率

• ノイズ取得: ランダムシードからオンザフライで摂動を再構築(ストレージオーバーヘッドなし)
• レイヤーレベルのインプレース摂動: 重みを順次修正、評価、復元(メモリ内で単一コピー)
• バッチGPU評価: スレッドを通じてGPイあたり複数の摂動されたモデルを評価
• 逆伝播なし: 勾配メソッド vs. ~50%メモリ削減

安定性特性

• ESの更新はランクベースのユーティリティ加重(報酬外れ値とスケールに対してロバスト)
• 明示的なKLペナルティなし; ESは人口多様性を通じて報酬ハッキングを自然に回避
• 分散削減: 同一の問題でGRPOを通じた実行全体で15.5倍低い

Implementation

1. 環境セットアップ

Python環境を準備し、分散GPU評価用の依存関係をインストールします。

# 仮想環境を作成およびアクティベート
python3.10 -m venv es_env
source es_env/bin/activate

# 依存関係をインストール(リポジトリから)
pip install -r requirements.txt

# 主要パッケージ:
# - torch>=2.0.0
# - transformers>=4.40.0
# - accelerate>=0.27.0 (分散訓練)
# - datasets>=2.18.0 (データ読み込み)
# - numpy, pandas (ユーティリティ)

2. 報酬関数を定義

報酬関数はモデルを受け取りスカラー値を返します。ESはこれを直接最適化します。勾配は不要です。

def compute_reward(model, tokenizer, examples):
    """
    タスクに対してモデルを評価しスカラー報酬を返す。

    Args:
        model: LLMインスタンス(既にロード済み)
        tokenizer: モデル用のトークナイザー
        examples: {input, expected_output}辞書のリスト

    Returns:
        float: 集約された報酬(0-1範囲推奨)
    """
    correct = 0
    for example in examples:
        # 貪欲デコーディングで応答を生成
        inputs = tokenizer(example["input"], return_tensors="pt").to(model.device)
        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=256,
                do_sample=False,  # 貪欲
                pad_token_id=tokenizer.eos_token_id
            )
        response = tokenizer.decode(output[0], skip_special_tokens=True)

        # 正確性をチェック(タスク固有)
        if is_correct(response, example["expected_output"]):
            correct += 1

    # 正解率を返す
    return correct / len(examples)


def is_correct(response, expected):
    """タスク固有の正確性チェック。"""
    # 例: 完全一致
    return response.strip() == expected.strip()

3. 人口と状態を初期化

ES状態を設定します:平均パラメータ、ステップサイズ、および人口ユーティリティ。

import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer

# モデルとトークナイザーをロード
model_name = "Qwen/Qwen2.5-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# パラメータを単一ベクトルに平坦化(ES状態用)
params_init = torch.nn.utils.parameters_to_vector(model.parameters()).detach().clone()
num_params = params_init.numel()

# ESハイパーパラメータ
population_size = 30  # エンジニアリング最適化により小さな人口
learning_rate = 0.001
sigma = 0.017  # 摂動の標準偏差(タスクごとにチューニング)

print(f"Model parameters: {num_params:,} | Population: {population_size}")

# ユーティリティを初期化(メンバーあたりの加重)
utilities = np.array([max(0, np.log(population_size/2 + 1) - np.log(i+1))
                       for i in range(population_size)])
utilities /= np.sum(utilities)  # 正規化

4. メインESループ: 変異、評価、および更新

複数の世代にわたってESを実行し、報酬を累積してパラメータを更新します。

def es_train_loop(
    model, tokenizer, params_init, reward_fn,
    generations=100, population_size=30, sigma=0.017, lr=0.001,
    seed_base=42, device="cuda"
):
    """
    メインEvolution Strategiesトレーニングループ。

    Args:
        model: 微調整するLLM
        tokenizer: モデルトークナイザー
        params_init: 初期パラメータベクトル
        reward_fn: 関数(model, tokenizer) -> float
        generations: ESの反復数
        population_size: 反復ごとの人口メンバー
        sigma: 摂動の標準偏差(探索度を制御)
        lr: 自然勾配のステップサイズ
        seed_base: 再現可能性のためのRNGシード
        device: "cuda"または"cpu"
    """
    params_current = params_init.clone()
    rewards_history = []

    for gen in range(generations):
        gen_rewards = []
        param_updates = np.zeros(params_init.numel())

        # 人口を生成および評価
        for member_id in range(population_size):
            # シードからの決定論的ノイズ(ストレージオーバーヘッドなし)
            seed = seed_base + gen * population_size + member_id
            np.random.seed(seed)
            noise = torch.tensor(
                np.random.randn(params_init.numel()),
                dtype=params_init.dtype,
                device=device
            )

            # 摂動されたパラメータ
            params_perturbed = params_current + sigma * noise

            # モデル重みをインプレースで更新(レイヤーごと)
            offset = 0
            for param in model.parameters():
                param_size = param.numel()
                param.data = params_perturbed[offset:offset+param_size].reshape(param.shape)
                offset += param_size

            # 評価(報酬のみ、勾配なし)
            reward = reward_fn(model, tokenizer)
            gen_rewards.append(reward)

            # 更新用のユーティリティ加重ノイズを累積
            param_updates += utilities[member_id] * noise.cpu().numpy()

        # 報酬を正規化してパラメータを更新
        rewards_array = np.array(gen_rewards)
        rewards_normalized = (rewards_array - np.mean(rewards_array)) / (np.std(rewards_array) + 1e-8)

        # 自然勾配更新: θ ← θ + α * (1/σ) * Σ util_i * noise_i * (r_i - mean_r)
        param_updates_weighted = np.zeros_like(param_updates)
        for member_id in range(population_size):
            seed = seed_base + gen * population_size + member_id
            np.random.seed(seed)
            noise_update = np.random.randn(params_init.numel())
            param_updates_weighted += utilities[member_id] * noise_update * rewards_normalized[member_id]

        params_current = params_current.cpu() + (lr / sigma) * torch.tensor(param_updates_weighted, dtype=params_current.dtype)
        params_current = params_current.to(device)

        # 進捗をログ
        best_reward = np.max(gen_rewards)
        mean_reward = np.mean(gen_rewards)
        rewards_history.append(best_reward)

        if (gen + 1) % 10 == 0:
            print(f"Gen {gen+1:3d} | Best: {best_reward:.4f} | Mean: {mean_reward:.4f} | Std: {np.std(gen_rewards):.4f}")

    return params_current, rewards_history

5. 微調整モデルを保存および評価

訓練後、最終パラメータを復元してパフォーマンスをテストします。

def save_finetuned_model(model, params_final, output_path):
    """
    最終パラメータをモデルに書き戻してディスクに保存する。

    Args:
        model: 保存するアーキテクチャのLLM
        params_final: ESからの最終パラメータベクトル
        output_path: 保存先ディレクトリ(model.save_pretrainedで作成)
    """
    # 最終パラメータを復元
    offset = 0
    for param in model.parameters():
        param_size = param.numel()
        param.data = params_final[offset:offset+param_size].reshape(param.shape)
        offset += param_size

    # ディスクに保存
    model.save_pretrained(output_path)
    print(f"Fine-tuned model saved to {output_path}")


# 使用例
if __name__ == "__main__":
    # データを読み込み(例: 数学推論)
    train_examples = [
        {"input": "Solve: 2x + 3 = 7", "expected_output": "x = 2"},
        # ... より多くの例
    ]

    # 報酬関数を定義
    def reward_fn(m, t):
        return compute_reward(m, t, train_examples[:20])  # 速度のため部分集合

    # ES微調整を実行
    params_final, history = es_train_loop(
        model, tokenizer, params_init, reward_fn,
        generations=100,
        population_size=30,
        sigma=0.017,
        lr=0.001
    )

    # 保存および評価
    save_finetuned_model(model, params_final, "./model_finetuned")

6. 分散マルチGPUセットアップ(Accelerate経由)

大規模モデルの場合、複数のGPUまたはマシン全体で人口評価を分配します。

from accelerate import Accelerator

def es_train_distributed(
    model_name, reward_fn,
    generations=100, population_size=30,
    num_processes=2, gpu_threads=15
):
    """
    Hugging Face Accelerateを使用したマルチGPU ES訓練。
    総並列評価 = num_processes * gpu_threads。

    Args:
        model_name: HuggingFaceモデルID
        reward_fn: 報酬関数(プロセスごとに呼び出し)
        num_processes: GPU数(またはマシン数)
        gpu_threads: GPU当たりのスレッド(GPU当たりのモデルコピー)
    """
    accelerator = Accelerator()

    # 各プロセスは独立してモデルをロード
    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = accelerator.prepare(model)

    # 各プロセスは人口の部分集合を評価
    local_pop_size = population_size // num_processes

    # メインESループ(単一GPU同様だが、報酬は集約される)
    # ...

    print(f"Rank {accelerator.process_index}: evaluating {local_pop_size} members")

Practical Guidance

ハイパーパラメータの推奨

パラメータ	典型的な範囲	注記
population_size	20–50	RLバッチサイズより小さい; 30がデフォルト。難しいタスクの場合は増やす。
sigma(ノイズ標準偏差)	0.01–0.05	探索と搾取のバランスを制御。0.017で開始; 最終調整では低下。
learning_rate	0.0001–0.01	パラメータ更新のステップサイズ。0.001が標準; 振動する場合は削減。
generations	50–500	タスク依存; 報酬曲線を監視してプラトーを検出。
seed_base	任意	再現可能性を確保; 複数試行が必要な場合は実行ごとに増加。

ES微調整をいつ使用するか

• 疎で遅延報酬を伴う推論タスク(数学、論理、パズル解法)
• 異質な基本モデル: Qwen、Llama、Mistralなど全体で動作するメソッドが必要
• ロバストネス重要: 限界報酬ゲイン以上に訓練安定性が重要
• 報酬仕様が難しい: 結果ラベルはあるが中間教師信号がない
• 小規模データセット: ESはサンプル効率的(RLデータの20%未満しばしば必要)
• 長期タスク: 中間ステップが少なく、最終回答のみが評価可能

ES微調整をいつ使用しないか

• 密な報酬信号: 損失勾配または詳細な中間教師信号がある場合、勾配ベースRL(PPO、DPO)が高速
• 連続アクション空間: ESは大規模離散パラメータ空間に優れており、アクション微調整ではRLがより直接的
• 極度のスピード要求: ESは更新ごとに複数の順方向パスが必要; レイテンシが重要な場合、SFTまたは単一パスメソッドを優先
• 高度にモデル固有の最適化: 単一モデルのチューニングで勾配チューニングに無制限コンピュートがある場合、RLが追加パフォーマンスを絞り出すかもしれない
• 評価予算が限定: 各世代はpopulation_size回の完全モデル評価が必要; 評価が高価(例: ヒューマンインザループ)な場合、より小さい人口またはRL重要度加重を使用

一般的な落とし穴

1. σが高すぎるまたは低すぎる: ノイズが大きすぎると更新がランダムになる。小さすぎるとローカル最適に固執。タスクごとにσを調整(0.017から開始、報酬がプラトーの場合は半減)。

2. 報酬スケールを無視: 世代ごとに報酬を正規化することは安定更新に重要。報酬が0–1 vs. 0–1000の場合、学習率は調整が必要; アルゴリズムはz-スコア正規化で処理します。

3. 大規模タスクで小さい人口: population_size < 15の場合、勾配推定はノイズ。複雑な推論では30以上を使用。

4. 貪欲デコーディングを無視: ESは決定論的報酬を前提(同じ入力 → 同じ出力)。生成中のサンプリングはノイズを追加; 貪欲デコーディングまたは固定シードを使用。

5. 中途の訓練チェックポイントから開始: ESは現在のパラメータポイントから検索; 基本モデルが訓練不足の場合、ESは弱い行動に最適化する可能性。強い基本モデルを微調整。

6. 不正なユーティリティ重み: ユーティリティベクトルは報酬で人口メンバーをランク付け。世代ごとに再計算されることを確認(異なるタスク全体で再利用しない)。

Reference

論文: Evolution Strategies at Scale: LLM Fine-Tuning Beyond Reinforcement Learning 著者: Xin Qiu, Yulu Gan, Conor F. Hayes, Qiyao Liang, Elliot Meyerson, Babak Hodjat, Risto Miikkulainen ArXiv: 2509.24372 コード: GitHub – Cognizant AI Lab

引用されたベースライン: PPO (Schulman et al., 2017)、GRPO (Xu et al., 2024)、DPO (Rafailov et al., 2023)