【Diffusion勉強会】Representation Autoencoder

Representation Autoencoder(とその周辺) 2025.11.05 Shumpei Takezaki

1 紹介する論文と概要 ICLR2026 submitted • Diffusion Transformers with Representation Autoencoders • DiTに視覚エンコーダー(DINOなど)と学習済みのデコーダーを組み合わせた Representation Autoencoders (RAE)を導入 • ImageNetの生成においてFID 1.51 (256×256)と 1.13(512×512)を達成 視覚エンコーダー デコーダー

2 拡散モデルはLatent spaceで動いている 今日はDiTメイン • Latent Diffusion Models (LDM)やDiffusion Transformer (DiT) が高い生成精度と処理効率を実現 • ほとんどの大規模な拡散モデル[1,2,3]はLatent spaceで動いているといっていい Stage1: Autoencoder for Latent Space Stage 2: Diffusion Model in Latent space Corrupt Latent VAE Enc. VAE Dec. VAE Enc. DiT SD-VAE Reconstruction Denoising

3 視覚エンコーダーを用いた拡散モデルの改善 • 視覚エンコーダー (DINO[4,5], MAE[6], CLIP[7]/SigLIP[8]など)を使って 拡散モデルを改善するような手法が提案 • 意味構造を持つ特徴空間を拡散モデルの学習に活用！ DINO[4,5] MAE[6] VFM=Vision Foundation Model (視覚エンコーダーと呼んでます) CLIP[7] / SigLIP[8]

4 視覚エンコーダーを用いた拡散モデルの改善 • 視覚エンコーダーの特徴と”Alignment”を取るのが常套手段 • Stage1: Latent spaceでAlignment • Stage2: DiT中間特徴でAlignment

5 視覚エンコーダーを用いた拡散モデルの改善 • 視覚エンコーダーの特徴と”Alignment”を取るのが常套手段 • Stage1: Latent spaceでAlignment • Stage2: DiTの中間特徴でAlignment Latent VAE Enc. VAE Dec. VA-VAE[9] Align ReDi[10] VFM Feature

6 視覚エンコーダーを用いた拡散モデルの改善 • 視覚エンコーダーの特徴と”Alignment”を取るのが常套手段 • Stage1: Latent spaceでAlignment • Stage2: DiTの中間特徴でAlignment REPA[11] REG[13] REPA-E[12] DDT[14] Corrupt VAE Enc. DiT Hidden feature VFM Align Feature

7 提案手法: Representation Autoencoder (RAE) • RAE: 視覚エンコーダーと学習済みのデコーダーを組み合わせたAE • RAEのLatent spaceに適したDiTを提案 (small) Corrupt RAE Dec. VFM VFM DiT RAE Reconstruction Denoising DDT Head Corrupt

8 提案手法: Representation Autoencoder (RAE) • RAE: 視覚エンコーダーと学習済みのデコーダーを組み合わせたAE • RAEのLatent spaceに適したDiTを提案 3. Noise Scheduler (small) Corrupt Corrupt RAE Dec. VFM RAE Reconstruction 1. RAE Recon. 2.DiT Arch. VFM DiT DDT Head 4. Noise-Aug. Decoding Denoising 5. DDT Head

9 1. RAE Reconstruction • RAEで高精度な再構成が可能であるかを評価 • VFM＆Decoderの大きさと再構成精度との関係を調査 DINOv2, SegLIP, MAE ViT-B, L, XL A e Reconstruction SD-VAEとほとんどおなじ (KL lossだけなし)

10 1. RAE Reconstruction A e • RAEで高精度な再構成が可能であるかを評価 • VFM＆Decoderの大きさと再構成精度との関係を調査 基本的にはこの組み合わせ - Enc.: DINOv2-B - Dec.: ViT-XL 再構成評価 (Dec.はViTXL) Decoderの大きさで比較 (Enc.はDINOv2-S) VFMの大きさで比較 (Dec.はViT-XL) 意味的な空間の評価 (Linea prob.) EncoderにVFM使っても 高精度な再構成が可能！ Decoderは大きいほど 精度がいい VFMは小さなモデルで 十分 SD-VAEは意味的な特徴を 捉えていない ※1 rFID=Reconstruction FID. Lower is better ※2 ImageNet-1kで評価

11 t RAEをDiTで使ってみよう i • RAEのLaten spaceでも従来のDiTと同じ計算量 • 𝑛 = VFM特徴空間の次元 (DINOv2-Bでは768) • 𝑑 = DiT特徴空間の次元 (DiT-Bでは768) • 𝑙 = DiT blockの数 (深さ） i ここの大きさが従来のDiTと同じ=同じ計算量 n at i at l i i t パッチサイズ 𝑝=

12 t RAEをそのまま使うと精度が悪い.. i • ImageNetのgFIDが，RAE > SD-VAEとなってしまう • 著者たちはRAEのLatent spaceを用いたDiTのために以下の3つを提案 • DiTアーキテクチャの改善 • 高次元なLatent spaceのためのNoise schedule • Noisy latentsに対応したDecoder ImageNetのgFID (50 steps Euler sampler) ※1 gFID= Generation FID. Lower is better

ここに注目 t 2. DiT Architecture i • DiT特徴次元𝑑とVFM特徴次元𝑛の関係 𝑑 < 𝑛: サンプルの品質が劣化 ） (画像1枚すら再現できない) “画像1枚のみ”を用いて学習したときの結果 𝒅 ≥） 𝒏: 品質が急激に改善 𝑑 < 𝑛 の時，DiT blockの数𝑙を 大きく(深く)しても品質は改善しない 13

ここに注目 t 2. DiT Architecture i • DiT特徴次元𝑑 ≥ VFM特徴次元𝑛であることが必須 • VFMの大きさに合わせたDiTを用意する必要がある • Encoder: DINOv2-B (𝑛 = 7 8)，Decoder: DiT-XL (𝑑 = VFM (DINOv2)とDiTの対応 (損失の値) 𝑑≥𝑛 𝑑<𝑛 )を使用 14

ここに注目 t 2. DiT Architecture • 𝑑 ≥ 𝑛でないといけない理論的な理由づけも行っている • 𝑑 < 𝑛 の時， 𝑛次元データにおける拡散モデルの損失は0に収束しない • 𝑑 ≥ 𝑛だと，０に収束することが保証される i 15

ここに注目 16 t 3. Noise schedule i • 高次元データを適切に拡散するNoise scheduleを使用 • RAE (𝑛 = 7 8)は，SD-VAE (𝑛 ≤ )に比べて次元が高い • 高次元データを低次元データと同じだけ壊すには多くのノイズ付与が必要 • 時刻をシフトし早い時刻でより多くのノイズを付与 (SD3[3]で提案されているものを そのまま使用) Corrupt 猫 猫 ×猫 𝑡𝑏 = + ( 𝑏Τ𝑎 − )𝑡𝑎 𝑏 Τ𝑎 = 3とすると， → 𝑡𝑏 ≈ 0.4 → 時刻0 2で時刻0 4分のノイズ付与 𝑎: シフト前の次元. 409 に設定 𝑏: シフト後の次元. 9 , 08 (= 7 8)に設定 𝑡𝑎 , 𝑡𝑏 : シフト前，シフト後の時刻．[0, ]の範囲． 3 猫 3 低次元 高次元 高次元だと 情報が残る 𝑏Τ𝑎 𝑡𝑎 𝑡𝑎 = 0. , g I : 23 08→4 81に改善

17 4. Noise-Augmented Decoding • DecoderでNoisy latentの再構成を学習 • Noisy latent: Latentに少量のノイズを付与 • 生成されるLatentは必ずしも綺麗ではない • ノイズ付与ありの学習で再構成の頑健性と汎化性を改善 ここに注目 (small) Corrupt RAE Dec. VFM Reconstruction ノイズ付与なし 〃 あり 再構成精度 ( I )は悪化するが， 生成精度 (g I )が改善

18 改善点を組み込んだDiT with RAEを評価 • 1.~4.を組み合わせてモデルで従来手法 (SiT[15], REPA[11])と比較 SiTの47倍, REPAの16倍 効率的な学習 g I =2 39を達成

5. DDT Head i 19 ea ここに注目 t • 高次元なLatentを効率よく計算するためにDDT Headを利用 • 𝑑 ≥ 𝑛の制約のために層を深くすると計算量が急激に増加 • 浅いが幅の広いTransformerモジュールであるDDT Headを導入 • DDT[14]から着想を得たモジュール • DDT Headで扱う特徴だけ高次元に設定 (𝑑′ = 048を採用) 𝑑 𝑑′ 𝑙 i Unpatchify DDT Final l i at 𝑛 ea DDT Block i DiT Block t 𝑑 𝑛

5. DDT Head i ea の導入により，大幅な生成精度の改善を実現 • • 特に高次元なLatent s a e ( INOv2-L)において大幅な改善 • 学習効率も大幅な改善が見られた 学習効率 生成精度 ↑ DDT Headなし 〃 あり ImageNetのg I ↑ ea ここに注目 t 20

21 SOTA DiT手法との比較 • 収束速度，スケーリングの観点で比較 収束速度UP↑ スケーリング性↑

22 ImageNetの画像生成でSOTA手法と比較 lass- n iti nal 256x256 gFID 1.51 (1.13 w/ Autoguidance[16]) Class-conditional 512x512 gFID 1.13 w/ Autoguidance

23 拡散モデルにおける視覚エンコーダーを使ったLatent space • RAEの他にも同時期に同じコンセプトの研究が多数投稿 • 全てICLR2026へ投稿．若干，手法や実装が異なる • RAE以外はVFMの最適化やAdapterを使用 T2Iを実装! G i, et al [19] en, et al [17] Bi, et al.[18] Shi, et al.[20]

24 まとめ • まとめ • DiTのLatent spaceとしてVFMの特徴空間をそのまま使用 • VFMの特徴空間は，再構成精度 (rFID)，生成精度 (gFID)ともに改善 • 同じようなモチベーションの研究が複数報告されている (T2Iも可能） • 感想 • VFMとDiTを“直接つなげた”初めての研究で興味深かった • これからのスタンダードな拡散モデルとなり得るので今後の発展に期待 • AEも含めてすべてがTransformerで成立しているところも大きい

25 参考文献 [1] bin mba , et al , “ ig - es l ti n Image S nt esis wit Latent [2] la est Labs, L X, 2024, tt s://git b [3] at i sse , et al , “S aling e ti ie ans [4] at il e a [5] axime Oq ab, et al , “ INOv2: Lea ning [6] Kaiming n, et al , “ me ging l w [7] Ale a [8] Xia a Z ai, et al , “Sigm i L ss [9] Jing eng Ya , et al , “ e [10] e [11] Si sK nst zelis, et al , “ n Y , et al , “ e [13] Ge W , et al , “ e [14] S [16] e [17] me s,” I 2021 is al eat es wit tS e visi n,” L , 2024 isi n Lea ne s,” els Lang age Image e- aining,” I Gene ati n: Gene ati n: i si n a, et al , “Si : x l ing l w an i si n en, et al , “Aligning is al n ati n [19] ing G i, et al , “A a ting Sel -S [20] inglei S i, et al , “Latent i si n i ans aining els an e G e st e si n ilemma in Latent i si n i si n i i si n els,” 2025 e S nt esis,” Ne I S2025 ans me s Is asie si n ans an Y in ”, I L 2025 ans me s”, I 2025 me s Is asie an Y in ,” a xiv e int, 2025 els wit S alable Inte lant ans me s,” Itsel ,” Ne I S2024 enize s enize s i Latent si n i e esentati ns as a Latent S a e el wit L2021 2023 ning wit Latent si n-base Gene ative n ati n n e vise aining me ,” a xiv el wit a a e visi n,” I eling via J int Image- eat esentati n ntanglement [18] ian i i, et al , “ isi n 2022 m Nat al Lang age S sting Gene ative Image le L2024 ans n -t - n Ka as, et al , “G i ing a wei ig - es l ti n Image S nt esis”, I ti n vs Gene ati n: aming O timizati n e 2022 isi n e s A e S alable nl els”, e vise ans e able is al A- : : si n est-labs/ l x ing A ai Wang et al , “ [15] Nan e b st esentati n Alignment [12] Xingjian Leng, et al , “ - me s e ties in Sel -S e, et al , “ as e A t en , et al , “Lea ning m/bla i t a iati nal A t en e ”, a xiv els,” a xiv si n e int, 2025 els”, a xiv e int, 2025 i ient Gene ati n,” a xiv e int, 2025 e int, 2025 2024 e int, 2025

• https://github.com/black-forest-labs/flux

26 Appendix

27 RAEを用いた効率的な高解像度画像生成 • RAEのEncoderには256ピクセル，Decoderには512を対応させることで， 高解像度画像生成におけるトークン数を削減可能 512 (small) 256 esize 512 t A e トークン数は256のまま (そのままなら1024) e nst ti n 256ピクセル用のLatent s a eで 学習した i をそのまま使える！ 若干精度は下がるが トークン数は4倍効率が良い

28 Ablation study 1: DDT Head単体の効果検証 • DDT HeadはRAEのLatent spaceを使用した時に効果を発揮する • SD-VAEのLatent spaceの場合はむしろ劣化をもたらす SD-VAE RAE Down↓ UP↑ ea なし 〃 あり

29 Ablation study 2: DDT HeadをPixel spaceで検証 • ea は高次元な空間である ixel s a eでも効果があるのか？ • ixel空間でも ea は改善をもたらす • ただし， A のLatent s a eと ixel s a eの間には大きな差が存在 DDT Headなし 〃 あり UP↑ UP↑