PRX データ戦略: VLM 再キャプションと Mosaic Streaming による事前学習のスケールアップ

PRX データ戦略: VLM 再キャプションと Mosaic Streaming による事前学習のスケールアップ

Photoroom は、7B パラメータモデルである PRX の学習に使用されたデータ戦略の詳細な内訳を公開しました。核となるアプローチには、公開および内部ソースから多様なコーパスを組み立て、高精度な記述を保証するために Vision Language Model (VLM) を使用してすべての画像を再キャプションし、データセットの探索と学習スループットのバランスを取るためにデュアルフォーマットのストレージシステムを利用することが含まれます。

事前学習の原則: 完璧さよりも広範さを

事前学習フェーズにおいて、Photoroom は個々の画像の美学よりも、カバー範囲と多様性を優先しました。目標は、視覚的な世界の幅広い分布にわたって、モデルに視覚的概念、構図、およびライティングを教えることでした。

  • 多様性 vs 嗜好: チームは、事前学習中に美学のために過度なフィルタリングを行うと、分布が狭まり、モデルに不可欠な構図の多様性が失われると主張しています。生成物の洗練は、より小規模で厳選されたセットでのファインチューニングと好みの調整(preference alignment)に委ねられます。
  • 実用的なソーシング: データは公開データセットと内部データセットの混合から組み立てられ、プロセスを加速させるために、利用可能な場合は既存のキュレーション(NSFW および PII フィルタリング)を活用しました。
  • キャプションの哲学: 長く正確なキャプションが極めて重要です。ロゴ、テキスト、またはアーティファクトを含む画像内のすべてを記述することで、これらの要素は無条件なノイズではなく、条件付けされた制御可能な属性になります。

データインフラストラクチャ: Lance と Mosaic Data Shards (MDS)

Photoroom は、データセットのライフサイクルを管理するために 2 つの異なるデータ形式を採用しています。キュレーションには Lance、ストリーミングには MDS を使用します。

探索のための Lance

Lance は、数十億行のデータセットを構築および探索するためのカラム型データフォーマットとして使用されます。これは、安価な述語プッシュダウン、スカラーインデックス、およびベクトル検索をサポートしており、チームが以下のことを可能にします:

  • データのプロファイリング: カットオフを設定するために解像度の分布を分析する(例:384² ピクセル未満の画像を破棄する)。
  • インタラクティブなブラウジング: カスタム UI を介して全文検索と最近傍類似性検索を使用し、多様性を定性的に評価し、非写真的なコンテンツなどの問題を発見する。
  • フラグメント管理: チームは、Lance テーブルを過度に断片化すること(例:1 フラグメントあたり 10 万行)がクエリを遅延させることを指摘しました。1 フラグメントあたり約 100 万行に圧縮することで、パフォーマンスが大幅に向上しました。

学習のための MDS

Mosaic Data Shards (MDS) は分散学習に使用され、決定論的な再開可能なシャッフル、エポック途中での弾力的なチェックポイント作成、および複数のストリームの重み付き混合を提供します。パイプラインを最適化するために:

  • オンザフライ Latents: テキスト Latents を事前計算する代わりに、Photoroom は学習中に Qwen3-VL エンコーダーを使用してそれらを計算します。これにより、MDS シャードのサイズが削減され、テラバイト単位のデータを書き換えることなくエンコーダーの入れ替えが可能になりました。7B スケールでのスループットへの影響はわずか 3–4% でした。
  • JPEG エンコーディング: すべての画像は品質 92 の JPEG として保存されます。実証テストにより、これは PNG と知覚的に区別できず、生成画像の品質に悪影響を与えない一方で、ストレージ要件を 3–10 倍削減できることが示されました。

VLM 再キャプション戦略

Photoroom は、長く詳細なキャプションがサンプル品質を大幅に向上させることを発見しました。ベンチマークでは、Qwen2.5-VL-7B キャプションで学習された小規模な拡散モデルが、より短い LLaVA-1.5-LLaMA3-8B キャプションで学習されたモデルよりも、FID、CMMD、および DINO-MMD メトリクスにおいて優れた性能を示しました。

キャプショナーの選定

いくつかの候補をベンチマークした後、Photoroom は主要なキャプショナーとして Qwen3-VL-8B を選択しました。選定は、品質とスループットのバランスに基づいています:

  • パフォーマンス: Qwen3.5-9B がメトリクスで最高のパフォーマンスを示しましたが、Qwen3-VL-8B はそれに近く、大幅に高速でした(H200 GPU 1 枚あたり 20 img/s)。
  • プロンプティング: チームは、推測や解釈をせず、正確な視覚的グラウンディングに焦点を当てた、100–200 語の単一の流れるような段落を要求する密なシステムプロンプトを使用しました。

データセットの精緻化と重複排除

VLM キャプションが生成された後、Photoroom はデータセット全体を書き換えることなく、コーパスを精緻化するために軽微なフィルタリングと重複排除を適用しました。

キャプションベースのフィルタリング

高コストなピクセルレベルの分類器を使用する代わりに、Photoroom はテキスト専用モードの Qwen3-8B を使用して、キャプションを「visual(視覚的)」、「text(テキスト)」、「nsfw」として分類しました。このアプローチは計算効率が高く(GPU 1 枚あたり 200 キャプション/s)、スクリーンショット、ドキュメント、およびインフォグラフィックを効果的に除去します。

知覚ハッシュによる重複排除

ニア・デュプリケート(酷似した重複)を除去するために、チームは DCT ベースの知覚ハッシュを実装しました。2 つの画像のハミング距離がゼロである場合、それらは重複とみなされます。異なる解像度に重複が存在する場合、最高解像度のコピーが保持されます。

スキップリスト・メカニズム

フィルタリングや重複排除のパスのたびに巨大な MDS コーパスを書き換えるのを避けるため、Photoroom はデータローダーに スキップリスト機能 を実装しました。シャードごとのサイドカーファイルにスキップするインデックスがリストされており、ロード時に動的にデータ(例:ユーザーのオプトアウトや品質の問題のため)を除外することができます。

アスペクト比バケツ分け (Aspect-Ratio Bucketing)

画像あたりの計算量を一定に保ち、無駄なパディングやクロッピングを避けるために、Photoroom はアスペクト比バケツ分けを使用しています:

  1. 解像度ティア: 画像はピクセル数に基づいてティア(512px、1024px、2048px、4096px)に割り当てられます。
  2. アスペクト比バケツ: 各ティア内で、画像は一定のパッチ・トークン・予算を維持する 13 のパッチ整合形状(AR 0.5 から 2.0)のいずれかにスナップされます。
  3. 処理: 画像は Lanczos フィルタを使用してリサイズされ、個別のストリーミングのために解像度とアスペクト比によってキー付けされたディレクトリツリーに保存されます。

Sources