4B規模の画像生成・編集モデル Mage-Flow──Mage-VAEと融合により学習スループットを2.5倍に向上

論文概要

大規模な視覚生成モデルは表現力が高い一方で、学習、ファインチューニング、およびデプロイにおけるコストが非常に高いという課題がある。本論文では、効率的なテキストからの画像生成および指示ベースの画像編集を行うための、4B規模のコンパクトな生成スタック「Mage-Flow」を提案する。本手法は、高忠実度な軽量潜在トークナイザ「Mage-VAE」と、整流フローマッチング（rectified flow matching）を用いて学習された「Native-Resolution Multimodal Diffusion Transformer」の2つの共同設計されたコンポーネントで構成されている。

新規性と貢献

Mage-Flowの主な貢献は、トークナイザ、バックボーン、システム（CUDAカーネル融合）の徹底的な共同設計（Co-design）にある。Mage-VAEによりトークン化コストを1桁以上削減しつつ、ネイティブ解像度のパッキングとスタックレベルのCUDAカーネル融合を組み合わせることで、エンドツーエンドの学習スループットを約2.5倍に向上させた。また、Base、RL整列済み、Turboの各バリエーションを展開し、1024^2解像度において単一のNVIDIA A100 GPUで生成0.59秒、編集1.02秒という低レイテンシを達成している。

提案手法の詳細

Mage-Flowのアーキテクチャは、Mage-VAEとNative-Resolution Multimodal Diffusion Transformerを核とする。Mage-VAEは、1ステップの拡散スタイルエンコーディングとデコーディングにアンカー潜在正則化を適用し、強力な既存のパブリックVAEと同等の再構築品質を保持しながら、トークン化コストの大幅な削減を達成する。また、モデルファミリーとしてBaseモデルのほか、指示追従性や美的品質を高めるためのDiffusion-NFT学習モデル、および敵対的知覚ガイダンスによる数ステップ蒸留を行った4ステップのTurboモデルが提供されている。

評価・考察

標準的な生成および編集ベンチマークにおいて、Mage-FlowおよびMage-Flow-Editはコンパクトな4B規模でありながら優れた性能を実証している。特に推論性能の面では、NVIDIA A100 GPU環境下において、Mage-Flow-Turboは1024^2解像度の画像を0.59秒で生成し、Mage-Flow-Edit-Turboは同解像度の画像編集を1.02秒で完了させる。これにより、高解像度な生成・編集処理をインタラクティブな用途に耐えうる実用的な速度と省メモリで実行可能であることを示している。

応用例と今後の展望

本手法の低レイテンシで高解像度な画像生成・編集能力は、日本のWebデザイン業やゲーム開発分野におけるアセット制作プロセスの効率化に直結する実務インパクトを持つ。特にインタラクティブなクリエイティブツールへの組み込みにおいて、GPUリソースを抑えた高速な試行錯誤が可能になる。今後の課題としては、さらに多様なドメインやマルチモーダル入力への拡張、およびエッジデバイス向けへのスケーリング検証が挙げられる。

結論

Mage-Flowは、トークナイザ・バックボーン・システムの緻密な共同設計により、4B規模のコンパクトなモデルファミリーでありながら高い生成品質と高速な推論を両立できることを示した。インタラクティブ利用を見据えたTurboバリエーションの実用性は、今後の視覚生成モデルの軽量化と高効率化の新たな方向性を示すものである。

注釈

Rectified Flow Matching: 確率的軌道を直線的に結ぶことで生成プロセスを効率化するフローベースの生成モデル学習手法。
CUDAカーネル融合 (Kernel Fusion): 複数のGPU演算カーネルを一体化し、メモリ帯域のボトルネックを削減して実行速度を高める最適化技術。