AMDが画像生成AI「Nitro-T」をリリース、32基のInstinct MI300Xでゼロから1日未満でトレーニング可能

半導体大手のAMDが、独自の画像生成AIである「Nitro-T」を発表しました。Nitro-Tは高効率な学習に重点を置いた、テキストから画像を生成することができるDiffusion Transformerモデル(DiTモデル)です。AMDは自社GPUであるInstinct MI300Xを32基使うことで、Nitro-Tをゼロからわずか1日未満でトレーニングすることができたとアピールしています。

Nitro-T: Training a Text-to-Image Diffusion Model from Scratch in 1 Day — ROCm Blogs
https://rocm.blogs.amd.com/artificial-intelligence/nitro-t-diffusion/README.html

AMDは2024年11月に、高速推論に重点を置いた拡散モデルのNitro-1をリリースしました。AMDはNitro-1から、「テキストから画像を生成する画像生成AIを、リソース効率の高い方法でゼロからトレーニングする方法」について研究を続けてきたそうです。

この「トレーニング効率の向上」を進め、AMDはついにNitro-Tで32基のInstinct MI300Xでトレーニングを1日未満にまで短縮することに成功。これにより、オープンソースのDiTモデルであるPixArt-αと比較すると、トレーニングコストは約14分の1にまで削減することが可能となります。AMDはNitro-Tの高いトレーニング効率について、「スマートなアーキテクチャの選択やシステム最適化など、複数の最先端手法を組み合わせることで実現しました」と説明しています。

Nitro-Tには2つのバリエーションがあります。一辺512ピクセルの画像生成に最適化されたパラメータ数が6億のDiTモデルと、一辺1024ピクセルという高解像度の画像生成に最適化されたパラメーター数が12億のMultimodal Diffusion Transformer(MMDiT)モデルです。これらのモデルはテキストコンディショニングを提供するためにLlama 3.2 1Bを利用しており、トレーニングにはパッチシーケンス長の短縮、収束の高速化、トレーニングスループットの最適化を可能にする戦略と設計選択が組み込まれています。

以下は左が「画像生成AIのトレーニングにかかる時間」(縦軸)を比較したグラフで、右が「推論レイテンシ(出力を生成するまでにかかる遅延時間)」(縦軸)を比較したグラフです。競合画像生成AIと比べ、Nitro-T(赤)はトレーニング効率が非常に高く、推論レイテンシが圧倒的に短いです。

Transformerベースのアーキテクチャにおける主要な計算ボトルネックのひとつは、自己注意メカニズムです。自己注意メカニズムの複雑さは入力トークンの数に比例して増大するため、トレーニング中にトークンシーケンスの長さを短縮することは、モデルのトレーニングを高速化する自然な方法となります。

拡散Transformerは画像パッチから生成されたトークンを操作します。画像には本質的に空間的な冗長性が存在するため、トークンをランダムにドロップまたはマスクすることで、トークンシーケンスの長さを短縮できます。しかし、単純なトークンドロップは、高レベルの構造をトレーニングするために必要な重要な情報を破棄することで、モデルのパフォーマンスを大幅に低下させる可能性があります。この問題を解決するため、MicroDiTで導入された遅延マスキング戦略を採用。

この手法では、事前にトークンマスクを適用する代わりに、いくつかのTransformerブロックで構成されるパッチミキサーと呼ばれる軽量の中間モジュールを導入します。重要なアイデアは、パッチミキサーステージの後までマスキング操作を遅らせることで、モデルが最初に画像全体の情報を集約してミックスできるようにすることです。これにより、トークンのサブセットがドロップされた場合でも、残りのトークンが入力のグローバルビューをエンコードし、重要なコンテキストと構造の信号が保持されます。

実際にNitro-Tで出力した画像のサンプルが以下。

Nitro-T-1.2Bで出力した画像

Nitro-T-0.6Bで出力した画像

AMDは「リソース効率の高いアプローチを用いてトレーニングされたテキストから画像を出力する拡散モデル群であるNitro-Tのトレーニングにおける経験を共有し、トレーニング時間を大幅に短縮しながら競争力のあるパフォーマンスを実現することに成功しました。遅延パッチマスキング、ディープ圧縮オートエンコーダー、表現アライメントなどの革新的な技術と、Instinct MI300X上の最新のROCmソフトウェアスタックを活用することで、これらのモデルを1日もかからずにゼロからトレーニングできることを実証しました。最終的に、これらの進歩により、研究者はアイデアをより速く反復できるようになり、独立した開発者や小規模なチームがニーズや制約にぴったり合ったモデルをトレーニングまたは微調整するための障壁が低くなります。完全なトレーニングコードとモデルの重みを公開することで、私たちの研究が実験を促し、生成AIツールへのアクセスを民主化し、この分野のさらなる研究の発展に役立つことを願っています」と記しました。

なお、Nitro-T-1.2BおよびNitro-T-0.6BのモデルデータはHugging Faceで配布されています。

amd/Nitro-T-1.2B · Hugging Face
https://huggingface.co/amd/Nitro-T-1.2B

amd/Nitro-T-0.6B · Hugging Face
https://huggingface.co/amd/Nitro-T-0.6B