11月3日、AWSが「Build efficient agents with knowledge distillation on Amazon SageMaker AI」と題した記事を公開した。この記事では、Amazon SageMaker AIで知識蒸留（Knowledge Distillation）を使い、フロンティアモデル並みの品質を持つ小型エージェントを低コストで構築する手法について詳しく紹介されている。

AIエージェントの「コスト問題」を知識蒸留で解く

AIエージェントが実用段階に入ると、コスト構造の問題が表面化する。エージェントは1回のタスクで複数回モデルを呼び出すループ構造になっているため、フロンティアモデル（GPT-4やClaude等の最高性能モデル）をそのまま使い続けると、1ユーザー・1タスクあたりの費用が積み重なり、プロダクトのROIを圧迫する。

この記事では、ショッピングアシスタントを例に、その解決策としてsequence-level knowledge distillation（SeqKD）を使ったアプローチを解説している。SeqKDとは、大きな「教師モデル」が生成した高品質な回答テキストを使って、小さな「生徒モデル」をファインチューニングする手法だ。教師モデルの重みやlogit（内部確率分布）へのアクセスは不要で、APIを通じたテキスト出力だけで成立する。

DeepSeek-R1の蒸留モデル（QwenやLlamaをベースにした小型モデル群）もこの手法を大規模に活用したことで知られており、実績のある出発点と言える。

なお本記事では、教師モデルとしてAmazon Bedrock上で利用可能なフロンティアモデルを、生徒モデルとして小型のオープンモデルを使用している。元記事に記載された具体的なモデル名については、元記事のコードおよびウォークスルーセクションを直接参照されたい。

パイプラインの全体像

記事が紹介するパイプラインは次の構成だ。

1. 教師モデルによる軌跡生成：Amazon Bedrock上のフロンティアモデルをShoppingBenchエージェントとして動かし、約427万件の商品カタログから条件に合う商品を探させる
2. リジェクションサンプリング：教師を各問題につき2回実行（temperature=1.0）し、ベンチマークのスコアラーが正解した軌跡だけを残す
3. データエンコード：成功した軌跡をステップ単位のJSONL形式に変換し、S3に格納
4. サーバーレスSFT：SageMaker AI model customizationで小型モデルをファインチューニング（LoRA）
5. 評価・推論：ホールドアウトしたテストセットで評価し、SageMaker推論エンドポイントにデプロイ

データ準備の核心：「解けた軌跡だけ」を学ばせる

最もエンジニアが参考にできる部分はデータ準備のプロセスだ。

ShoppingBenchは750問（商品・店舗・クーポンの3カテゴリ、各250問）のショッピング課題を提供している。このうち426問を学習用、324問を評価用に分割し、教師モデルを各問題で2回実行することで850軌跡を生成。ベンチマークのスコアラーで正解軌跡だけを残すことで、質の高い学習データを得る仕組みだ。

temperature=1.0で2回サンプリングするのは、異なる検索クエリや推論パスを探索させ、少なくとも1回は正解できる確率を高めるためだ。「教師は常に正解する必要はなく、正解した軌跡が模倣に値すればよい」という設計思想が明快だ。

エンコード形式：ステップごとに1行、損失はassistantターンのみ

軌跡をファインチューニング用データに変換する際、1軌跡を複数ステップ分の行に展開する。各行の構造はこうだ。

{
  "messages": [
    {
      "role": "user",
      "content": "You are a shopping assistant. ...\n# Task\nFind a black waist bag under $20 from a shop rated 4.5 or higher.\n# Observations so far\n(none yet)"
    },
    {
      "role": "assistant",
      "content": "<think>Start with a broad search and let the price filter work before I narrow by shop rating.</think>\n<tool_call>[{\"name\": \"find_product\", \"parameters\": {\"q\": \"black waist bag\", \"price\": \"0-20\"}}]</tool_call>"
    }
  ]
}

損失計算はassistantターンのみに適用し、モデルは「次の推論とツール呼び出し」を学ぶ。環境が返した観測値（検索結果）は予測対象に含めない。

トークン予算を守るため、ツール観測値の圧縮が重要になる。商品検索結果はid・価格・評価点・制約チェックに必要なフィールドのみ残し、生のリストテキストは削除する。シーケンス長の上限は16,384トークン。

また、データの分割は「行単位」ではなく「問題単位」で行う。同一問題から複数行が展開されるため、行単位で分割すると学習・検証セット間でデータが漏れる。

ファインチューニングはサーバーレスで

SageMaker AI model customizationを使えば、学習スクリプトもコンテナイメージもインスタンス管理も不要だ。コードは以下のように簡潔に書ける。

from sagemaker.train.sft_trainer import SFTTrainer
from sagemaker.train.common import TrainingType

sft_trainer = SFTTrainer(
    model="huggingface-vlm-qwen3-5-4b",
    training_type=TrainingType.LORA,
    training_dataset=train_dataset.arn,
    validation_dataset=val_dataset.arn,
    model_package_group=model_package_group,
    mlflow_experiment_name="seqkd-shopping-exp",
    s3_output_path="s3://<amzn-s3-demo-bucket>/seqkd/output/",
    accept_eula=True,
)

LoRAのランクは64、エポック数は3。学習メトリクスはAmazon SageMaker with MLflowでリアルタイムに追跡できる。

SageMaker AI model customizationを選ぶ利点は、インフラ管理の省力化にとどまらない。学習ジョブの起動・停止・スケーリングがマネージドで行われるため、オンプレミス環境でのGPUクラスタ構築やGCPのVertex AI Custom Trainingのような別クラウドサービスと比べて、AWSエコシステム（S3・Bedrock・MLflow統合）との連携コストが低く抑えられる点も実務上の利点だ。既存のAWSインフラを持つチームにとっては、追加の認証・ネットワーク設定なしにパイプライン全体を完結させられる。

必要な前提環境

• IAMロール付きAWSアカウント（S3読み書き権限 + 教師モデルの呼び出し権限）
• Amazon Bedrockで利用できるフロンティアモデル
• SageMaker AI model customizationが有効なリージョン
• SageMaker Studio（またはAWS CLI設定済みのローカルIDE）、Python 3.12以上、SageMaker Python SDK v3、boto3
• ml.m5.4xlarge以上のノートブックインスタンス（データ準備とローカルスコアリングに使用）

SeqKDはAPIのみアクセス可能なフロンティアモデルでも動作し、蒸留後は小型モデルで推論できるため、運用コストの削減効果が直接的だ。記事にはGitHubリポジトリと、ShoppingBenchを使った完全なウォークスルーコードも含まれている。

詳細はBuild efficient agents with knowledge distillation on Amazon SageMaker AIを参照していただきたい。