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

大規模言語モデル（LLM）は、新しい企業向けタスクに適応させるのが難しいです。プロンプトはエラーが発生しやすく、品質向上の効果は限られている一方で、ファインチューニングには大量の人手でラベル付けされたデータが必要であり、これはほとんどの企業タスクには存在しません。今回は、新しいモデル調整方法を紹介します。この方法はラベルなしの使用データのみを必要とし、企業が既に持っているデータだけでAIの品質とコストを改善できるようにします。この方法、Test-time Adaptive Optimization（TAO）は、テストタイムコンピュート（o1やR1で普及した技術）と強化学習（RL）を活用して、過去の入力例のみを基にモデルにタスクを改善させる方法です。つまり、TAOは人手によるラベル付けではなく、調整可能な計算リソースでスケールします。重要なのは、TAOはテストタイムコンピュートを使用しますが、それをモデル訓練プロセスの一部として活用することです。その後、訓練されたモデルは低い推論コストでタスクを直接実行します（推論時に追加の計算リソースは必要ありません）。驚くべきことに、ラベル付きデータがなくても、TAOは従来のファインチューニングよりも優れたモデル品質を達成でき、Llamaのような低コストのオープンソースモデルを、GPT-4oやo3-miniのような高価な専有モデルと同等の品質に引き上げることができます。

TAOは、企業がすでに持っているデータを使ってAIを特定のドメインで優れた性能を発揮させる問題に取り組む、私たちのデータインテリジェンス研究チームのプログラムの一環です。TAOを使用すると、以下の三つのエキサイティングな結果が得られます：

1. 文書質問応答やSQL生成などの専門的な企業タスクにおいて、TAOは何千件ものラベル付き例で行う従来のファインチューニングを上回ります。Llama 8Bや70Bのような効率的なオープンソースモデルを、GPT-4oやo3-miniのような高価なモデルと同じ品質にまで引き上げます。

2. 複数のタスクに対して多目的TAOを使用すると、LLMの広範なパフォーマンス向上が可能です。ラベルなしで、TAOはLlama 3.3 70Bの性能を、広範な企業向けベンチマークで2.4％向上させます。

3. 調整時にTAOの計算リソースを増やすと、同じデータを使ってより良いモデル品質が得られますが、調整後のモデルの推論コストは変わりません。

図1は、TAOがLlamaモデルを3つの企業タスク（FinanceBench、DB Enterprise Arena、BIRD-SQL）でどのように改善するかを示しています。LLMの入力のみを利用できるにもかかわらず、TAOは従来のファインチューニング（FT）を上回り、Llamaを高価な専有モデルと同等の品質に引き上げます。

図1：Llama 3.1 8BおよびLlama 3.3 70BにおけるTAOの性能向上。TAOは品質の大幅な改善をもたらし、ファインチューニングを上回り、高価なLLMに挑戦しています。

TAOは現在、Llamaのチューニングを希望するDatabricksのお客様向けにプレビュー版として提供されており、今後のいくつかの製品に搭載される予定です。プライベートプレビューで自身のタスクにTAOを試してみたい場合は、このフォームに記入してご興味をお知らせください。本記事では、TAOの仕組みとその成果についてさらに詳しく説明します。

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

TAOの主なアイデアは、人手で注釈を付けた出力データを必要とせず、テストタイムコンピュートを使用してモデルにタスクに対する適切な応答を探索させ、その後、強化学習を用いてこれらの応答を評価し、LLMを更新することです。このパイプラインは、品質向上のために高価な人手労力の代わりにテストタイムコンピュートを使用してスケールできます。さらに、タスク固有のインサイト（例えばカスタムルール）を使用して簡単にカスタマイズすることが可能です。驚くべきことに、このスケーリングを高品質なオープンソースモデルに適用することで、多くのケースで人手によるラベル付けよりも優れた結果が得られます。

TAOは、以下の4つの段階で構成されています。

1. 応答生成: この段階では、タスクのための入力プロンプトやクエリの例を収集します。Databricksでは、これらのプロンプトはAI Gatewayを使用して、任意のAIアプリケーションから自動的に収集できます。収集されたプロンプトは、さまざまな候補応答を生成するために使用されます。ここでは、単純な思考の連鎖を用いたプロンプトから、洗練された推論や構造的なプロンプト技術に至るまで、幅広い生成戦略が適用できます。

2. 応答スコアリング: この段階では、生成された応答が体系的に評価されます。スコアリング手法には、報酬モデル、優先順位に基づいたスコアリング、またはタスク固有の検証（LLMジャッジやカスタムルールを使用）が含まれます。この段階では、生成された応答が品質と基準にどれほど一致しているかが定量的に評価されます。

3. 強化学習（RL）トレーニング: 最終段階では、RLベースのアプローチを使用して、LLMを更新し、前の段階で特定された高スコアの応答に密接に一致するような出力を生成するようモデルを誘導します。この適応的な学習プロセスを通じて、モデルは予測を洗練し、品質を向上させます。

4. 継続的改善: TAOが必要とするデータは、例としてのLLM入力のみです。ユーザーはLLMと対話することで自然にこのデータを生成します。LLMがデプロイされると、次のTAOサイクルのためのトレーニングデータが生成されます。Databricksでは、TAOのおかげでLLMを使えば使うほど改善されます。

重要なことに、TAOはテスト時の計算を使用しますが、それを使ってモデルを訓練し、そのモデルがタスクを直接実行できるようにします。これにより、TAOで生成されたモデルは、元のモデルと同じ推論コストと速度で実行でき、o1、o3、R1のようなテスト時の計算モデルよりもはるかに低い推論コストで動作します。私たちの結果が示すように、TAOで訓練された効率的なオープンソースモデルは、品質で優れた商用モデルに挑戦することができます。

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

TAOは非常に柔軟な手法で、必要に応じてカスタマイズも可能ですが、Databricksのデフォルト実装は、さまざまな企業向けタスクに対してそのままでも十分に効果を発揮します。
この実装の核となっているのは、私たちのチームが開発した新しい強化学習と報酬モデリングの技術です。これにより、TAOは探索を通じて学習し、その後、強化学習を使って基盤モデルをチューニングすることが可能になります。
たとえば、TAOの中心要素の一つとして、企業向けタスクに特化して私たちが訓練したカスタム報酬モデル「DBRM」があり、幅広いタスクに対して高精度なスコアリング信号を提供します。

TAOによるタスク性能の向上

このセクションでは、TAOを使ってLLMを企業向けの専門タスクに最適化した方法について、より詳しく説明します。
私たちは、一般的なオープンソースベンチマークに加えて、社内で開発した「Domain Intelligence Benchmark Suite（DIBS）」の一部である独自ベンチマークも含む、3つの代表的なベンチマークを選定しました。

各タスクに対して、以下のアプローチを評価しました：

1. オープンソースのLlamaモデル（Llama 3.1-8BまたはLlama 3.3-70B）をそのまま使用する。

2. Llamaのファインチューニング。これには、数千の例を含む大規模で現実的な入力と出力のデータセットを使用または作成しました。ファインチューニングで良いパフォーマンスを得るためには通常これが必要です。具体的には以下のデータセットを使用しました：

   • FinanceBenchのSEC文書に関する7200件の合成質問

   • DB Enterprise Arenaのための4800件の人間が書いた入力

   • Databricks SQLダイアレクトに合わせて修正したBIRD-SQLトレーニングセットからの8137件の例

3. Llamaに対するTAO。ここでは、ファインチューニングデータセットからの例示的な入力のみを使用し、出力は使用せず、DBRMという企業向けの報酬モデルを使用します。DBRM自体はこれらのベンチマークで訓練されていません。

4. 高品質な商用LLM – GPT 4o-mini、GPT 4o、およびo3-mini。

表3に示されているように、すべてのベンチマークと両方のLlamaモデルにおいて、TAOはLlamaのベースラインパフォーマンスを大幅に改善し、ファインチューニングを超える性能を実現しています。

従来のテスト時計算と同様に、TAOはより多くの計算リソースにアクセスできると、より高品質な結果を生み出します（例として図3を参照）。しかし、テスト時計算とは異なり、この追加の計算リソースはチューニングフェーズでのみ使用され、最終的なLLMは元のLLMと同じ推論コストとなります。例えば、o3-miniは私たちのタスクで他のモデルの5～10倍の出力トークンを生成し、それに比例して推論コストが高くなりますが、TAOは元のLlamaモデルと同じ推論コストです。

 TAOによるマルチタスクインテリジェンスの向上

これまで、私たちはTAOを使って、SQL生成のような個別の狭いタスクにおけるLLMの性能を向上させてきました。しかし、エージェントがより複雑になるにつれて、企業は複数のタスクをこなせるLLMをますます必要としています。このセクションでは、TAOがどのようにして企業向けのさまざまなタスクでモデルのパフォーマンスを広範囲にわたって向上させるかを示します。

この実験では、コーディング、数学、質問応答、文書理解、チャットなど、さまざまな企業向けタスクを反映した175,000件のプロンプトを収集しました。その後、TAOをLlama 3.1 70BおよびLlama 3.3 70Bに実行しました。最後に、企業関連のタスクを一通りテストしました。これには、一般的なLLMベンチマーク（例：Arena Hard、LiveBench、GPQA Diamond、MMLU Pro、HumanEval、MATH）と、企業に関連する複数の分野の社内ベンチマークが含まれます。

TAOは、両方のモデルのパフォーマンスを意味のある形で向上させました^[t][u]。Llama 3.3 70BとLlama 3.1 70Bはそれぞれ2.4％と4.0％向上しました。TAOは、Llama 3.3 70BをGPT-4oに大きく近づけ^[v][w]、企業向けタスクでのパフォーマンスを改善しました。これらすべては、人間によるラベリングコストなしで、代表的なLLM使用データとTAOの生産実装のみで達成されました。品質は、コーディングを除くすべてのサブスコアで向上し、コーディングに関してはパフォーマンスが一定のままでした。

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を開発しました。それは、標準的なファインチューニングに必要なラベル付きデータが不足しているという問題です。TAOはテスト時計算と強化学習を利用して、企業がすでに持っているデータ（例えば、入力例）を使ってモデルを改善します。この方法により、デプロイ済みのAIアプリケーションの品質を簡単に向上させ、小さなモデルを使用することでコストを削減できます。TAOは非常に柔軟な手法であり、専門的なAI開発におけるテスト時計算の力を示しています。私たちは、TAOがプロンプトエンジニアリングやファインチューニングとともに、開発者に強力でシンプルな新しいツールを提供するものだと考えています。

Databricksの顧客はすでにLlamaでTAOをプライベートプレビューで使用しています。このプライベートプレビューの一環として、自身のタスクに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ベンチマークです。