SageMaker AIのマルチターン強化学習をWordleで解説

何が発表されたか

  • AWS Machine Learning Blogが、Amazon SageMaker AIの「マルチターン強化学習(MTRL)」を実際に試す手順を公開した。
  • エージェントはAmazon Bedrock AgentCore上、またはEKS・EC2・Fargateなど自前の環境に置ける。学習対象のモデルを呼び出し、複数ターンで行動し、終了時に1つの報酬を返す。
  • 題材はWordle。gpt-oss-20bの正解率は、デフォルトの推論強度で61%から84%、学習時の低い推論強度で15%から69%に上がった(未使用の単語での評価)。
  • 学習トークンの費用は約339ドル。AgentCoreランタイムの利用料は別に約79ドルかかった。
  • データ、エージェント、学習済みモデルは利用者のAWSアカウント内に置かれ、ベースモデルの学習には使われないと記載されている。

ポイント解説

なぜ「マルチターン」が必要なのか

一般的なファインチューニングは、1回の質問に対してより良い回答を出せるよう調整するものです。一方、経費精算やサポートチケットの振り分けのようなエージェント型のタスクでは、行動を何手も重ねます。各手は直前の結果に左右され、「正しい手順」が1通りに決まらないことも多くあります。

わかるのは、最後まで通してうまくいったかどうかだけ、という場面です。記事によると、教師あり学習にはこうした場面に使えるラベルがなく、1回完結型の強化学習では数手前の判断に結果を結び付けられません。MTRLは、最終結果の採点だけを手がかりに学習する仕組みです。利用者が指定するのは「正しい手順」ではなく「結果の採点方法」です。

仕組みの全体像

役割は2つに分かれています。SageMaker AIは学習ジョブと学習対象のモデルを担当します。AgentCoreは、エージェント、Wordleの環境、報酬の計算を1つのPythonファイルとして担当します。

学習データの1行ごとにエージェントが呼び出され、モデルが最大6回まで推測し、最後に報酬が返ります。同じ単語を複数回プレイした結果の報酬のばらつきを使ってモデルを更新する手法がGRPOで、更新対象はLoRAアダプターです。LoRAは、モデル全体ではなく小さな追加部品だけを学習する、一般的な軽量チューニング手法です。

報酬設計と費用

報酬は結果重視で、正解なら速さに応じて1.05〜1.55、失敗なら-0.5、有効な推測が1度もなければ-1.5です。この上に、戦略の良さを評価する小さな加点が付いています。

費用はトークン数に応じた従量課金で、インスタンスやクラスターの設計は不要です。記事では、学習側のトークンが費用の中心で、対戦が長いほど増えると説明されています。したがって、ゲームを短くしたり推論の出力を簡潔にしたりすると、費用を抑えられるとされています。

つまずきやすい点

記事は、実際に遭遇した落とし穴も挙げています。

  • 学習ジョブ用のIAMロールの信頼先は、job.sagemaker.amazonaws.comにする必要がある。汎用のプリンシパルだとエラーになる。
  • エージェントのハンドラーは、BedrockAgentCoreAppとentrypointで包まないと、呼び出しが30秒でタイムアウトする。
  • ジョブの実行中にエージェントを再デプロイすると、ジョブ全体が失敗する。
  • MLflowアプリは必須で、学習の進み具合の確認に使う。

評価用のジョブも用意されています。推論強度や温度を変えて、ベースモデルと学習済みモデルを同じ条件で比べられます。

読むときの注意点

Wordleは説明用の題材で、記事自身も最先端のモデルはすでにこのゲームを得意としていると述べています。本来の狙いは、小さく安価なモデルを自社固有の複数手順タスクに強くすることです。ただし、示された数字はこの1回の実行結果です。自社の業務で同じ改善が得られるかどうかは、試してみないとわからないと考えられます。なお、掲載した本文は途中までで、評価結果の詳細な比較は確認できていません。

どんな人に関係があるか

  • 社内の定型業務をAIエージェントに任せたいが、公開モデルでは精度が足りない開発者。
  • 大型モデルのAPI利用料を抑えるため、小型モデルを自社タスク向けに鍛えたい人。
  • GPUクラスターを自分で用意せずに、強化学習を試したい人。

具体的には、複数のツールを順番に呼び出す手続き型の業務で、完了したかどうかを機械的に判定できるものが向いていると考えられます。その場合、判定ロジックがそのまま報酬関数になります。一方、結果の良し悪しを自動で採点しにくい業務では、報酬の設計自体が難しくなるかもしれません。

まとめ

SageMaker AIのMTRLは、何手も続くエージェントの作業を、最終結果の採点だけで学習させる仕組みです。Wordleの例では、gpt-oss-20bで正解率の向上が示され、学習トークンの費用は約339ドルでした。IAMロールや再デプロイなど実運用での注意点も公開されているため、自社タスクで試す際の参考になります。


出典: Multi-Turn Reinforcement Learning on Amazon SageMaker AI: A hands-on guide(AWS Machine Learning Blog)

※この記事はAIを活用して作成しています。正確な情報は出典元をご確認ください。

コメント

タイトルとURLをコピーしました