Hugging Face⁠⁠、小型AIモデルにデータ分析を学習させる強化学習環境「SmolDataEnvs」を公開

Hugging FaceのAdithya S Kolavi氏は9月24日、コードを使ったデータ分析を小型AIモデルに学習させるための強化学習(RL)環境「SmolDataEnvs」の公開をXで告知した。データ分析プラットフォームKaggleのデータセットやコンペティションを基に、5,000件以上の課題を用意した。課題を解くための実行環境と解答を自動採点する仕組みも提供する。開発者は、課題や採点プログラムを一から用意する手間を省き、小型モデルの学習・評価に取り組める。

Adithya氏は、強化学習の入門例に多い単語当てゲームのWordleや算数問題のGSM8Kでは、学習後のモデルが実用につながりにくいと述べている。そこで、100億パラメーター未満の小型モデルに、実際のデータ分析に役立つ能力を身に付けさせるため、SmolDataEnvsを開発したという。

データ分析の課題をコードで解き⁠、正解と照合

SmolDataEnvsの各課題では、分析対象のデータファイルと、そのデータに関する質問を与える。AIモデルはツールを使うエージェントとして、データの内容を調べ、コードを実行して答えを求める。収録された課題は、難易度に応じてeasy、medium、hardの3段階に分類されている。具体例として、サッカーの試合データから引き分けの割合を計算する課題がある。

採点には、エージェントが出した解答を、あらかじめ用意した正解と照合するプログラムを使う。文字列の完全一致だけでなく、数値が許容誤差の範囲に収まるか、リストやパーセントの表記が異なっていても同じ答えを示しているかなどを判定する。採点にLLMは用いず、プログラムによる採点結果を報酬として、モデルの強化学習を進める。

課題の作成時にも、エージェントに実際に課題を解かせ、用意した正解と同じ答えが得られるかを確かめた。曖昧な問題や解答を検証できない問題は除外したという。

学習⁠・評価用の実行環境と⁠、解答過程を学ぶためのデータを提供

エージェントの評価・学習向けフレームワーク「Harbor」の課題形式に沿った実行環境も提供する。課題ごとに、コードを動かすコンテナー環境、分析対象のデータ、採点プログラムをまとめている。Adithya氏は、学習と評価で同じ環境構成を使える点を挙げている。

学習用の課題として5,000件のほか、モデルの能力を確かめるテスト用に250件、学習中の検証用に144件を用意した。テスト用と検証用は、学習用より難しい課題の割合が高い。

また、教師ありファインチューニング(SFT)用のデータには、学習用5,000件のうち4,677件の課題について、エージェントが正解に至るまでの解答過程を1件ずつ収録している。データを調べるためのツール呼び出しやその実行結果、最終的な解答までの記録を手本に、モデルを学習させられる。

関連コードは、複数の強化学習環境のプロジェクトをまとめたGitHubのFineEnvsリポジトリで公開している。SmolDataEnvsはその一つで、専用ディレクトリにSFTと強化学習を試すノートブックやスクリプトを収めている。

SmolDataEnvsのREADMEでは、学習前のモデルを評価し、SFTと強化学習を経て再び評価する手順を案内している。公開ノートブックでは、SFTはデフォルトの設定でGoogle ColabのT4 GPU、強化学習はA100 GPUを1基使う構成となっている。

ライセンスは公開物によって異なっている。GitHubのFineEnvsリポジトリはApache License 2.0を採用している。Hugging Faceで配布するSmolDataEnvsの課題データセット、SFT用データ、Harbor形式の学習用実行環境は、各配布ページのライセンス欄にMITと記載されている。

おすすめ記事

記事・ニュース一覧