PRXのデータ戦略:多様性と網羅性を追求したPre-training

Hugging Faceのブログで詳細が明かされたPhotoroomの画像生成モデルPRXのデータ戦略は、その開発プロセスにおいてデータがいかに重要かを再認識させます。彼らのデータパイプラインは、「パブリックおよび内部データセットの混合、VLM(Vision-Language Model)による画像のリキャプション、そしてストリーミング可能なコーパスへの変換」という一連の流れで構成されています (出典)。

特に注目すべきは、Pre-training(事前学習)におけるデータセット構築の指針です。PRXチームは、モデルに「世界の様子」を学習させるため、多様で広範なデータセットを重視しています。

視覚的概念、物体、シーン、構成、照明など、画像が包含するあらゆる要素を網羅することが目的です。彼らは、個々の画像の品質よりも、分布の広さと多様性がモデルの学習においてより重要であると考えています。

Pre-trainingは網羅性のためにあり、ファインチューニングは品質のためにある。

過度な美しさのフィルタリングは、逆に分布を狭め、モデルが後から回復できない概念や構成の多様性を失うことに繋がりかねないという見解は、データセットのキュレーションにおいて非常に示唆に富んでいると感じます。これは、モデルが現実世界の複雑さを理解するために、初期段階では多少のノイズも許容し、広範な知識を獲得することが不可欠である、というアプローチだと考えられます。


混合データソースと高精度キャプションが鍵

PRXのデータソースは、パブリックデータセットと内部データセットの混合で構成されています。既存のキュレーションされたデータセットを活用し、ゼロから全てを構築するのではなく、効率的なアプローチを採用したとのことです。これは、限られたリソースの中で迅速にモデルを開発するスタートアップらしい現実的な戦略ではないでしょうか。

データセットの品質を左右するもう一つの重要な要素は、キャプションの精度です。PRXチームの経験では、「画像内の全てを正確に記述する長いキャプション」がPre-trainingにとって最も重要だとしています。これは、短すぎるキャプションから長いキャプションへ切り替えたことで、モデルの生成品質が大幅に向上したという過去の経験に基づいています (出典)。

キャプションタイプ Pre-trainingにおける効果
短いキャプション 生成品質が低い傾向
長いキャプション サンプル品質が大幅に向上

正確なキャプションは、画像に含まれるスクリーンショット、広告、ロゴ、テキストといった「ノイズ」を、モデルが条件付け可能で制御可能な属性として学習させる効果があります。これにより、ノイズだと思われた要素をプロンプトで制御したり、逆にプロンプトで排除したりすることが可能になります。この「ノイズを情報に変える」考え方は、私の普段のコーディングでも、エラーログから有用な情報を引き出すアプローチと通じるものがあるように感じます。


データ管理を支える Lance と MDS の活用

大規模なデータセットを効率的に管理するため、PRXのデータパイプラインでは2種類のデータフォーマットが使い分けられています。分散学習には、これまで「Mosaic Streaming」と「Mosaic Data Shards (MDS)」が使用されてきました。これらはS3やGCSのようなオブジェクトストレージから直接分散学習を可能にし、低メンテナンスで柔軟、高性能なフレームワークとして機能します (出典)。

しかし、MDSにはカラムの追加やサブセット作成の際にデータセット全体のスキャンと再書き込みが必要になるという課題があります。この柔軟性の低さを補うために採用されたのが、Lance です。Lanceは、述語プッシュダウン、スカラーインデックス、ベクトル検索機能を備えたカラムナデータフォーマットであり、何十億もの行を持つデータセットの構築と探索に適しています (出典)。

フォーマット 特徴 用途
MDS 低メンテナンス、柔軟、高性能な分散学習 データストリーミング
Lance カラムナ、高速クエリ、ベクトル検索、柔軟なスキーマ変更 データ構築、キュレーション、特徴量エンジニアリング

つまり、Lanceでデータセットを構築・キュレーションし、MDSで効率的にストリーミングするという役割分担が行われているわけです。この戦略的なフォーマットの使い分けは、大規模AIモデル開発におけるデータ管理のベストプラクティスを示していると感じます。特に、GCPを主軸に普段BigQueryやCloud Runを扱っている身としては、こういったデータ基盤の選択がモデルのパフォーマンスや開発効率に直結すると強く感じます。

NVIDIAも「AI-ready dataこそが真の優位性」と指摘しており、モデルやGPUの数よりも、いかに適切にデータを準備するかが重要であると述べています (出典)。PRXの事例は、まさにこの考え方を実践していると言えるでしょう。