表現オートエンコーダを備えた拡散トランスフォーマー

要旨

潜在生成モデリングにおいて、事前学習されたオートエンコーダがピクセルを拡散プロセスのための潜在空間にマッピングする手法は、Diffusion Transformers（DiT）の標準的な戦略となっている。しかし、オートエンコーダのコンポーネントはほとんど進化していない。ほとんどのDiTは、依然として元のVAEエンコーダに依存しており、これにはいくつかの制限が伴う：アーキテクチャの簡潔性を損なう時代遅れのバックボーン、情報容量を制限する低次元の潜在空間、そして純粋に再構成ベースのトレーニングに起因する弱い表現力が生成品質を最終的に制限する。本研究では、VAEを事前学習された表現エンコーダ（例：DINO、SigLIP、MAE）と訓練されたデコーダを組み合わせた、我々が「表現オートエンコーダ（RAE）」と呼ぶモデルに置き換えることを探求する。これらのモデルは、高品質な再構成と意味的に豊かな潜在空間を提供し、スケーラブルなトランスフォーマーベースのアーキテクチャを可能にする。これらの潜在空間は通常高次元であるため、拡散トランスフォーマーがそれら内で効果的に動作するための主要な課題を分析し、理論的に動機付けられた解決策を提案し、それらを実証的に検証する。我々のアプローチは、補助的な表現アライメント損失なしでより速い収束を達成する。軽量で広範なDDTヘッドを備えたDiTバリアントを使用して、ImageNetにおいて強力な画像生成結果を達成する：256x256で1.51 FID（ガイダンスなし）、256x256および512x512で1.13（ガイダンスあり）。RAEは明確な利点を提供し、拡散トランスフォーマーのトレーニングにおける新たなデフォルトとなるべきである。

English

Latent generative modeling, where a pretrained autoencoder maps pixels into a latent space for the diffusion process, has become the standard strategy for Diffusion Transformers (DiT); however, the autoencoder component has barely evolved. Most DiTs continue to rely on the original VAE encoder, which introduces several limitations: outdated backbones that compromise architectural simplicity, low-dimensional latent spaces that restrict information capacity, and weak representations that result from purely reconstruction-based training and ultimately limit generative quality. In this work, we explore replacing the VAE with pretrained representation encoders (e.g., DINO, SigLIP, MAE) paired with trained decoders, forming what we term Representation Autoencoders (RAEs). These models provide both high-quality reconstructions and semantically rich latent spaces, while allowing for a scalable transformer-based architecture. Since these latent spaces are typically high-dimensional, a key challenge is enabling diffusion transformers to operate effectively within them. We analyze the sources of this difficulty, propose theoretically motivated solutions, and validate them empirically. Our approach achieves faster convergence without auxiliary representation alignment losses. Using a DiT variant equipped with a lightweight, wide DDT head, we achieve strong image generation results on ImageNet: 1.51 FID at 256x256 (no guidance) and 1.13 at both 256x256 and 512x512 (with guidance). RAE offers clear advantages and should be the new default for diffusion transformer training.

表現オートエンコーダを備えた拡散トランスフォーマー

Diffusion Transformers with Representation Autoencoders

要旨

Support