PixelDiT2는 동결 비전 표현을 잡음제거 조건으로 쓰는 픽셀확산 트랜스포머다.

무엇인가

이 논문은 픽셀 공간 확산 모델이 잠재 공간 확산 모델과의 이미지 품질 격차를 좁혔지만 여전히 수렴이 느리고 최종 품질이 뒤처진다는 문제를 다룬다. 저자들은 핵심 원인을 명시적 표현 사전(representation prior)의 부재로 본다. 잠재 확산은 사전학습된 오토인코더가 만든 압축되고 구조화된 잠재 공간에서 잡음을 제거하지만, 픽셀 확산은 원시 RGB 공간에서 잡음 제거에 유용한 표현 학습과 픽셀 생성을 동시에 해야 한다는 것이다. PixelDiT2는 오토인코더나 잠재 재구성 병목 없이 표현 학습과 픽셀 생성을 분리하려는 엔드투엔드 픽셀 공간 확산 모델이다.

어떻게 동작하나

방법의 뼈대는 픽셀 공간 rectified flow다. 깨끗한 이미지 x와 잡음 ε에 대해 x_t = t x + (1-t) ε, v = x - ε로 두고, 네트워크가 (x_t, t, 클래스 라벨 y)를 받아 깨끗한 이미지 추정 x_hat을 출력한다. 예측 속도는 v_hat = (x_hat - x_t) / max(1-t, τ)로 계산하며 τ=0.05를 사용하고, 손실은 L_diff = E ||v_hat - v||^2다. 백본은 PixelDiT의 패치 수준과 유사한 단일 경로 패치 수준 확산 트랜스포머이며 x 예측과 속도 공간 손실을 쓴다.

무엇과 다른가

핵심 제안은 표현 그라운딩이다. 동결된 DINOv3 백본 E가 현재 잡음 샘플 x_t를 인코딩해 패치별 토큰 g_bar_t = E(x_t) ∈ R^(B×L×d_E)를 만든다. 여기서 L=(H/p)(W/p)이고 p=16은 픽셀 확산 트랜스포머의 패치 크기와 맞춘다. 이 토큰은 타임스텝 조건부 투영 P_g를 거쳐 확산 트랜스포머의 은닉 차원 D로 사상된다. 즉 g_t = P_g(g_bar_t, t) ∈ R^(B×L×D)다. P_g는 폭 D의 단일 DiT 스타일 트랜스포머 블록으로, 선형 입력 투영, 셀프 어텐션, 타임스텝 임베딩 e_t의 AdaLN으로 변조되는 정규화와 MLP, 최종 선형 층으로 구성된다. 이 그라운딩 토큰은 픽셀 확산 트랜스포머 블록에 공간 AdaLN 조건으로 주입되고 추론 시에도 활성화된다. 인코더는 완전히 동결되며, 잡음 수준 적응은 P_g가 담당한다.

어떻게 쓰나

저자들은 REPA 같은 훈련 시점 정렬과의 차이를 강조한다. REPA는 사전학습 특징을 보조 훈련 목표로 쓰고 샘플링 중에는 교사 인코더를 평가하지 않지만, PixelDiT2는 사전학습 표현을 잡음 제거 과정의 일부로 유지하며 패치별 조건으로 공간 구조를 보존한다. 또한 Latent Forcing이 잠재 스트림을 픽셀과 함께 생성해 표현 안내를 제공하는 것과 달리, PixelDiT2는 각 단계의 현재 잡음 이미지에서 직접 조건 특징을 얻고 확산 과정은 전적으로 픽셀 공간에 남는다. CFG를 위해 클래스 토큰을 확률 p_drop=0.1로 널 토큰으로 바꾸지만, 그라운딩 토큰이 클래스 판별 정보를 담을 수 있어 클래스 드롭만으로는 무조건 분기가 완전히 무조건적이지 않다. 그래서 E와 P_g를 우회하고 타임스텝과 널 클래스 임베딩의 브로드캐스트로 대체하는 널 그라운딩 상태를 훈련에 노출하고, 추론의 CFG 무조건 분기도 같은 널 그라운딩 토큰을 쓴다.

전제와 한계

실험은 클래스 조건부 ImageNet 256×256과 512×512에서 수행됐다. 아키텍처는 패치 크기 16의 B/L/H 세 규모이고, 훈련은 AdamW, 5에폭 선형 워밍업 후 상수 학습률 2×10^-4, 가중치 감쇠 없음, 전역 배치 1024, EMA 0.9999, 그래디언트 클립 1.0, bf16을 사용한다. 샘플링은 Heun ODE 솔버 50스텝과 구간 제한 CFG를 쓰고 FID-50K와 Inception Score를 50,000개 샘플로 보고한다. 256×256에서 초기부터 독립적 클래스 및 그라운딩 드롭아웃을 적용한 훈련은 600에폭에 FID 1.46과 IS 301.6을 달성했다. 지연된 그라운딩 드롭아웃을 쓰면 더 이른 480에폭에 FID 1.48에 도달했고, 이는 PixelDiT-XL의 320에폭 FID 1.61과 JiT-G/16의 600에폭 FID 1.82(약 2B 파라미터)를 앞선다. 512×512에서는 PixelDiT2-H/16이 패치 16으로 32×32 토큰 그리드를 만들고, 그라운딩 경로는 동결 DINOv3-B/16, REPA 목표는 DINOv2-B/14를 쓴다. 이 구성은 680에폭에 FID 1.48과 IS 295.7을 기록했고 표에 있는 픽셀 공간 방법 중 가장 낮은 FID였다. 논문 기여에는 PixelDiT2-H/16이 256에서 600에폭 FID 1.46, 512에서 680에폭 FID 1.48로 PixelDiT의 850에폭 FID 1.81을 능가한다고 정리한다.

절제 실험은 그라운딩의 기여와 설계 선택을 분리한다. REPA와 표현 그라운딩을 함께 쓰면 200에폭에는 REPA 단독과 거의 같은 FID 1.89와 1.90이었고, 그라운딩 단독은 2.28로 그라운딩 없는 백본 2.53보다 좋았다. 320에폭 이후부터는 REPA와 그라운딩 결합이 평가한 모든 에폭에서 가장 낮은 FID를 보였고 480에폭까지 1.70에서 1.59로 개선됐다. 같은 구간에서 REPA 단독은 1.7 근처에서 정체했고 그라운딩 단독은 2.08에서 1.83으로 계속 개선됐다. 주입 방식은 공간 AdaLN이 FID 1.462로 가장 좋았고 입력 덧셈 1.521, 레이어 0 토큰 연결 1.621, 레이어 8 토큰 연결 1.603이었다. 동결 인코더 선택에서는 DINOv3-S/B/L, DINOv2-S/B, MAE-B가 FID 1.462–1.624 범위였고 DINOv3-S/16이 가장 좋았다. DINOv3 규모를 S에서 B나 L로 키우면 600에폭 공유 평가에서 FID가 각각 1.547, 1.552로 나빠졌다. 인코더 적응 실험에서는 동결 인코더가 FID 1.462인 반면 네 가지 적응 변형은 1.581–1.625에 머물렀다. P_g 형태를 본 파라미터 맞춤 B/16 진단에서는 DiT 스타일 P_g가 FID 5.29로 선형 투영 6.79보다 좋았고, 총 훈련 파라미터는 143M이었다. P_g 용량 실험에서는 DINOv3-L을 고정하고 폭을 두 배로 늘리면 200에폭 FID가 1.82에서 1.78, 320에폭이 1.73에서 1.64로 개선됐고, 깊이를 4블록으로 늘리면 200에폭은 그대로였고 320에폭에 1.69에 도달했다. 그러나 두 변형 모두 320에폭 DINOv3-S/16 기준 1.59에는 못 미쳤다.

그라운딩 드롭아웃 연구는 CFG와 후반 안정성에 관한 실용적 관찰을 준다. 초기부터 독립적 클래스 및 그라운딩 드롭아웃을 쓰면 FID가 320에폭 1.59에서 600에폭 1.46으로 단조 감소했다. 클래스 드롭아웃만 유지하고 그라운딩을 남기면 500에폭 FID 1.54 이후 600에폭 1.555로 올라가 경미한 후반 과적합 양상을 보였다. 지연된 그라운딩 드롭아웃은 160에폭까지 클래스 드롭아웃만 쓰고 그라운딩을 유지한 뒤 이후 독립적 클래스 및 그라운딩 드롭아웃을 켜는 방식으로, 200, 320, 480에폭에 FID 1.66, 1.50, 1.48을 기록했다. 같은 에폭에서 그라운딩을 유지한 실행은 1.75, 1.59, 1.54였다. 내부 분석에서는 동결 인코더가 t가 0에 가까워질수록 급격히 드리프트하지만 DINOv3-S/16 뒤에 훈련된 P_g를 붙이면 깨끗한 입력 출력과의 코사인 유사도가 모든 t에서 0.77 이상으로 유지됐다. 200에폭 PixelDiT2-B/16의 깨끗한 입력 프로브는 널 클래스 토큰을 쓰고, 그라운딩+REPA가 모든 블록에서 REPA 단독 기준을 앞섰으며 둘 다 인컨텍스트 주입 블록 i=4에서 최고였다. 가장 이른 블록 i=0에서 격차는 약 +38 퍼센트 포인트였고 더 깊은 절반에서는 약 +8에서 +10 포인트의 안정적 차이를 보였다.

개발자 관점에서 이 논문은 픽셀 공간 생성을 오토인코더 없이 구성하려는 경우, 사전학습 비전 인코더를 단순한 훈련 목표가 아니라 추론 시에도 동작하는 공간 조건 신호로 쓰는 설계를 제안한다. 다만 인코더를 각 잡음 제거 단계에서 실행해야 하므로 추가 추론 비용이 발생하고, 인코더와 확산 트랜스포머의 패치 그리드를 16×16으로 맞춰야 하며, 잡음 입력에 대한 인코더 드리프트를 흡수할 타임스텝 조건부 투영과 CFG용 널 그라운딩 분기를 준비해야 한다. 저자들이 밝힌 한계는 PixelDiT2가 최고 잠재 확산 베이스라인과의 완전한 격차를 닫지 못했고 동결 인코더 실행이라는 추가 추론 비용을 도입한다는 점이다. 또한 실험은 클래스 조건부 ImageNet 256과 512에 한정되며, 향후 과제로 확장된 텍스트-투-이미지 생성에 적용할 계획을 제시한다. 256픽셀에서는 더 작은 그라운딩 인코더가 항상 좋다는 해상도 독립적 규칙을 지지하지 않으며, 512픽셀에서는 다른 최적화 양상이 나타난다고 논문은 밝힌다.