LLMは自分の間違いを修正するのがどの程度得意か？KerasとTPUを用いたチャットボット・アリーナの実験

TL;DR

Hugging Faceは、LLMがカレンダー・ツールのためのAPIコールを生成し、その後の短い英語のフィードバックに基づいて間違いを修正する、チャットボット・アリーナの実験を実施しました。Gemma 2 9Bが最も優れたパフォーマンスを示し、エラーを確実に修正しました。Llama 3.1 8Bもそれに近い結果でしたが、追加のプロンプトを必要としました。小型モデル（1-2B）やVicunaのような古いモデルは、失敗するか、支離滅裂な出力を生成することがよくありました。

実験設定

テストには、2つの関数を持つPython API action.add_calendar_entry(title, date="YYYY-MM-DD", time="HH:MM", duration=m) と action.remove_calendar_entry(title, date, time) が使用されました。システムプロンプトは、モデルに対して親切な音声アシスタントとして振る舞い、実行可能なコードを1行で返答するように指示しました。対話は、年が省略された不正確なリクエスト（「Add a meeting with Fred on Nov 11 at 5PM」）から始まり、モデルに間違いを犯させ、その後の「The current year is 2024.」という発言で修正できるようにしました。その後のターンで、イベントの追加、時間の変更、エントリーのキャンセルが行われ、間違いを修正する機会が複数提供されました。

技術的実装: TPUs, JAX, Keras, およびモデル・シャーディング

アリーナはGradioを使用してHugging Face Spaces上に構築されました。これはTPU v5e 2x4上で動作し、8つのコアとコアあたり16GBのRAM（合計128GB）を提供します。このメモリ容量により、著者はすべてのコアにモデルをシャーディング（分割）することで、複数のモデルを同時にロードすることができました。Keras（現在はJAX上でネイティブに動作）を使用し、bfloat16形式で最大5つの約8Bパラメータのモデルと3つの約2Bパラメータのモデルをロードし、合計7つのLLMを同時にメモリに常駐させました。モデル並列性は、Kerasの組み込みレイアウト・マップ（例：keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)）に依存し、各モデルを8つのTPUコアにシャーディングしました。著者は、安定したロードを実現するために、デバッグと時折のレイアウト・マップの調整が必要であったと述べています。

評価されたモデル

実験は、TPUメモリに収まり、タスクが十分に単純であることから選ばれた、10Bパラメータ未満のインストラクション・チューニング済みモデルに焦点を当てました。テストされたファミリーは、Gemma、Llama 3、Mistral、Vicuna、およびCodeGemmaです。具体的なチェックポイントは以下の通りです：

  • Gemma 2 9B-instr
  • Llama 3.1 8B-instr
  • Llama 3.2 3B-instr
  • Llama 3.2 1B-instr
  • Gemma 2B-instr
  • CodeGemma 7B-instr
  • Vicuna 1.5 7B-instr
  • Mistral 7B-instr すべてのモデルは、投稿に記載されているHugging Face HubのURL経由でアクセスされました。

最初の質問に対する信頼性

各モデルに最初のプロンプトを5回ずつ投げかけました。結果（✓ = 正しいAPIコール、🍄 = ほとんど正しいが間違いがある、🔥 = ゴミ/認識可能なコールなし）：

  • Gemma 2 9B-instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.1 8B-instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.2 3B-instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.2 1B-instr: 🔥 🍄 🔥 🔥 🔥
  • Gemma 2B-instr: 🍄 🍄 🍄 🍄 ✓
  • CodeGemma 7B-instr: ✓ ✓ ✓ ✓ ✓
  • Vicuna 1.5 7B-instr: ✓ 🔥 🔥 ✓ 🔥
  • Mistral 7B-instr: ✓ ✓ 🍄 ✓ 🍄 1-2Bのモデルと古いVicunaだけが継続的に失敗し、大型モデルは毎回正解しました。

完全な対話 – 間違いの修正

6ターンの完全な対話を実行した際、著者は間違いのたびにフィードバックを与え、モデルがそれを修正できたかどうかを記録しました（🥦 = 修正成功）。主な結果：

  • Gemma 2 9B-instrとLlama 3.1 8B-instrは、1つか2つの軽微なミスだけで対話を完了しました。Llamaは「fix it」という追加のプロンプトを必要としました。
  • オンラインのGemini（はるかに大型のモデル）は、APIコールを出力するために特別なプロンプトを必要とし、依然としていくつかの間違いを犯しました。これは、サイズだけでこのタスクにおけるパフォーマンス向上が保証されないことを示しています。
  • 小型モデルの中では、Gemma 2B-instrだけが対話を完了できましたが、おしゃべりなテキスト（「Sure, here’s the updated code…」）を追加したり、日付や時間を混同したりする傾向がありました。しかし、指示されれば間違いを修正できました。
  • Vicuna 1.5 7B-instrは、エラーを修正する前に、繰り返し出力やジャンク出力（🔥🔥）に陥ることがよくありました。
  • Mistral 7B-instrは多くの間違いを犯しましたが、修正することができ、多数の修正ターンの後に6つのブロッコリー記号を獲得しました。
  • CodeGemma 7B-instrは、執拗な間違いを犯しました。年にスペースを挿入し（「20 24」）、修正されてもそれを排除できませんでした。

追加の間違い修正テスト

修正能力を分離するため、著者は他のモデルから生成された誤った出力を各モデルに与え、モデルがそれを修復できるか確認しました。結果（🥦 = 正しい修正、🍄 = エラーが継続、🔥 = 複数のエラー）：

  • Gemma 2 9B-instr: 間違った時間と「API only」プロンプトを完璧に修正しました。「Wrong API」プロンプトには苦戦し、正しいコールの他に謝罪を出力しました（これは「API call only please」と頼むことで排除できる繰り返しの間違いです）。
  • Llama 3.1 8B-instr: Gemmaと同様に、ほとんどのエラーを修正しましたが、間違ったAPIコールを確実に修正することに苦労しました。
  • CodeGemma 7B-instr: Gemmaのように振る舞い、謝罪のパターンを繰り返しました。
  • Mistral 7B-instr: 時間のエラーはうまく修正しましたが、「API-only」および「wrong-API」プロンプトに対して一貫性が見られませんでした。
  • すべての小型モデル（Llama 3.2 3B, Llama 3.2 1B, Gemma 2B）およびVicunaは、多くの🔥または🍄の結果を出しました。これは、提供された間違いを確実に修復できないことを示しています。

まとめと示唆

著者は、単純な2つのAPI設定はすべてのモデルにとって容易であると予想していましたが、結果は明確な差を示しました。Gemma 2 9B-instrは、最小限の修正プロンプトで、ほぼ完璧な正確さで全対話をナビゲートできた唯一のモデルでした。Llama 3.1 8B-instrは、わずかに多くのガイダンスを必要としたものの、僅差で2位でした。モデルのサイズは役に立ちましたが、はるかに大型のGeminiモデルはフォーマットの問題によりパフォーマンスが低下しました。これは、指示への忠実度とトークンレベルの制御が、生のスケールよりも重要であることを示しています。小型モデル（≤2B）や古いモデルは、有効なAPIコールを生成できないことが頻繁にあり、生成できたとしても、広範なプロンプトなしに修正フィードバックを取り入れるのに苦労しました。この実験は、ユーザーが平易な英語で間違いを修正できる単純なツール使用シナリオにおいて、現在のミドルサイズのインストラクション・チューニング済みLLMは有用になり得るものの、信頼性は依然としてモデルのスケール、学習データ、および余計なテキストを抑制する能力に依存することを示しています。

読者は、公開されているKeras Chatbot Arena Spaceを使用して、この研究を再現または拡張することができ、間違い修正行動を改善するためのファインチューニングの実験が推奨されます。

Sources