TinyDiT、シングルGPU向け学習レシピを公開：210M画像モデルではregisterがcross-attentionの9割を吸収

TinyDiTは、大規模な商用画像生成モデルに対抗することを目的としたものではなく、現代的なtext-to-image学習を検証可能なシングルGPU実験へと縮小したものだ。作者は1基のRTX PRO 6000上で、420万枚の256²画像を使って210Mパラメータのdiffusion transformerを学習した。400,000ステップに要した時間は3.5日。FLUX.2 autoencoderとflan-t5-baseテキストエンコーダーは凍結され、独自部分にはDiT、データパイプライン、学習レシピ、評価が含まれる。

アーキテクチャにはcross-attention、2D RoPE、QK normalization、SwiGLU、adaLN-singleを採用した。adaLN-singleは作者の旧設計と比べてパラメータ予算を27%節約し、その分を16層、hidden width 896の構成に振り向けている。さらに注目すべきなのがregisterだ。第3層で16個の学習可能な画像tokenを追加し、各cross-attention層には別途2個のnull key/value slotを設けている。計測によると、中間層かつ中程度のノイズ条件では、2個のnull slotがcross-attentionの約90%を吸収し、EOSは約4%まで低下した。一方、registerベクトルのノルムは画像tokenの4〜13倍に達しており、registerがグローバルな一時記憶領域として機能するという解釈を支持している。ただし、これは因果関係を示すablationの証拠ではない。

学習にはrectified flowを採用した。FLUX.2 latentが32 channelsであることから、timestep shiftを2.8に設定し、サンプルの半数がノイズ率74%以上の領域に位置するようにした。`torch.compile`は実際の学習ステップで2.4倍の高速化と約半分のメモリ使用量を実現したが、5種類のアスペクト比bucketごとに静的グラフが1つ必要になる。最終的にFIDは33.7から27.0へ、FD-DINOv2は570から218へ改善し、物体検出精度は65%から90%へ上昇した。同期間にflow lossは0.805から0.754へ下がっただけであり、視覚品質の判定よりも学習状態の健全性を監視する用途に適していることが示された。

制約も明確だ。このモデルは依然として、判読可能な文字、近距離の顔、群衆、3個を超える数量のカウント、複雑な幾何形状を苦手としている。指標は作者が独自に構築した固定プロンプトのワークフローに基づいており、第三者による再現もまだ行われていない。また、重みは凍結された外部コンポーネントに依存している。エンジニアは依存関係とライセンスの全体を先に確認し、「シングルGPUで学習可能」という表現を、エンドツーエンドで完全に独立しているという意味に解釈しないよう注意すべきだ。