Saining Xie 님의 엄청난 논문이 나왔다
VAE를 완전 안써버릴 대단한 논문.. RAE
https://arxiv.org/abs/2510.11690
Diffusion Transformers with Representation Autoencoders
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 contin
arxiv.org
ABSTRACT
Latent generative modeling에서 pretrained autoencoder가 pixel을 latent space로 매핑하여 diffusion 과정을 수행하는 방식은 Diffusion Transformer(DiT)의 표준 전략이 되었음.
그러나 autoencoder 자체는 거의 발전하지 않음. 대부분의 DiT는 여전히 기존 VAE encoder를 사용함.
이로 인해 구조 단순성이 저해되는 오래된 backbone, 정보 용량이 제한되는 저차원 latent space, 단순 reconstruction 기반 학습으로 인한 약한 representation 등이 존재하여 최종적으로 generative 품질이 제한됨.
본 연구는 pretrained representation encoder(DINO, SigLIP, MAE 등)와 학습된 decoder를 조합한 새로운 구조인 Representation Autoencoder (RAE)를 제안함.
- RAE는 high-quality reconstruction과 의미적으로 풍부한 latent space를 동시에 제공하며, transformer 기반 구조로 확장 가능함.
이러한 latent space는 일반적으로 고차원이므로, diffusion transformer가 이 공간 내에서 효율적으로 작동하도록 하는 것이 핵심 과제임. 본 연구는 그 어려움의 원인을 분석하고, 이론적으로 근거 있는 해결책을 제안하며, 실험적으로 이를 검증함. - 본 접근법은 auxiliary representation alignment loss 없이 faster convergence을 달성함.
- lightweight하고 넓은 DDT head를 장착한 DiT variant를 사용하여 ImageNet에서 강력한 이미지 생성 성능을 보여줌 — FID 1.51 @256×256 (no guidance), FID 1.13 @256×256 및 512×512 (with guidance).
- RAE는 clear advantages 을 제공하며, diffusion transformer training의 new default가 되어야 함.
1 INTRODUCTION
생성 모델링의 발전은 데이터 표현을 학습하는 “위치와 방식(where and how)”의 지속적인 재정의에 의해 주도되어 옴.
초기의 pixel-space 모델들은 이미지 통계를 직접적으로 포착하려 했으나, latent diffusion(Vahdat et al., 2021; Rombach et al., 2022)의 등장으로 생성 과정을 학습된 압축 표현 공간(learned compact representation space) 내에서 수행하는 방식으로 재구성됨.
이 공간에서 diffusion을 수행함으로써, Latent Diffusion Model (LDM)(Rombach et al., 2022)과 Diffusion Transformer (DiT)(Peebles & Xie, 2023; Ma et al., 2024) 같은 모델들이 더 높은 시각적 품질(visual fidelity)과 효율성을 달성하여, 현재 가장 강력한 이미지 및 비디오 생성기를 구동하게 됨.
Diffusion backbone의 발전에도 불구하고, latent space를 정의하는 autoencoder는 거의 변하지 않았음.
- 널리 사용되는 SD-VAE(Rombach et al., 2022)는 여전히 과도한 channel-wise 압축과 reconstruction-only objective에 의존함.
- 이로 인해 지역적 외형(local appearance)은 포착하지만, 일반화와 생성 성능에 중요한 전역적 의미 구조(global semantic structure)는 포착하지 못하는 저용량(latent capacity가 낮은) latent를 생성함.
- 또한 SD-VAE는 이전 세대의 convolutional 설계에 기반하여 계산 효율성 측면에서도 비효율적임(Fig. 2 참고).
한편, 시각적 표현 학습(visual representation learning)은 급격한 변화를 겪고 있음.
Self-supervised 및 multi-modal encoder들—예: DINO(Oquab et al., 2023), MAE(He et al., 2021), JEPA(Assran et al., 2023), CLIP / SigLIP(Radford et al., 2021; Tschannen et al., 2025)—은 다양한 작업과 스케일에 일반화되는 semantically structured visual features 을 학습하며, 시각적 이해의 자연스러운 기반을 제공함.
그러나 latent diffusion은 이러한 진전으로부터 거의 분리된 채 남아 있음. 여전히 의미적으로 풍부한 표현 공간 대신, 단순 reconstruction 기반 VAE 공간에서 diffusion을 수행함.
최근 연구들은 REPA(Yu et al., 2025; Yao et al., 2025; Leng et al., 2025)류의 외부 encoder 정렬(alignment)을 통해 latent 품질을 간접적으로 개선하려 시도했지만, 이는 추가적인 학습 단계, 보조 손실(auxiliary loss), 튜닝 복잡성을 초래함.
이러한 분리는 semantic objective와 generative objective 간의 비호환성이라는 오래된 가정에서 비롯됨.
즉, 의미 구조를 포착하도록 학습된 encoder는 세밀한 재구성(reconstruction)에 적합하지 않다는 인식이 일반적임 — “이들은 고수준 정보(high-level information)에 초점을 맞추며, 세부적 픽셀 유사도보다는 의미적 유사성만 복원할 수 있다”(Yu et al., 2024b).
또한 diffusion 모델이 고차원 latent space에서는 잘 작동하지 않는다는 믿음(Skorokhodov et al., 2025; Yao et al., 2025; Esser et al., 2024; Liu et al., 2024) 역시 이러한 경향을 강화시켜, 실무자들이 일반적으로 저차원 VAE latent를 선호하게 만듦.
본 연구에서는 이러한 두 가지 가정이 모두 틀릴 수 있음을 보임.
구체적으로, pretrained representation encoder가 재구성을 위한 학습이 아닌 semantics 중심으로 최적화되었더라도, 적절히 frozen된 상태에서 powerful autoencoder로 재활용될 수 있음을 증명함.
이러한 encoder는 추가적인 구조 복잡성이나 보조 손실 없이도 SD-VAE보다 우수한 reconstruction 품질을 달성함.
또한 diffusion transformer 학습은 이러한 고차원 latent space에서도 안정적이고 효율적으로 수행될 수 있음.
적절한 구조적 조정을 통해, 고차원 표현은 오히려 이점이 되어 더 풍부한 구조, 빠른 수렴, 더 나은 생성 품질을 제공함.
특히 token 수는 patch size에 의해 고정되고, 채널은 첫 layer에서 DiT hidden dimension으로 projection되므로 추가적인 계산 또는 메모리 비용이 발생하지 않음.
이를 기반으로 본 연구는 Representation Autoencoder (RAE)를 공식화함.
RAE는 기존 VAE를 pretrained representation encoder (예: DINO)와 학습된 decoder로 대체한 새로운 형태의 autoencoder로, semantic richness, structural coherence, diffusion 친화성(diffusion-friendliness)을 동시에 지님.
이로써 semantic modeling과 generative modeling을 하나의 shared latent representation으로 연결함.
그러나 이러한 고차원 semantic latent를 diffusion transformer에 적용하기 위해서는 세심한 설계가 필요함.
기존의 DiT는 compact한 SD-VAE latent를 전제로 설계되었기에, 고차원 latent에서는 다음의 어려움이 발생함:
(1) Transformer design: 모델 폭(width)이 token 차원보다 작으면 한 장의 이미지를 학습할 수도 없음 → 폭이 latent 차원에 맞게 확장되어야 함.
(2)Noise scheduling: 기존 resolution-based schedule shift(Chen, 2023; Hoogeboom et al., 2023; Esser et al., 2024)는 token 차원을 고려하지 않으므로, 차원 의존적 조정이 필요함.
(3) Decoder robustness: 기존 VAE는 연속적 분포 N(μ, σ²I)(Kingma & Welling, 2014)에서 latent를 학습하지만, RAE decoder는 discrete한 latent 분포에서 학습되므로 diffusion이 생성한 continuous latent를 복원하기 어려움 → 이를 해결하기 위해 noise-augmented decoder training을 도입함.
마지막으로, **DDT(Wang et al., 2025c)**에서 영감을 받은 새로운 DiT variant인 DiTDH를 제안함.
이는 얕고 넓은(shallow yet wide) head를 추가하여, 폭 확장을 통한 학습 안정성을 유지하면서도 quadratic한 연산 비용 증가는 방지함.
실험적으로 이러한 구조는 고차원 RAE latent에서 diffusion transformer 학습을 더욱 향상시킴.
결과적으로, RAE 기반 DiTDH는 ImageNet에서 FID 1.51(256×256, no guidance), FID 1.13(256×256 및 512×512, AutoGuidance) 성능을 달성함.
이는 RAE가 기존 VAE의 대안으로 diffusion transformer 학습에 매우 효과적임을 입증함.
더 나아가 autoencoding을 단순한 압축 메커니즘(compression mechanism)이 아닌 표현 기반(representation foundation)으로 재정의함으로써, diffusion transformer 학습을 더 효율적이고 효과적으로 만드는 새로운 패러다임을 제시함.
2 RELATED WORKS
본 절에서는 representation learning 및 reconstruction/generation 관련 기존 연구를 논의함. 보다 자세한 내용은 Appendix A에 기술함.
Representation for Reconstruction.
최근 연구들은 semantic representation을 활용하여 VAE를 향상시키는 방향을 탐구함.
예를 들어, VA-VAE(Yao et al., 2025)는 pretrained representation encoder와 VAE latent를 정렬(alignment)시키며, MAETok(Chen et al., 2025a), DC-AE 1.5(Chen et al., 2025d), l-DEtok(Yang et al., 2025)은 MAE 혹은 DAE(Vincent et al., 2008)에서 영감을 받은 objective를 VAE 학습 과정에 통합함.
이러한 정렬 방식은 VAE의 reconstruction 및 generation 성능을 크게 향상시키지만, 여전히 **강한 압축(low-dimensional latent)**에 의존하기 때문에 재구성 충실도와 표현 품질 모두 제한됨.
이에 반해 본 연구는 representation encoder feature에서 직접 복원하며, 압축을 수행하지 않음.
Frozen된 representation encoder feature 위에 단순한 ViT decoder를 얹는 방식만으로도 **SD-VAE(Rombach et al., 2022)**보다 우수하거나 동등한 reconstruction 품질을 달성하면서 훨씬 강력한 표현을 유지함.
Representation for Generation.
최근에는 semantic representation을 직접 generative modeling에 활용하려는 시도도 이루어짐.
예를 들어, REPA(Yu et al., 2025)는 DiT의 middle block을 representation encoder feature와 정렬시켜 수렴 속도를 가속화함.
DDT(Wang et al., 2025c)는 DiT를 encoder–decoder 구조로 분리하고 encoder 출력에 REPA loss를 적용함으로써 추가적인 성능 향상을 보임.
REG(Wu et al., 2025)는 DiT sequence에 learnable token을 도입하고 이를 representation encoder의 feature와 명시적으로 정렬함.
ReDi(Kouzelis et al., 2025b)는 diffusion model 내에서 VAE latent와 DINOv2 feature의 PCA component를 함께 생성함.
이에 반해 본 연구는 representation encoder의 feature 자체를 기반으로 diffusion model을 직접 학습하며, 훨씬 빠른 수렴을 달성함.
또한 보다 폭넓은 관점에서의 논의는 **Dieleman (2025)**을 참고함.
3 HIGH FIDELITY RECONSTRUCTION FROM FROZEN ENCODERS
본 절에서는 pretrained representation encoder(예: DINOv2(Oquab et al., 2023), SigLIP2(Tschannen et al., 2025))가 reconstruction task에는 부적합하다는 일반적 가정을 비판함.
이들 encoder는 “고수준 의미(high-level semantics)에 집중하고 저수준 세부(low-level detail)를 약화시킨다”(Tang et al., 2025; Yu et al., 2024b)는 이유로 재구성에 부적절하다고 여겨졌음.
그러나 본 연구는 적절히 학습된 decoder를 결합할 경우, frozen representation encoder가 diffusion latent space의 강력한 encoder로 기능할 수 있음을 보임.
제안하는 Representation Autoencoder (RAE)는 frozen pretrained encoder와 ViT 기반 decoder를 조합하여 SD-VAE와 동등하거나 더 우수한 재구성을 달성함.
더 중요한 점은, RAE는 강한 압축으로 인해 표현 품질이 제한되는 VAE의 근본적 한계(Kingma & Welling, 2014)를 완화함.
예를 들어 SD-VAE는 256² 이미지를 32²×4 latent로 매핑(Esser et al., 2021; Rombach et al., 2022)함으로써 재구성 품질뿐 아니라 표현력 역시 제한함.
RAE decoder 학습 과정은 다음과 같음.
입력 \(x \in \mathbb{R}^{3 \times H \times W}\) 에 대해, patch size \((p_e)\)와 hidden size \(d\)를 갖는 frozen encoder E를 통해 \(N = HW / p_e^2\)개의 token을 얻음.
ViT decoder (D)는 patch size (p_d)를 사용해 이를 다시 pixel로 복원함.
기본적으로 p_d = p_e로 설정하여 입력 해상도와 동일한 크기를 유지함.
256×256 이미지 실험에서 encoder는 256개의 token을 생성하며, 이는 대부분의 SD-VAE 기반 DiT(Peebles & Xie, 2023; Yu et al., 2025; Ma et al., 2024)에서 사용된 token 수와 동일함.
Decoder는 **L1 + LPIPS(Zhang et al., 2018) + adversarial loss(Goodfellow et al., 2014)**를 조합하여 학습함:
\(z = E(x), \hat{x} = D(z)\)
\(\mathcal{L}{rec}(x) = \omega_L \text{LPIPS}(\hat{x}, x) + L_1(\hat{x}, x) + \omega_G \lambda{GAN}(\hat{x}, x)\)
Decoder 구조, 하이퍼파라미터, loss 계수, GAN 학습 세부 사항은 Appendix C에 기술함.
RAE는 세 가지 대표 encoder를 실험에 사용함:
- DINOv2-B(Oquab et al., 2023): self-distillation 기반 self-supervised 모델, (p_e=14, d=768)
- SigLIP2-B(Tschannen et al., 2025): language-supervised 모델, (p_e=16, d=768)
- MAE-B(He et al., 2021): masked autoencoder, (p_e=16, d=768)
DINOv2에 대해서는 S/B/L (d=384/768/1024) 모델 크기 비교도 수행함.
모든 RAE는 기본적으로 ViT-XL decoder를 사용함.
Reconstruction 품질 평가는 ImageNet validation set에서 **FID(Heusel et al., 2017)**를 사용하며, 이를 rFID로 표기함.
Reconstruction, scaling, and representation.
Table 1a에서 보듯,
frozen encoder를 사용한 RAE는 **SD-VAE보다 일관적으로 더 나은 재구성 품질(rFID)**을 달성함.
예를 들어, MAE-B/16 기반 RAE는 rFID 0.16을 기록하여 SD-VAE를 명확히 능가함과 동시에
“representation encoder는 pixel-level 세부를 복원할 수 없다”는 기존 가정을 반박함.
다음으로, encoder와 decoder의 scaling behavior를 분석함.
Table 1c에서 보듯,
DINOv2-S, B, L 간의 재구성 품질은 안정적으로 유지되며, 이는 작은 representation encoder도 충분한 low-level detail을 보존함을 의미함.
Decoder 측면(Table 1b)에서는 모델 용량을 늘릴수록 rFID가 지속적으로 향상됨 — ViT-B에서 0.58 → ViT-XL에서 0.49로 개선됨.
특히 **ViT-B는 SD-VAE보다 14배 효율적(GFLOPs 기준)**이면서도 이미 더 높은 품질을 달성하며, ViT-XL은 SD-VAE의 연산 비용의 1/3만으로 품질을 추가로 향상시킴.
또한 ImageNet-1K에서 linear probing을 수행하여 표현 품질을 평가함(Table 1d).
RAE는 frozen pretrained encoder를 사용하므로, 기저 representation encoder의 강력한 표현력을 직접 계승함.
반면 SD-VAE는 약 8% 수준의 정확도만을 보임.
4 TAMING DIFFUSION TRANSFORMERS FOR RAE
RAE가 우수한 reconstruction 품질을 보여주었으므로, 이제 그 latent space의 diffusability(Skorokhodov et al., 2025), 즉 diffusion model이 얼마나 쉽게 그 latent 분포를 모델링할 수 있으며, 생성 품질이 얼마나 좋은지를 조사함.
Before turning to generation, encoder를 fix한 채 latent space의 generation capabilities을 탐구함.
Table 1a에 따르면 MAE, SigLIP2, DINOv2 모두 SD-VAE보다 낮은 reconstruction rFID를 달성하며, 이 중 MAE가 가장 우수함.
그러나 단순히 reconstruction이 좋다고 해서 생성 성능이 보장되지는 않음.
실험적으로는 DINOv2가 가장 강력한 generation 성능을 보여주므로, 이후 모든 실험에서는 DINOv2를 기본 encoder로 채택함(단, 전체 비교는 Section G.1에 제시함).
표준 방식에 따라, 본 연구는 Flow Matching Objective(Lipman et al., 2023; Liu et al., 2023)를 사용함.
즉, \(x_t = (1 - t)x + t\varepsilon) ((x \sim p(x)), (\varepsilon \sim N(0, I))\) 형태의 linear interpolation을 수행하고,
모델이 velocity \((v(x_t, t))\) 를 예측하도록 학습함(Section J 참고).
모델 backbone은 LightningDiT(Yao et al., 2025)로, 이는 DiT(Peebles & Xie, 2023)의 변형 구조임.
Patch size는 1로 설정하여, 256×256 이미지에서 모든 RAE가 sequence length = 256이 되도록 함.
이는 VAE 기반 DiT(Peebles & Xie, 2023; Yu et al., 2025; Yao et al., 2025)에서 사용하는 token 수와 동일함.
DiT의 계산량은 sequence length에 주로 의존하므로, RAE latent를 사용하는 것은 VAE 기반 방법에 비해 추가적인 연산 오버헤드를 유발하지 않음.
모델 평가는 **Euler sampler (50 steps)**로 50K 이미지를 생성한 뒤 **FID(gFID)**로 측정함.
모든 정량 실험은 특별히 명시되지 않는 한 ImageNet 256×256에서 80 epoch 동안 학습함.
세부적인 학습 설정은 Section D에 기술함.
그러나 놀랍게도, 표준 diffusion recipe는 RAE latent에 그대로 적용될 경우 실패함(Table 2 참고).
RAE latent에서 직접 학습할 경우, 작은 backbone(예: DiT-S)은 완전히 학습이 불가능하며,
큰 backbone(예: DiT-XL)조차 SD-VAE latent 기반 counterpart보다 성능이 크게 떨어짐.
이러한 현상을 분석하기 위해 세 가지 가설을 설정함:
- • Suboptimal design for diffusion transformers.
RAE token은 high-dimensional feature이므로, diffusion transformer의 최적 구조가
저차원 VAE token에 맞춰 설계된 기존 DiT와 달라질 수 있음. - • Suboptimal noise scheduling.
기존의 noise schedule 및 loss re-weighting 기법은 image 또는 VAE latent를 전제로 설계됨.
이 방식이 high-dimensional semantic token에 그대로 적용 가능한지는 명확하지 않음. - • Diffusion generates noisy latents.
VAE decoder는 noisy latent에서도 이미지를 복원하도록 학습되어 있어 diffusion output에 강건함.
반면 RAE decoder는 clean latent만으로 학습되므로, noise가 포함된 latent에 대해 일반화가 어려움.
4.1 SCALING DIT WIDTH TO MATCH TOKEN DIMENSIONALITY
RAE latent 기반 diffusion transformer의 training dynamics을 이해하기 위해 simplified experiment 을 구성함.
ImageNet 전체가 아닌, a single image를 randomly select하여 RAE로 인코딩하고, diffusion model이 이를 reconstruct할 수 있는지를 테스트함.
Table 2에 따르면 RAE는 SD-VAE보다 reconstruction 성능이 낮지만, DiT의 폭(width)을 늘릴수록 성능이 향상됨.
이를 분석하기 위해 depth를 고정하고 width만 변화시킴.
DiT-S 구조에서 hidden dimension을 384 → 784로 증가시키면,
Figure 3에서 보이듯 **모델 폭이 token dimension(n = 768)**보다 작을 때는 품질이 매우 낮지만,
폭이 token dimension 이상이 되면 학습이 급격히 안정화되며 입력 이미지를 거의 완벽히 복원함.
학습 loss 역시 같은 경향을 보이며, (d ≥ n)일 때만 수렴함.
이 개선이 단순히 parameter 수 증가 때문인지 확인하기 위해, 폭을 (d = 384)로 고정하고 depth만 12 → 24로 두 배 늘림.
그러나 Figure 3에서 보이듯이 여전히 생성 이미지에는 artifact가 남고, 학습 loss도 폭을 늘린 경우만큼 수렴하지 않음.
즉, RAE latent에서의 안정적인 generation을 위해서는 diffusion 모델의 폭이 최소한 token 차원 이상이어야 함.\
이는 데이터 manifold가 저차원 구조를 가진다는 일반적인 믿음(Pope et al., 2021; Sauer et al., 2022)과 상충되는 듯 보이지만,
diffusion model의 수식적 구조 때문임.
diffusion은 학습 과정에서 데이터에 Gaussian noise를 직접 주입하므로, 데이터 manifold의 support가 전체 공간으로 확장됨.
결과적으로 모델은 전체 차원에 걸친 표현 능력을 가져야 하며, 따라서 폭이 token 차원 이상이어야 함.
이 현상은 다음의 Theorem 1로 이론적으로 정당화됨:
여기서 (d < n)이고, (f)는 폭이 token dimension보다 작은 DiT block의 집합을 나타냄.
이때, 모든 (g \in G_d)에 대해 다음이 성립함:
[
L(g, θ) = \int_0^1 E_{x, \varepsilon}[| g(x_t, t) - (\varepsilon - x) |^2] dt \geq \sum_{i=d+1}^{n} \lambda_i
]
여기서 (\lambda_i)는 (W = \varepsilon - x)의 공분산 행렬의 고유값임.
특히 (d ≥ n)일 때 (G_d)는 (L(g, θ))의 유일한 최소해를 포함함 (증명은 Section B.1 참조).
toy setting에서 (p(x) = \delta(x - x_0))일 경우 (W ∼ N(-x, I_n)), 따라서 (\lambda_i = 1).
이에 따라 평균 loss의 하한은 \( \tilde{L}(g, θ) ≥ (n-d)/n \).
이는 Figure 3에서의 실험 결과와 일치함.
확장된 실험
보다 실제적인 설정으로, 폭이 서로 다른 세 모델 {DiT-S, DiT-B, DiT-L}과
encoder 크기가 다른 {DINOv2-S, DINOv2-B, DINOv2-L} 조합을 각각 overfit시킴.
Table 3에서 보이듯, 모델 폭이 token dimension 이상일 때만 수렴이 발생함(예: DiT-B + DINOv2-B),
그 외 조합(예: DiT-S + DINOv2-B)에서는 loss가 수렴하지 않음.
이에 따라 이후 모든 실험에서는 DiT 폭을 RAE token dimension 이상으로 설정하고,
DINOv2-B encoder에는 DiT-XL을 사용함.
4.2 DIMENSION-DEPENDENT NOISE SCHEDULE SHIFT
이전 연구들(Teng et al., 2023; Chen, 2023; Hoogeboom et al., 2023; Esser et al., 2024)은
입력 \(z ∈ \mathbb{R}^{C×H×W}\) 에서 해상도(H×W)가 커질수록 동일한 noise level에서 정보 손실이 줄어들어 diffusion 학습을 방해한다고 보고함.
하지만 이러한 결과는 대부분 pixel-space 혹은 저차원 VAE latent with few channels(C ≤ 16)에 한정됨.
실제 Gaussian noise는 spatial과 channel 양쪽 차원에 모두 적용되므로, channel 수가 늘어날수록
각 token이 담고 있는 유효 해상도(effective resolution)도 커지고 정보 손실은 더 줄어듦.
따라서 기존의 “해상도 기반 스케줄 조정(resolution-dependent schedule shift)”은
보다 일반화된 형태인 effective data dimension (token 수 × channel dimension) 기반으로 확장되어야 함.
이를 위해 **Esser et al. (2024)**의 shifting 전략을 채택함.
base schedule \(t_n ∈ [0, 1]\)과 input dimensions n, m에 대해
shifted timestep \(t_m = \frac{α t_n}{1 + (α - 1)t_n}, \quad α = \sqrt{m/n}\)
로 정의함.
(n = 4096)을 기준 차원으로 두고, (m)을 RAE latent의 실제 유효 차원으로 설정함.
Table 4에서 보이듯, 이 dimension-dependent 스케줄은 성능 향상을 크게 유도하며,
특히 high-dimensional RAE latent에서 diffusion 학습 안정성과 품질을 동시에 개선함.
이후 모든 실험에서는 차원 의존형 noise schedule을 기본으로 사용함.
내 이해 :
여기서 m= 실제 차원, n= 기준 차원(4096).
차원이 커질수록 α>1, 따라서 t_m<t_n(즉, 시간 축이 “늘어남”).
이는 diffusion이 노이즈를 천천히 추가하도록 시프트하는 효과를 냄.
결과적으로:
“노이즈가 약하게 느껴지는” 고차원 공간에서도
“노이즈가 적당히 섞인 상태”를 유지하도록 스케줄을 재보정하는 것임.
###########################
차원이 커질수록 노이즈 약하게 적용해야한다는것임
고차원에서는 정보 손실이 덜 발생되는데
노이즈를 차원마다 독립적으로 계속 비례해서 증가해주면 안된다는것
고차원에서 여러 차원에서 노이즈가 상쇄되면서 약하게 느껴짐
디퓨전모델이 timestep t 클수록 더 강한 노이즈 섞음
근데 고차원에서는 같은 t값에서도 데이터가 덜 깨지니까, 기존 출력대로 가면 너무 빨리 데이터 파괴 구간에 도달한ㄷ는것
아그니까 고차원에서는 노이즈가 약하게 느껴지는데 원본이미지같은데 / 기존대로 노이즈 강하게 섞으면 너무 빨리 이미지가 망가져버려
그래서 t 증가속도를 늦춰서 노이즈 더 천천히 누적시키는것
4.3 NOISE-AUGMENTED DECODING
VAE에서는 latent가 연속적 분포 (N(μ, σ^2I))로 표현되지만(Kingma & Welling, 2014),
RAE decoder D는 discrete 분포 \(p(z) = \sum_i δ(x - z_i)\)에서 이미지를 복원하도록 학습됨.
여기서 ({z_i})는 RAE encoder E로부터 생성된 training sample.
하지만 inference 시 diffusion model이 생성하는 latent는 완벽히 clean하지 않고,
학습 분포에서 약간 벗어나거나 noisy할 수 있음(Abuduweili et al., 2024).
이는 decoder에 out-of-distribution 문제를 일으켜 샘플 품질을 저하시킴.
이를 완화하기 위해 Normalizing Flow 관련 연구(Dinh et al., 2017; Ho et al., 2019; Zhai et al., 2025)에서 영감을 받아
decoder 학습 시 additive Gaussian noise (n ∼ N(0, σ^2I))를 추가함.
즉, clean latent p(z) 대신 smoothed distribution
\(p_n(z) = \int p(z - n) N(0, σ^2I)(n) , dn\)
에서 학습함으로써 diffusion model의 연속적 output에 대한 일반화 능력을 향상시킴.
또한 σ를 고정하지 않고 |N(0, τ²)|에서 stochastic하게 샘플링하여학습 regularization과 robustness를 동시에 개선함.
4.3 NOISE-AUGMENTED DECODING (continued)
(p_n(z))가 reconstruction과 generation에 어떤 영향을 미치는지 분석함.
Table 5에 따르면, (p_n(z))는 gFID를 개선하지만 rFID는 다소 악화시킴.
이러한 trade-off는 예상된 결과임.
노이즈를 추가하면 latent 분포가 부드러워지고, decoder의 out-of-distribution(OOD) 문제를 줄여줌.
하지만 동시에 세밀한 디테일이 손실되어 reconstruction 품질이 약간 저하됨.
Section G.2에서 τ 값과 다양한 encoder 조합에 대한 추가적인 ablation 실험을 수행함.
-->가우시안으로 퍼져서 매끄러운 연속분포로 변해짐!!!
p_n(z): ~~~ ~~~~ ~~~~ ~~~~
(각 점이 주변으로 퍼져서 매끄러운 연속 분포가 됨)
• Diffusion generates noisy latents.
이후 모든 실험에서는 noise-augmented decoding을 기본적으로 적용함.
이상의 모든 기법을 결합하여 RAE latent 기반 DiT-XL을 학습함.
개선된 diffusion recipe는 단 80 epoch 학습만으로 **gFID 4.28 (Figure 4)**을 달성하며,
720 epoch 시점에는 gFID 2.39로 수렴함.
같은 모델 크기 기준으로, 이는 VAE latent 기반으로 학습된 기존 diffusion 모델(예: SiT-XL (Ma et al., 2024))보다 47배 빠른 학습 속도를 보이며,
representation alignment 기반의 최신 방법(REPA-XL (Yu et al., 2025))보다도 약 16배 빠른 수렴 속도를 달성함.
이후 절에서는 RAE 기반 generation을 더욱 효율적이고 효과적으로 만드는 방안을 분석함으로써,
state-of-the-art 수준의 성능으로 확장하는 방법을 제시함.
5 IMPROVING THE MODEL SCALABILITY WITH WIDE DIFFUSION HEAD
앞선 Section 4에서 논의했듯, 표준 DiT 프레임워크에서 RAE latent의 **고차원성(high-dimensionality)**을 처리하려면
backbone 전체의 폭(width)을 확장해야 함.
그러나 이는 연산 비용이 급격히 증가함.
이를 해결하기 위해 DDT (Wang et al., 2025c)에서 영감을 받아 Wide Diffusion Head (DDT head)를 도입함.
이는 얕지만 폭이 넓은 transformer module로, diffusion 과정에서의 denoising 전용 역할을 수행함.
이 head를 표준 DiT에 결합하면, 전체 FLOPs의 quadratic 증가 없이 폭을 효과적으로 확장할 수 있음.
이 확장된 구조를 DiTDH (Diffusion Transformer with DDT Head)라 부르며,
Section G.3에서 DDT head의 구체적인 설계 선택에 대한 실험을 제시함.
Wide DDT head.
형식적으로, DiTDH 모델은 기본 DiT backbone M과 얕고 넓은 transformer head H로 구성됨.
주어진 noisy 입력 (x_t), timestep (t), 선택적 class label (y)에 대해, 모델은 다음과 같이 velocity (v_t)를 예측함:
- xtx_t: 현재 timestep tt에서의 noisy latent (RAE latent에 노이즈가 더해진 상태)
- M: 기본 DiT backbone
- : wide DDT head (얕지만 폭이 넓은 transformer)
- y: 선택적 class conditioning
- z_t: DiT backbone이 추출한 feature embedding
- v_t: 모델이 최종적으로 예측하는 velocity (노이즈 제거 방향)
DiTDH converges faster than DiT.
RAE latent 상에서 backbone 크기를 달리한 일련의 DiTDH 모델(DiTDH-S, B, L, XL)을 학습함.
모든 DiTDH 모델에서 DDT head는 2-layer, 2048-dim 구조를 사용함.
성능은 표준 DiT-XL baseline과 비교함.
Figure 6a에 따르면, DiTDH는 DiT보다 FLOPs 효율이 훨씬 높음.
예를 들어, DiTDH-B는 DiT-XL 대비 약 40%의 학습 FLOPs만 사용하면서도 훨씬 더 우수한 성능을 보임.
비슷한 연산 예산 하에서 DiTDH-XL은 FID 2.16을 달성하며,
이는 DiT-XL의 약 절반 수준의 FID에 해당함.
DiTDH maintains its advantage across RAE scales.
세 가지 RAE encoder (DINOv2-S, DINOv2-B, DINOv2-L)를 사용하여 DiTDH-XL과 DiT-XL을 비교함.
Table 6에 따르면, DiTDH는 모든 encoder 스케일에서 DiT보다 일관되게 우수하며,
encoder 규모가 커질수록 그 격차가 더욱 확대됨.
예를 들어, DINOv2-L 기반 실험에서 DiT-XL의 FID는 6.09,
DiTDH는 이를 2.73으로 개선함.
이러한 성능 향상은 DDT head의 구조적 강건성에 기인함.
대형 encoder는 더 높은 차원의 latent를 생성하므로,
기존 DiT의 폭 제한(width bottleneck) 문제가 더욱 심화됨.
DiTDH는 Section 4에서 논의된 폭 조건을 충족시키면서 feature 표현을 compact하게 유지함.
또한, 고차원 latent에서 더 두드러지는 noisy 정보를 효과적으로 필터링함으로써
RAE 기반 diffusion 모델의 안정성과 효율성을 동시에 개선함.
5.1 STATE-OF-THE-ART DIFFUSION TRANSFORMERS
Convergence.
DiTDH-XL의 수렴 거동을 기존 state-of-the-art diffusion 모델들과 비교함 (Peebles & Xie, 2023; Ma et al., 2024; Yu et al., 2025; Gao et al., 2023; Yao et al., 2025).
비교는 guidance 없이 FID 기준으로 수행함.
Figure 6b는 DiTDH-XL의 FID 수렴 곡선을 training epoch 및 GFLOPs 단위로 나타냄.
기존 baseline 모델들은 각 논문에서 보고된 최종 성능 지점으로 표시함.
DiTDH-XL은 약 5 × 10¹⁰ GFLOPs 시점에서 이미 REPA-XL, MDTv2-XL, SiT-XL을 모두 초월하며,
5 × 10¹¹ GFLOPs 시점에서는 전체 모델 중 최저 FID를 기록함.
이는 기존 모델 대비 40배 이상 적은 연산량으로 동일하거나 더 나은 성능을 달성함을 의미함.
Scaling.
서로 다른 모델 규모에서의 DiTDH와 최신 방법들을 비교함.
Figure 6c에 따르면, DiTDH의 크기를 키울수록 FID가 일관되게 개선됨.
가장 작은 모델인 DiTDH-S는 이미 FID 6.07을 달성하며,
이는 훨씬 더 큰 REPA-XL보다도 우수함.
DiTDH-S에서 DiTDH-B로 스케일링하면 FID가 6.07 → 3.38로 크게 향상되며,
이는 유사하거나 더 큰 규모의 모든 기존 모델을 능가함.
성능은 모델 크기가 커질수록 지속적으로 개선되어,
DiTDH-XL은 80 epoch 학습만으로 FID 2.16을 달성하며 새로운 state-of-the-art 결과를 확립함.
Performance.
마지막으로, 가장 강력한 모델인 DiTDH-XL을
최신 diffusion 모델들과 ImageNet 256×256 및 512×512 해상도에서 정량적으로 비교함 (Table 7, 8 참조).
결과적으로, 본 방법은 모든 기존 diffusion 모델을 큰 폭으로 초월하며,
256×256 해상도에서 FID 1.51 (no guidance), **FID 1.13 (with guidance)**를 기록함.
이는 기존 SOTA보다 상당히 낮은 수치임.
또한 512×512 해상도에서 400 epoch 학습 시 FID 1.13 (with guidance)를 달성하며, 이전 최고 성능인 **EDM-2 (FID 1.25)**를 초과 달성함.
정량적 결과 이후, Figure 7에 DiTDH-XL의 생성 샘플을 제시함.
이 시각적 결과들은 높은 의미적 다양성(semantic diversity)**과 정교한 세부 표현(fine-grained detail을 모두 보여주며,
실제 ImageNet 이미지와 유사한 수준의 품질을 달성함.
추가적인 시각화 예시는 Section M, unconditional generation 결과는 Section L에 포함됨.
Remarks on FID evaluation.
conditional FID 평가를 위해 사용된 50,000개 샘플의 구성 방식을 명시함.
기존 연구들(DDT(Wang et al., 2025c), VAR(Tian et al., 2024), MAR(Li et al., 2024b), xAR(Ren et al., 2025))은
일반적으로 클래스당 50개 이미지를 균등하게 샘플링함.
반면 일부 연구들은 1,000개 클래스 전체에서 uniform random sampling을 사용함.
이론적으로 uniform sampling은 class-balanced sampling의 근사값이지만,
실험적으로는 class-balanced sampling이 항상 약 0.1 정도 더 낮은 FID를 기록함을 관찰함.
공정한 비교를 위해, 접근 가능한 checkpoint가 존재하는 일부 최근 모델들(SiT(Peebles & Xie, 2023), REPA(Yu et al., 2025), REPA-E(Leng et al., 2025))을
모두 class-balanced sampling으로 재평가하였으며,
이에 따라 FID 점수를 업데이트함.
보다 상세한 비교 내용은 Section E에 제시함.
6 DISCUSSIONS
6.1 HOW CAN RAE EXTEND TO HIGH-RESOLUTION SYNTHESIS EFFICIENTLY?
고해상도 이미지를 생성할 때의 주요 과제는 해상도가 토큰 수와 함께 증가한다는 점임.
즉, 이미지의 가로·세로 크기를 2배로 늘리면 필요한 토큰 수는 약 4배 증가함.
이를 해결하기 위해, **decoder가 해상도 확장(resolution scaling)**을 담당하도록 설계함.
즉, decoder의 **patch size (p_d)**를 encoder의 **patch size (p_e)**와 다르게 설정함.
- (p_d = p_e)이면 출력 해상도는 입력과 동일함.
- (p_d = 2p_e)로 설정하면, 동일한 latent token을 사용해 2× 업샘플링된 512×512 이미지를 복원함.
decoder는 encoder와 diffusion 과정 모두와 **완전히 분리(decoupled)**되어 있으므로,
256×256 해상도에서 학습된 diffusion 모델을 그대로 재사용 가능함.
즉, 기존 diffusion 모델은 그대로 두고 decoder만 업샘플링용으로 교체하면
재학습 없이 512×512 출력을 얻을 수 있음.
Table 9에 따르면, 이 방법은 rFID가 약간 증가하지만 gFID는 경쟁력 있는 수준을 유지하며,
토큰 수를 4배로 늘리는 방식보다 4배 이상 효율적임.
6.2 DOES DiTDH WORK WITHOUT RAE?
본 연구에서는 RAE와 DiTDH를 함께 제안함.
Section 4에서 이미, RAE만으로도 DiT의 성능을 크게 향상시킬 수 있음을 보였음.
그렇다면 반대로, RAE 없이도 DiTDH가 여전히 유효한가?라는 질문이 생김.
이를 검증하기 위해, SD-VAE latent(patch size 2)를 사용하여
DiT-XL과 DiTDH-XL을 각각 학습함(비교를 위해 DINOv2-B도 포함).
학습은 80 epoch 동안 수행하고, unguided FID로 평가함.
Table 10에서 보듯, SD-VAE latent에서는 DiTDH-XL이 오히려 DiT-XL보다 성능이 나쁨.
이는 diffusion head가 추가적인 연산량을 초래했음에도 불구하고 개선 효과를 내지 못함을 의미함.
따라서 DDT head는 저차원 latent 공간에서는 거의 이점을 제공하지 않으며,
그 주요 강점은 RAE가 제공하는 고차원 latent 기반의 diffusion task에서 발휘됨.
6.3 THE ROLE OF STRUCTURED REPRESENTATION IN HIGH-DIMENSIONAL DIFFUSION?
DiTDH는 RAE의 고차원 latent 공간과 결합될 때 매우 강력한 성능을 보임.
따라서 중요한 질문이 제기됨
“RAE의 구조적 표현(structured representation)이 필수적인가,
아니면 unstructured high-dimensional 입력(예: raw pixel)에서도 DiTDH가 잘 작동하는가?”
This raises a key question:
is the structured representation of RAE essential, or would DiTDH work equally well on unstructured high-dimensional inputs such as raw pixels?
이를 평가하기 위해, DiT-XL과 DiTDH-XL을 raw pixel에 직접 학습함.
256×256 이미지에서 patch size를 16으로 설정하면, DiT 입력 토큰의 차원은 (16 × 16 × 3 = 768)이 되어 DINOv2-B latent의 차원과 동일함.
이 설정에서 80 epoch 학습 후 unguided FID를 측정함.
Table 11에 따르면, DiTDH는 pixel 입력에서도 DiT보다 약간 우수하지만, 두 모델 모두 RAE latent에서 학습된 경우보다 훨씬 열악한 성능을 보임.
이 결과는 단순히 차원이 높다고 해서 좋은 성능을 얻을 수 있는 것이 아님을 보여줌.
-> RAE가 제공하는 구조화된 표현(structured representation)이 고차원 diffusion에서의 성능 향상에 핵심적인 역할을 함.
7 CONCLUSION
본 연구에서는 pretrained representation encoder는 재구성(reconstruction)이나 생성(generation)에 적합하지 않다는 기존의 통념을 반박함.
우리는 frozen representation encoder에 경량의 학습된 decoder를 결합하여,
효율적이고 표현력 있는 Representation Autoencoder (RAE)를 구성할 수 있음을 보임.
이 latent space 위에서 diffusion transformer를 안정적이고 효율적으로 학습하기 위해 세 가지 핵심 구성 요소를 도입함:
(1) DiT의 폭(width)을 encoder의 token 차원에 맞춤,
(2) noise schedule에 차원 의존형(dimension-dependent) shift를 적용,
(3) decoder가 diffusion 출력을 처리할 수 있도록 noise augmentation을 추가함.
또한, 계산량의 제곱적 증가(quadratic compute) 없이 폭을 확장할 수 있는
얕지만 넓은 구조의 diffusion transformer head인 DiTDH (Diffusion Transformer Dual Head)를 제안함.
실험적으로, RAE 기반 DiTDH는 강력한 시각 생성 능력을 입증함.
ImageNet에서 RAE 기반 DiTDH-XL은
256×256 해상도에서 FID 1.51 (no guidance), 256×256 및 512×512 모두에서 FID 1.13 (with guidance)를 달성함.
이러한 결과는 RAE latent가 향후 diffusion transformer 학습을 위한 강력한 기반 후보(strong candidate)가 될 수 있음을 시사함.
