メインコンテンツへジャンプ

TAO:ラベル付きデータなしで効率的なLLMを訓練するためのテストタイムコンピュートの使用

モザイクリサーチチーム による投稿

in


この投稿を共有する
Test-time Adaptive Optimization (TAO)

大規模な言語モデルは、新しいエンタープライズタスクに適応するのが難しいです。プロンプト作成はエラーが発生しやすく、限定的な品質向上しか達成できません。一方、ファインチューニングは、ほとんどのエンタープライズタスクで利用可能な人間によるラベル付けデータを大量に必要とします。今日、我々は新しいモデルチューニング方法を紹介します。これは未ラベルの使用データのみを必要とし、企業が既に持っているデータだけを使用してAIの品質とコストを改善することを可能にします。私たちの方法、テスト時適応最適化(TAO)は、テストタイムコンピュート(o1とR1によって普及)と強化学習(RL)を活用して、過去の入力例だけに基づいてモデルがタスクをより良く行うように教えることができます。つまり、人間のラベリング努力ではなく、調整可能なチューニング計算予算に応じてスケールします。重要なことに、TAOはテストタイムコンピュートを使用しますが、それはモデルを訓練するプロセスの一部として使用します。そのモデルは、低い推論コスト(つまり、推論時に追加の計算を必要としない)で直接タスクを実行します。驚くべきことに、ラベル付けされたデータがなくても、TAOは伝統的なファインチューニングよりも優れたモデル品質を達成することができ、Llamaのような安価なオープンソースモデルを、GPT-4oやo3-miniのような高価な独自モデルの品質に近づけることができます。

TAOは、既存の企業データを使用してAIが特定の領域で優れたパフォーマンスを発揮するという問題、つまりデータインテリジェンスに関する当社の研究チームのプログラムの一部です。TAOを使用することで、私たちは3つのエキサイティングな結果を達成します:

  • ドキュメントの質問応答やSQL生成などの特化したエンタープライズタスクでは、TAOは数千のラベル付けされた例で伝統的なファインチューニングを上回ります。それは、ラベルが必要ないLlama 8Bや70Bのような効率的なオープンソースモデルを、GPT-4oやo3-mini1のような高価なモデルと同等の品質に引き上げます。
  • 私たちはまた、多くのタスク全体でLLMを広範囲に改善するために、マルチタスクTAOを使用することもできます。ラベルを使用せずに、TAOは広範なエンタープライズベンチマークでLlama 3.3 70Bのパフォーマンスを2.4%改善します。
  • チューニング時にTAOの計算予算を増やすと、同じデータでより良いモデル品質が得られ、チューニングされたモデルの推論コストは変わりません。

図1は、TAOがFinanceBench、DB Enterprise Arena、およびBIRD-SQL(Databricks SQL方言を使用)という3つのエンタープライズタスクでLlamaモデルを改善する方法を示しています。LLM入力のみにアクセスできるにもかかわらず、TAOは何千ものラベル付き例を用いた伝統的なファインチューニング(FT)を上回り、Llamaを高価なプロプライエタリモデルと同等の範囲に持ってきます。

図1は、TAOがFinanceBench、DB Enterprise Arena、およびBIRD-SQLという3つの企業タスクでLlamaモデルを改善する方法を示しています。

図1は、TAOがFinanceBench、DB Enterprise Arena、およびBIRD-SQLという3つの企業タスクでLlamaモデルを改善する方法を示しています。

図1:3つのエンタープライズベンチマークでのTAO on Llama 3.1 8BとLlama 3.3 70B。TAOは品質の大幅な改善をもたらし、微調整を上回り、高価な独自のLLMに挑戦します。

TAOは現在、LlamaをチューニングしたいDatabricksの顧客にプレビューで利用可能であり、いくつかの今後の製品の推進力となる予定です。このフォームに記入して、プライベートプレビューの一部としてあなたのタスクで試すことに興味があることを表明してください。この投稿では、TAOの動作方法とそれによる結果について詳しく説明します。

TAOはどのように機能するのでしょうか?テストタイムコンピュートと強化学習を使用してモデルをチューニングする

TAOの主要なアイデアは、人間が注釈を付けた出力データを必要とする代わりに、テストタイムコンピュートを使用してモデルがタスクのための可能性のある応答を探索し、これらの応答を評価することに基づいてLLMを更新することです。このパイプラインは、高価な人間の努力ではなく、テストタイムコンピュートを使用してスケールアップすることができます。さらに、タスク固有の洞察(例えば、カスタムルール)を使用して簡単にカスタマイズすることができます。驚くべきことに、このスケーリングを高品質のオープンソースモデルに適用すると、多くの場合で人間のラベルよりも良い結果を得ることができます。

LIFTパイプライン。LIFTは、推論スケーリングを使用してタスクのためのレスポンスを自動的に生成し、スコアリングし、ノイズの多いフィードバックに基づいてモデルを調整する方法を学びます。
Figure 2: The TAO pipeline.

具体的には、TAOは4つのステージで構成されています:

  • レスポンス生成: このステージは、タスクのための例示的な入力プロンプトまたはクエリを収集することから始まります。Databricksでは、これらのプロンプトは、AIゲートウェイを使用して任意のAIアプリケーションから自動的に収集することができます。それぞれのプロンプトは、多様な候補レスポンスを生成するために使用されます。ここでは、単純な思考の連鎖から洗練された推論と構造化されたプロンプト技術まで、幅広い生成戦略を適用することができます。
  • レスポンススコアリング: このステージでは、生成されたレスポンスが系統的に評価されます。スコアリング方法には、報酬モデリング、優先度ベースのスコアリング、またはLLM判定者やカスタムルールを利用したタスク固有の検証など、さまざまな戦略が含まれます。このステージは、各生成レスポンスが品質と基準との整合性について定量的に評価されることを保証します。
  • 強化学習(RL)トレーニング: 最終段階では、RLベースのアプローチがLLMの更新に適用され、モデルが前のステップで特定された高スコアのレスポンスに密接に一致する出力を生成するように導きます。この適応学習プロセスを通じて、モデルはその予測を洗練して品質を向上させます。
  • 継続的な改善: TAOが必要とする唯一のデータは、LLMの例示的な入力です。ユーザーは、LLMとの対話を通じて自然にこのデータを作成します。あなたのLLMがデプロイされるとすぐに、次のTAOのラウンドのためのトレーニングデータの生成を開始します。Databricks上では、あなたのLLMはTAOのおかげで、それを使うほど良くなることができます。

重要なことは、TAOがテストタイムコンピュートを使用しているにもかかわらず、それを使用して訓練モデルがタスクを直接低推論コストで実行することです。これは、TAOによって生成されたモデルが、元のモデルと同じ推論コストと速度を持つことを意味します - o1、o3、R1のようなテストタイムコンピュートモデルよりも大幅に少ないです。私たちの結果が示すように、TAOで訓練された効率的なオープンソースモデルは、品質で先導的な独自モデルに挑戦することができます。

TAOは、AIモデルをチューニングするためのツールキットで新たな強力な方法を提供します。プロンプトエンジニアリングは遅く、エラーが発生しやすい一方、ファインチューニングは高価で高品質な人間のラベルを作成する必要がありますが、TAOはAIエンジニアがタスクの代表的な入力例を提供するだけで素晴らしい結果を達成することができます。

LLMチューニング方法の比較。
Table 1: Comparison of LLM tuning methods.

TAOは非常に柔軟な方法で、必要に応じてカスタマイズすることができますが、Databricksでのデフォルトの実装は、多様なエンタープライズタスクでそのままうまく機能します。私たちの実装の中核には、新しい強化学習と報酬モデリングの技術があり、これによりTAOは探索によって学習し、その後RLを使用して基礎となるモデルを調整することができます。例えば、TAOを支える要素の一つは、幅広いタスクで正確なスコアリング信号を生成できる、企業タスク向けに訓練したカスタム報酬モデル、DBRMです。

TAOを使用してタスクのパフォーマンスを向上させる

このセクションでは、特定のエンタープライズタスクでLLMを調整するためにTAOをどのように使用したかを詳しく説明します。人気のあるオープンソースのベンチマークと、私たちがドメインインテリジェンスベンチマークスイート(DIBS)の一部として開発した内部のものを含む、3つの代表的なベンチマークを選びました。

 表2:このブログで使用されたベンチマークの概要。
Table 2: Overview of benchmarks used in this blog.

各タスクについて、いくつかのアプローチを評価しました:

  • オープンソースのLlamaモデル(Llama 3.1-8BまたはLlama 3.3-70B)をそのまま使用します。
  • ラマでのファインチューニング。これを行うために、私たちは数千の例が含まれる大規模で現実的な入出力データセットを使用または作成しました。これは通常、ファインチューニングで良好なパフォーマンスを達成するために必要です。これらには以下のものが含まれます:
    • FinanceBenchのためのSEC文書に関する7200の合成質問。
    • DBエンタープライズアリーナのための4800の人間が書いた入力。
    • BIRD-SQLトレーニングセットからの8137の例、Databricks SQL方言に合わせて修正されました。
  • TAO on Llamaは、微調整のデータセットからの例示入力のみを使用し、出力は使用せず、DBRMエンタープライズ向け報酬モデルを使用します。DBRM自体はこれらのベンチマークで訓練されていません。
  • 高品質な独自のLLM - GPT 4o-mini、GPT 4o、およびo3-mini。

表3に示すように、すべての3つのベンチマークと両方のLlamaモデルにおいて、TAOは基準となるLlamaのパフォーマンスを大幅に向上させ、さらには微調整を超える結果を出しています。

表3:3つのエンタープライズベンチマークでのLlama 3.1 8BとLlama 3.3 70Bに対するTAO。
Table 3: TAO on Llama 3.1 8B and Llama 3.3 70B across three enterprise benchmarks.

古典的なテストタイムコンピュートと同様に、TAOはより多くの計算を利用できるときに高品質な結果を生み出します(例として図3を参照してください)。しかし、テストタイムコンピュートとは異なり、この追加の計算はチューニングフェーズ中にのみ使用され、最終的なLLMは元のLLMと同じ推論コストを持ちます。例えば、o3-miniは私たちのタスクで他のモデルよりも5-10倍多くの出力トークンを生成し、比例して推論コストが高くなりますが、TAOは元のLlamaモデルと同じ推論コストを持ちます。

図3: TAOは、チューニングプロセス中に使用されるテスト時間の計算量とともにスケールします。結果として得られるLLMを使用するための推論コストは、元のLLMと同じです。
Figure 3: TAO scales with the amount of test-time compute used during the tuning process. Inference cost to use the resulting LLM is the same as the original LLM.

TAOを用いたマルチタスクインテリジェンスの改善

これまで、私たちはTAOを使用して、SQL生成などの個々の狭いタスクでLLMを改善してきました。しかし、エージェントがより複雑になるにつれて、企業は1つ以上のタスクを実行できるLLMがますます必要となっています。このセクションでは、TAOがエンタープライズタスクの範囲でモデルのパフォーマンスを広範に改善する方法を示します。

この実験では、コーディング、数学、質問応答、ドキュメント理解、チャットを含む多様なエンタープライズタスクを反映した175,000のプロンプトを集めました。次に、Llama 3.1 70BとLlama 3.3 70BでTAOを実行しました。最後に、人気のあるLLMベンチマーク(例えば、アリーナハード、ライブベンチ、GPQAダイヤモンド、MMLUプロ、ヒューマンエバル、MATH)およびエンタープライズに関連する複数の領域での内部ベンチマーク。

TAOは両モデルのパフォーマンスを有意に向上させます[t][u]。ラマ3.3 70Bとラマ3.1 70Bはそれぞれ2.4ポイントと4.0ポイント向上します。TAOは、エンタープライズタスクにおいてLlama 3.3 70BをGPT-4oに大幅に近づけます[v][w]。これはすべて、人間のラベリングコストなしで、代表的なLLM使用データと私たちのTAOの製品実装を使用して達成されます。コーディングを除くすべてのサブスコアで品質が向上し、パフォーマンスは静的です。

表4:TAOを使用してマルチタスクエンタープライズインテリジェンスを改善する
Table 4: Improving multitask enterprise intelligence using TAO

実践でのTAOの使用

TAOは、テストタイムコンピュートを活用することで多くのタスクで驚くほどうまく機能する強力なチューニング方法です。自分のタスクで成功裏に使用するためには、以下が必要です:

  • あなたのタスクに対する十分な例入力(数千)、AIアプリケーション(例えば、エージェントに送られた質問)から収集したものや合成的に生成したもの。
  • 十分に正確なスコアリング方法:Databricksの顧客にとって、ここでの強力なツールの一つは、TAOの実装を支えるカスタム報酬モデル、DBRMですが、タスクに適用可能な場合は、カスタムスコアリングルールや検証者でDBRMを補完することができます。

TAOや他のモデル改善方法を可能にするベストプラクティスの一つは、AIアプリケーションのためのデータフライホイールを作成することです。AIアプリケーションをデプロイするとすぐに、Databricks Inference Tablesのようなサービスを通じて、入力、モデルの出力、その他のイベントを収集することができます。その後、入力だけを使用してTAOを実行することができます。あなたのアプリケーションをより多くの人々が使用するほど、それを調整するためのデータが増え、TAOのおかげであなたのLLMがより良くなります。

結論とDatabricksでの開始方法

このブログでは、ラベル付きデータを必要とせずに高品質な結果を達成する新しいモデルチューニング技術であるTest-time Adaptive Optimization(TAO)を紹介しました。我々は、エンタープライズの顧客が直面している主要な課題を解決するためにTAOを開発しました:彼らは標準的な微調整に必要なラベル付きデータを持っていませんでした。TAOは、テストタイムコンピュートと強化学習を使用して、企業がすでに持っているデータ、例えば入力例を使用してモデルを改善し、任意のデプロイされたAIアプリケーションを品質とコストの面で改善することを容易にします。TAOは、特化したAI開発のためのテストタイムコンピュートの力を示す非常に柔軟な方法であり、我々はそれが開発者にとって強力でシンプルな新しいツールとなると信じています。

Databricksの顧客はすでにプライベートプレビューでLlama上のTAOを使用しています。プライベートプレビューの一部として、あなたのタスクで試すことに興味があることを示すために、このフォームを記入してください。TAOは、私たちの今後のAI製品のアップデートとローンチの多くにも組み込まれています - 乞うご期待!


¹ 著者:Raj Ammanabrolu、Ashutosh Baheti、Jonathan Chang、Xing Chen、Ta-Chung Chi、Brian Chu、Brandon Cui、Erich Elsen、Jonathan Frankle、Ali Ghodsi、Pallavi Koppol、Sean Kulinski、Jonathan Li、Dipendra Misra、Jose Javier Gonzalez Ortiz、Sean Owen、Mihir Patel、Mansheej Paul、Cory Stephenson、Alex Trott、Ziyi Yang、Matei Zaharia、Andy Zhang、Ivan Zhou

² このブログ全体でo3-mini-mediumを使用しています。
³ これはDatabricksのSQL方言と製品に合わせて修正されたBIRD-SQLベンチマークです。