📌 Paper Info
- Title: Diffusion Transformers with Representation Autoencoders
- Authors: Boyang Zheng, Nanye Ma, Shengbang Tong, Saining Xie
- Venue / Year: arXiv / 2025
- Link: https://arxiv.org/abs/2510.11690
🧠 Summary
DINO, SigLIP, MAE같은 representation 기반 frozen encoder를 사용해서 decoding하자
✅ Key Points
Motivation
- 구식 CNN 기반 구조 — 비효율
- 저차원 latent — 저차원이라 세밀한 표현이 어려움
- recon 중심 학습 — semantic 구조가 약해서 표현력이 제한됨
Method
Stage 1: Representation Autoencoder (RAE)
- Encoder: DINOv2, SigLIP2, MAE 등 pretrained ViT 기반 representation encoder (freeze)
- Decoder: 학습 가능한 ViT 구조, L1 + LPIPS + GAN loss로 학습
- No compression: latent를 줄이지 않고 그대로 유지 (e.g., 256 tokens, 768-dim)
Stage 2: Training Diffusion Transformer on RAE latents
💭 Code
핵심코드 간단히 재구현 (Official 기반 chatgpt 재구현)
# ========================= RAE =========================
class RAE(nn.Module):
"""
Representation AutoEncoder (Stage 1)
- 논문 제안 요약:
• Pretrained representation encoder (DINOv2/SigLIP/MAE 등)을 freeze한 채 사용
• MAE-style ViT decoder만 학습
• Decoder 학습 시 Noise-Augmented Decoding 적용 (Section 4.3)
→ z' = z + σ·ε, σ ~ Uniform(0, τ), ε ~ N(0, I)
→ Decoder가 noisy latent에도 견고하도록 훈련
(ref: RAE-DiT §4.3, Table 5, Table 15, Fig. 4)
"""
def __init__(self, encoder, decoder, noise_tau=0.8, reshape_to_2d=True, stats=None):
super().__init__()
self.encoder = encoder.eval().requires_grad_(False) # pretrained DINOv2/SigLIP/MAE (frozen)
self.decoder = decoder # MAE-style ViT decoder (trainable)
self.noise_tau = noise_tau
self.reshape_to_2d = reshape_to_2d
self.stats = stats or {}
@torch.no_grad()
def encode(self, x):
if "enc_mean" in self.stats:
x = (x - self.stats["enc_mean"]) / (self.stats["enc_std"] + 1e-6)
return self.encoder(x)
def _maybe_noise(self, z):
# Noise-Augmented Decoding (RAE-DiT §4.3)
if self.training and self.noise_tau > 0:
sigma = self.noise_tau * torch.rand((z.size(0),) + (1,)*(z.ndim-1), device=z.device)
eps = torch.randn_like(z)
z = z + sigma * eps
return z
def decode(self, z):
z = self._maybe_noise(z)
if self.reshape_to_2d and z.dim() == 4:
b, c, h, w = z.shape
z = z.view(b, c, h*w).transpose(1, 2)
return self.decoder(z) # MAE-style ViT decoder# ======================== DiTDH =========================
class DiTDH(nn.Module):
"""
DiT with Deep Backbone + Shallow, Wide Head
- 논문 제안 요약:
• DDT(=DiTDH) head를 "wide & shallow"하게 설계해야 FID가 개선됨
• Head width는 ≥2048, depth는 2가 최적
• RAE latent 크기에 따라 head 효과가 달라짐
(ref: Appendix G.3, Table 16–17)
"""
def __init__(self, in_channels, base_hidden=1152, base_depth=28,
head_hidden=2048, head_depth=2, num_heads=16):
super().__init__()
self.backbone = LightningDiT(in_channels, hidden_size=base_hidden,
depth=base_depth, num_heads=num_heads)
self.proj_in = nn.Linear(base_hidden, head_hidden)
self.head = nn.ModuleList([LightningDDTBlock(head_hidden, num_heads)
for _ in range(head_depth)])
self.proj_out = nn.Linear(head_hidden, in_channels)
def forward(self, x_tokens, t, condition=None):
z = self.backbone(x_tokens, t, condition)
h = self.proj_in(z)
for blk in self.head:
h = blk(h, t, condition)
return self.proj_out(h)# =============== time-shift helpers (논문 §4.2) ===============
def compute_time_shift(latent_size, base_dim=4096):
"""
Dimension-dependent schedule shift (RAE-DiT §4.2)
α = sqrt(m_eff / m_base)
- m_eff = C×H×W (latent의 총 차원)
- m_base = 4096 (VAE 기반 diffusion에서 안정적이었던 기준 차원)
"""
C, H, W = latent_size
return ((C * H * W) / float(base_dim)) ** 0.5
def remap_time_grid(t0, t1, steps, shift):
"""
시간 리매핑 함수 (Eq. 11)
t' = α·t / (1 + (α - 1)·t)
- 고차원 latent에서 noise 분포를 보정하기 위해 시간축을 압축
(ref: RAE-DiT §4.2, Table 4)
"""
t = 1.0 - torch.linspace(t0, t1, steps, device="cuda")
return shift * t / (1 + (shift - 1) * t)
# ======================= Transport =======================
class Transport(nn.Module):
"""
Flow Matching 기반 Transport module (Stage 2)
- 논문 제안 요약:
• 항상 compute_time_shift()를 통해 α 계산 → dimension-dependent schedule shift 적용
• remap_time_grid()로 시간 스케줄 리매핑
• velocity prediction 기반 drift: dxt/dt = vθ(xt, t)
(ref: RAE-DiT §4.2, §4.3, Appendix J)
"""
def __init__(self, latent_size, base_dim=4096,
prediction="velocity", path_type="Linear", loss_weight="velocity",
time_dist_type="uniform", train_eps=0.0, sample_eps=0.0):
super().__init__()
self.model_type = prediction
self.path_type = path_type
self.loss_type = loss_weight
self.time_dist_type = time_dist_type
self.train_eps = train_eps
self.sample_eps = sample_eps
self.latent_size = tuple(latent_size)
self.base_dim = base_dim
self.time_dist_shift = float(compute_time_shift(self.latent_size, self.base_dim)) # α
def remap_t(self, t):
"""단일 t를 α 기반으로 리매핑"""
α = self.time_dist_shift
return α * t / (1 + (α - 1) * t)
def sample_ode(self, model, xT, steps=50, atol=1e-5, rtol=1e-5, condition=None):
"""
Neural ODE integration for generation
- 식: dx_t/dt = vθ(x_t, t)
- odeint로 t=1→0 적분 → latent 복원 → RAE decoder로 이미지 복원
(ref: Appendix J, Eq. 13–14)
"""
α = self.time_dist_shift
t = remap_time_grid(0.0, 1.0, steps, α)
samples = odeint(lambda x, tt: self._drift(x, tt, model, condition=condition),
xT, t, method="dopri5", atol=atol, rtol=rtol)
return samples[-1]
def _drift(self, x, t, model, condition=None):
"""velocity prediction 기반 drift (Flow Matching, Eq. 12)"""
return model(x, t, condition)
📌 Paper Info
🧠 Summary
✅ Key Points
Motivation
Method
Stage 1: Representation Autoencoder (RAE)
Stage 2: Training Diffusion Transformer on RAE latents
Width Scaling
Dimension-dependent Noise Schedule Shift
Noise-Augmented Decoding
DiTDH (Diffusion Transformer with DDT Head)
💭 Code
핵심코드 간단히 재구현 (Official 기반 chatgpt 재구현)