【Diffusion勉強会】Scaling Diffusion Transformers to 16 Billion Parameters

DEEP LEARNING JP [DL Papers] Scaling Diffusion Transformers to 16 Billion Parameters Jeong Seong Cheol, Matsuo Lab, University of Tokyo http://deeplearning.jp/

書誌情報 Title: Scaling Diffusion Transformers to 16 Billion Parameters Author: Zhengcong Fei, Mingyuan Fan, Changqian Yu Debang Li, Junshi Huang* Affiliation: Kunlun Inc. Beijing, China Status: Preprint Publish date(arxiv): 16 Jul 2024 Code: https://github.com/feizc/DiT-MoE/tree/main Abstract: DiT-MoEを提案した。5M枚の合成画像で学習させた16BのモデルでFIDがSOTA。 アブレーションでは、MoE layerの深さと選択されるExpertの関係を 1)条件付けたクラス 2)入力したパッチ位置 3)サンプリングステップ数 の観点で調査。示唆深い結果が得られた。 Figure 2: Overview of the DiT-MoE architecture. Generally, DiT-MoE is built upon the DiT and composed of MoE-inserted Transformer blocks. In between, we replace the MLP with a sparsely activated mixture of MLPs. The right subfigure demonstrates the details of MoE layer integration of the shared expert strategy.

研究の背景と目的 TransformerベースのLDMのスケーラブルな学習を低コスト（省メモリ、高速）で行いたい Unet Base チャネルサイズを増やす [1]Bigger is not Always Better: Scaling Properties of Latent Diffusion Models [1]では39M-5Bで検証 Transformer Base Transfomer Block (Dit Block)を増やす [2]Scalable Diffusion Models with Transformers DiT論文ではmax 118.6 Gflops [3]Scaling Rectified Flow Transformers for High-Resolution Image Synthesis SD3論文ではmax8B 本論文では 199MM-16.5B (max 163.51 Gflops)

提案手法：DiT-MoE アーキテクチャの違いはFFNがMoEに変わったのみ(赤枠) Figure 2: Overview of the DiT-MoE architecture. Generally, DiT-MoE is built upon the DiT and composed of MoE-inserted Transformer blocks. In between, we replace the MLP with a sparsely activated mixture of MLPs. The right subfigure demonstrates the details of MoE layer integration of the shared expert strategy. Dit[2] DiT-MoE

MoE layer Common knowledgeを学習するShared Expertsによって学習効率Up。無いと知識がoverlapping ここを説明してます→ 1 2 3 4 5 6

学習アルゴリズム Expert-level balance lossによりExpertの学習に偏りがないようにする ※Routing後に計算するLoss. わかりにくい&Notationミスありそうなので右参照→ つまり Expert-Level Balance Loss. We use an expert-level balance loss (Fedus et al., 2021; Lepikhin et al., 2021) to mitigate the risk of routing collapse: L_bal = α1 * sum(fi * Pi) (23) fi = (Nr / KrT) * sum(1(Token t selects Expert i)) (24) Pi = (1/T) * sum(Si,t) (25) where α1 is a hyper-parameter called expert-level balance factor; 1(·) denotes the indicator function; and T denotes the number of tokens in a sequence. [4]DeepseekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models α1: 重み係数(ハイパラ) Nr: Expertの総数(shared以外) Kr: Top kのExpertが選ばれる (Token t selects Expert i), Expert iが選ばれた回数 あるトークンに対して各Expertの選択される確率 (1行Nr列のベクトル,ベクトルの合計1) L_balance = α * sum(i=1 to N) (実際にexpert iが選ばれた割合) × (expert iが選択される平均確率) L_balance = (αn / KT) * (1/N) * (1/N) = (αn / KT) * (1/N) 均等にExpertが選ばれた場合 L_balance = (αn / KT) * 1 * 1 = (αn / KT) * 1 同じExpertのみ選ばれた場合

実験設定：評価方法と比較対象 Expert数を変えたりBlock数を変えたりしてモデルサイズをスケーリングする Task: Class-conditional image generation Datasets: ImageNet (256x256 and 512x512) Data Amount: 1,281,167 training images (1,000 classes) Baseline: GAN, Unet-base LDM, Transformer-base LDM Evaluation metrics: FID50K, sFID, Inception Score, Precision/Recall Sampling step: 250 (DDPM) Table 1: Scaling law model size. The model sizes, detailed hyperparameters settings, and inference burden for MoE scaling experiments. L/2-8E2A means pach size: p = 2, Expert amount: n = 8, and top K Expert= 2.

結果：Class conditional image generation 256x256と512x512どちらもFIDは最小。Activeなパラメータ数は1.5Bで、3Bや7Bモデルに勝った。 Table 2: Benchmarking class-conditional image generation on ImageNet 256x256 dataset. We can see that DiT-MoE-XL/2 achieves state-of-the-art FID metrics towards best competitors with less inference cost. Table 3: Benchmarking class-conditional image generation on ImageNet 512x512 dataset. DiT-MoE demonstrates a promising performance compared with dense networks for diffusion. Activeなパラメータ数は1.5B さらに SDXLとSD3-Mediumで512x512を5M枚合成画像生成。 元データ1M枚と合計6M枚を使ってDiT-MoE-G/2-16E2A (16.5B)学習。 FID50K 1.80 (SOTA)

Ablation①: Shared Expertの有無と数、モデルサイズ Shared Expertはあった方が良いが2個より増やしても変化は少ない、一方モデルサイズは大きいほど良い Figure 3: Ablation experiments on ImageNet dataset at 256x256 resolution. We report FID metrics on 50K generated samples without CFG. (a) Incorporation of shared expert routing can accelerate the training as well as optimize generated results. (b) Number of experts and (c) model parameters scaling. As we expected, increasing the expert number and the model size can consistently improve the generation performance. However, directly changing the share experts number influences the results marginally.

Ablation②: expertの数と学習安定性、収束速度 Expertを増やす(0 -> 8 ->16)と学習曲線にスパイクが発生しやすくなるが、収束は早くなる Figure 4: Training loss curves for small version variants. We can see that the incorporation of simple designs helps DiT-MoE training stably and consistently converge. On the other hand, increasing the expert number helps decrease while introducing more loss spikes.

Ablation③: MoE layerの深さとexpertの活性 特定のクラスと特定のExpertに相関はなさそう 縦軸: 条件付けた1,000クラス 横軸: Expertの番号(0~7) 色の濃くなっているところは そのExpertが選択されたことを 意味する Figure 5: Frequency for selected experts per image class. We show the 12 MoE layers of DiT-MoE-S/2-8E2A. The x-axis corresponds to the 8 experts in a MoE layer. The y-axis is the 1000 ImageNet classes. For each pair (expert e, image class i), we show the average routing frequency for the patches corresponding to all generated images with class i that particular expert e. The darker the color, the higher the frequency of selection. The larger the layer number, the deeper the MoE layers.

Ablation③: MoE layerの深さとexpertの活性 浅い層のMoE Layerではパッチ位置と渡されるExpertに偏りがある。層が深くなると均一に渡される。 縦軸: Imageのパッチ位置の番号 (256x256画像->256パッチ) 横軸: Expertの番号(0~7) 色の濃くなっているところは そのExpertが選択されたことを 意味する Figure 6: Frequency for selected experts per image patch position. We show the 12 MoE layers of DiT-MoE-S/2-8E2A. The x-axis corresponds to the 8 experts in a MoE layer. The y-axis are the 256 patches in ImageNet images with 32/2 * 32/2 = 256 sequence length of patch size 2, at 256x256 resolution with VAE compression 8. For each pair (expert e, image patch id i), we show the average routing frequency for all the patches with patch-id i that were assigned to that particular expert e.

Ablation③: MoE layerの深さとexpertの活性 層の深さに関係なく50 step未満では、渡されるExpertが偏っている(例えばLayer2ではExpert3)。 縦軸: Sampling step(250) 横軸: Expertの番号(0~7) 色の濃くなっているところは そのExpertが選択されたことを 意味する Figure 7: Frequency for selected experts per denoising time step. We show the 12 MoE layers of DiT-MoE-S/2-8E2A. The x-axis corresponds to the 8 experts in a layer. The y-axis are the 250 DDPM steps for sampling the synthesis image. For each pair (expert e, inference step i), we show the average routing frequency for all time step i that were assigned to that particular expert e.

Related works: Switch Diffusion Transfomer[5] どちらもSwitch Transfomerを参考にしているが、MoE箇所や検証しているモデルサイズが異なる DiT-MoE MoE箇所: FFN 検証モデルサイズ: 199M~16.5B Switch Diffusion Transfomer MoE箇所: MHAの前に新たに追加 検証モデルサイズ: 36M-749M (DiT-MoE論文ではスパースじゃないと指摘) Figure 2: Overview of the DiT-MoE architecture. Generally, DiT-MoE is built upon the DiT and composed of MoE-inserted Transformer blocks. In between, we replace the MLP with a sparsely activated mixture of MLPs. The right subfigure demonstrates the details of MoE layer integration of the shared expert strategy. Fig. 1: Switch Diffusion Transformer. ⊙ represents an element-wise multiplication.

まとめと感想 感想：DiT-MoEはスケーリングLDMのデファクトスタンダードになりそう。 DiT-MoEを提案した。 技術のキモは、 1) Common knowledgeを学習するShared Expert 2) Expertを均等に学習させるExpert-level balance loss 合計6M枚の合成画像で学習させた16BのモデルでImagenet(512x512) のFID50KがSOTA。 アブレーションでは、MoE layerの深さと選択されるExpertの関係を 1)条件付けたクラス 2)入力したパッチ位置 3)サンプリングステップ数 の観点で調査。示唆深い結果が得られた。

Reference [1] Mei, Kangfu et al. "Bigger is not Always Better: Scaling Properties of Latent Diffusion Models." ArXiv abs/2404.01367 (2024): n. pag. [2] Peebles, William, and Saining Xie. "Scalable diffusion models with transformers." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023. [3]Esser, Patrick, et al. "Scaling rectified flow transformers for high-resolution image synthesis." Forty-first International Conference on Machine Learning. 2024. [4]Dai, Damai, et al. "Deepseekmoe: Towards ultimate expert specialization in mixture-of- experts language models." arXiv preprint arXiv:2401.06066 (2024). [5]Park, Byeongjun, et al. "Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Experts." arXiv preprint arXiv:2403.09176 (2024).