nanoVLM: ビジョン言語モデルのトレーニングのためのミニマリスト PyTorch ツールキット

Hugging Face は、Vision Language Model（VLM）のトレーニングプロセスを分かりやすくすることを目的とした、純粋な PyTorch で書かれたミニマリストツールキット nanoVLM を導入しました。Andrej Karpathy の nanoGPT に触発され、nanoVLM は読みやすく軽量なコードベースを提供し、ユーザーが無料枠の Google Colab ノートブック上で VLM をトレーニングできるようにします。

Vision Language Model（VLM） の基礎

Vision Language Model は、画像とテキストの両方の入力を処理し、テキスト出力を生成できるマルチモーダルアーキテクチャです。VLM は画像キャプション生成、物体検出、セマンティックセグメンテーションなど多様なタスクに利用できますが、nanoVLM は主なトレーニング目的として Visual Question Answering（VQA） に特化しています。

技術アーキテクチャ

nanoVLM は、2 つの事前学習済みバックボーンを投影層で整合させるモジュラーアーキテクチャを採用しています:

  • Vision Backbone: Google の SigLIP (google/siglip-base-patch16-224) ビジョンエンコーダを使用します。
  • Language Backbone: Llama 3 アーキテクチャに従い、HuggingFaceTB/SmolLM2-135M を使用します。
  • Modality Projection Module: このモジュールはビジョンとテキストのモダリティを整合させます。ビジョンバックボーンからの画像埋め込みを、言語モデルの埋め込み層と互換性のある埋め込みに変換します。このプロセスは pixel shuffle 操作を含み（画像トークン数を減らして計算コストを低減し、トレーニング速度を向上させます）その後に線形層が続きます。

トレーニングワークフローと実装

トレーニングプロセスは train.py によって管理され、データセットのロード、モデルの初期化、最適化を行います。

データパイプライン

get_dataloaders 関数は Hugging Face の load_dataset API を活用してデータセットをロード、シャッフル、分割します。カスタムデータセット（VQADataset、MMStarDataset）とコラレータ（VQACollator、MMStarCollator）を使用してデータを準備します。

最適化戦略

事前学習済みバックボーンのトレーニングと新たに初期化されたプロジェクタのバランスを取るため、nanoVLM は二重学習率（LR）戦略を採用しています。

  • Higher LR: モダリティプロジェクタ（MP）に適用され、迅速な学習を促進します。
  • Lower LR: エンコーダ/デコーダスタックに適用され、バックボーンに既存の知識を保持します。

トレーニングループとモニタリング

トレーニングは混合精度のために torch.autocast を組み込み、線形ウォームアップ付きのコサイン学習率スケジュールを使用します。パフォーマンスはトークンスループット（tokens/sec）で監視され、設定されていれば Weights & Biases（wandb）でバッチロス、検証ロス、精度を追跡します。

推論と事前学習済みモデル

Hugging Face は、Hub に公開された事前学習済み nanoVLM モデル（nanoVLM-222M）を提供しています。このモデルは cauldron データセットから 170 万サンプルを使用し、単一の H100 GPU 上で約 6 時間トレーニングされました。

ユーザーは generate.py スクリプトを使用して推論を実行でき、以下の論理的な流れに従います：

  1. Initialization（初期化）: モデル、トークナイザー、画像プロセッサをロードします。
  2. Processing（処理）: テキストプロンプトをトークン化し、画像をテンソルに変換します。
  3. Generation（生成）: model.generate を実行してテキスト出力を生成します。
  4. Decoding（デコード）: batch_decode を使用して生成されたトークンを人間が読めるテキストに変換します。

はじめに

トレーニングを開始するには、リポジトリをクローンし、トレーニングスクリプトを実行します：

# Clone the repo
git clone https://github.com/huggingface/nanoVLM.git

# Execute the training script
python train.py

Sources