픽셀 확산 트랜스포머가 스스로 만든 표현으로 생성 품질을 높인다

무엇인가

이 논문은 픽셀 확산 트랜스포머(pDiT)에서 "생성에 쓸모 있는 표현이란 무엇인가"를 다룬다. 컴퓨터 비전의 표현 학습은 분류·검출·분할 같은 이해 과제 중심으로 발전해 왔지만, 생성은 의미 구조와 공간 구조를 함께 조직하면서 최종적으로 모든 픽셀을 복원해야 하므로 이해에 좋은 표현이 생성에도 좋으리라는 보장이 없다. 저자들은 외부 오토인코더로 미리 정의된 표현 공간을 쓰지 않고 원시 픽셀에서 표현 학습과 이미지 합성이 end-to-end로 일어나는 pDiT, 그중에서도 평범한 Transformer만으로 경쟁력 있는 픽셀 생성을 보여준 JiT를 분석 무대로 삼는다.

어떻게 동작하나

분석에서 두 가지 결론을 얻는다. 첫째, 표현 품질이 좋아진다고 생성 효용이 커지지 않는다. JiT-B/16의 중간층을 DINOv2에 정렬(REPA)하면 정렬한 층의 RMSC·LDS·선형 프로빙 정확도는 크게 오르지만, 이후 층으로 갈수록 그 이득이 빠르게 사라진다. 생성에서도 같은 불일치가 나타나 REPA-4는 초반 학습을 가속하지만 약 160 에폭 이후 베이스라인이 역전해 최종 FID가 4.56 대 4.80으로 더 좋았고, 정렬을 더 깊은 층으로 옮기면 REPA-6 5.38, REPA-8 8.33으로 오히려 나빠졌다. 픽셀 단위 감독은 국소 연산만으로 대부분 충족되므로 모델이 고수준 표현을 형성·활용할 유인이 없다는 설명이다. 둘째, 전용 픽셀 디코더로 구조를 분리하는 것은 표현을 개선하기보다 의미 모델링과 공간 모델링의 분담을 재배치할 뿐이다. PixelDiT는 파라미터 55%, FLOPs 67%를 더 쓰고도 FID 4.26으로 베이스라인 4.17과 비슷했고, DiP는 6.09, DeCo는 9.17이었다. 디코더의 공간 모델링 능력이 인코더가 무엇을 보존해야 하는지를 바꾸기 때문에, 디코더 설계는 표현을 소비할 뿐 아니라 표현의 내용을 규정한다.

무엇과 다른가

제안 방법인 Representation Guidance(RG)는 이 두 요구를 저차원 병목과 유도 목표로 나눠 해결한다. 구조는 DDT를 따라 표현 학습용 인코더와 픽셀 합성용 디코더로 나뉘지만 모든 블록은 JiT 아키텍처를 그대로 쓴다. 디코더는 자체 패치 임베딩과 소수의 Transformer 블록을 갖고, 인코더에는 표현 r을 통해서만 접근하며 이를 토큰별 adaLN 조건으로 받는다. 인코더 출력 h를 선형 병목으로 u = h·W_down (d≪D)로 압축했다가 r = u·W_up으로 되돌리는데, 이 u 공간을 암묵적 잠재 공간이라 부른다. VAE 잠재처럼 압축적이고 공간 구조를 갖지만, 매 스텝 z_t에서 계산되고 생성 목적으로 end-to-end 최적화되며 확산 자체는 픽셀 공간에서 일어난다.

어떻게 쓰나

유도 목표는 모델 자신의 표현이 만든 속도 차이를 학습 타깃에 반영한다. 학습 중 확률 0.1로 r을 0으로 만들어 표현 없는 기준 예측 v_θ(z_t,t,∅)을 얻고, 두 예측의 차이 Δv_rep = v_θ(z_t,t,r) − v_θ(z_t,t,∅)를 구한다. 그리고 v_RG = v + (s−1)·sg(Δv_rep)를 새 타깃으로 삼아 L_RG = ‖v_θ(z_t,t,r) − v_RG‖²를 최소화한다(sg는 stop-gradient). s=1이면 기존 flow matching 손실로 환원된다. 수렴 시점에 RG는 r이 예측 속도에 기여하는 몫을 1/(2−s)배로 증폭시켜, 학습 분포를 r과 일치하는 샘플 쪽으로 날카롭게 만든다. 손실의 그래디언트는 r을 거쳐 인코더로 흘러 표현 자체가 유도 목표 아래 최적화된다. Δv_rep은 안정적인 타깃을 위해 EMA 네트워크로 계산하고, RG는 t∈(0.3, 1] 구간에서 20 에폭 워밍업 후 적용하며 s=1.4를 쓴다. 중간 인코더 층에는 REPA를 함께 걸어 L = L_RG + λL_REPA로 학습한다. 추론 시 추가 비용은 없다.

전제와 한계

실험은 ImageNet 256×256과 512×512의 클래스 조건부 생성, 그리고 텍스트-투-이미지 생성에서 이뤄졌다. RG-XL/16은 725M 파라미터로 ImageNet 256×256에서 FID 1.55를 기록해 픽셀 공간 방법 중 가장 낮았고, PixelDiT-XL(1.61), DiP-XL/16(1.79), DeCo-XL/16(1.69)을 앞선다. 459M에 불과한 RG-L/16은 FID 1.79로 2B 파라미터의 JiT-G/16(1.82)과 953M의 PixelREPA-H/16(1.81)을 능가한다. 512×512에서는 FID 1.74를 얻었다. 텍스트-투-이미지에서는 RG-T2I-L/16이 GenEval 0.89, DPG-Bench 87.1을 기록해 SD3(8B), FLUX.1-dev(12B) 같은 대규모 잠재 공간 모델과 더 큰 텍스트 인코더를 쓰는 픽셀 공간 방법들을 앞섰다. 참고로 RG-T2I-L/16의 텍스트 인코더는 341M(FLAN-T5-Large)이고, PixelDiT-T2I와 DeCo-XXL/16은 각각 2.6B, 1.7B 텍스트 인코더를 쓴다.

절제 실험은 설계 선택의 근거를 보여준다. 병목 폭 d는 200 에폭에서는 3.18~3.36으로 비슷했지만 600 에폭에서는 d=128이 FID 2.64로 d=512(2.89)보다 좋아 기본값이 됐다. REPA 정렬 층은 3층·6층이 3.28·3.33인 반면 마지막 인코더 층(8층)에 정렬하면 8.97로 급등하는데, DINOv2 같은 사전학습 표현이 픽셀 생성에 맞지 않는다는 근거다. 12블록 예산에서 인코더/디코더를 4/8로 나누면 3.60, 8/4면 3.28, 10/2면 4.19로, 인코더가 깊어야 표현이 좋아지되 고주파 합성을 위한 디코더 깊이도 필요하다. 구성 요소별로는 REPA 단독 4.17, 병목 추가 3.79, 유도 추가 3.28, 600 에폭에서 전체 2.64였고, REPA나 병목 중 하나라도 빠지면 유도가 FID를 7.85~8.99로 악화시켰다. RG가 추가하는 비용은 파라미터 0.4M과 FLOPs 13%뿐이다. RG 스케일은 s=1.4, 구간 하한은 t_min=0.3이 최적이며 모든 노이즈 레벨에 적용하면 5.28로 나빠진다. 정렬 대상은 DINOv2-B 3.46, 레지스터를 붙인 DINOv2-B 3.28, DINOv3-B/16 3.30으로 큰 차이가 없다.

개발자 관점에서 이 논문의 실용적 메시지는 두 가지다. 첫째, 사전학습 표현을 중간층에 정렬하는 것만으로는 생성 품질이 보장되지 않으며, 정렬 층이 뒤로 갈수록 오히려 해가 될 수 있으므로 정렬 위치와 스케줄을 반드시 검증해야 한다. 둘째, 전용 픽셀 디코더를 도입하는 대신 표현 인터페이스에 저차원 병목을 두고 표현이 실제로 쓰이도록 학습 목표를 바꾸는 편이 파라미터와 연산 면에서 유리할 수 있다. 특히 RG는 추론 시 추가 비용이 없고 학습 시 EMA 네트워크의 그래디언트 없는 순전파 두 번만 더하므로, 기존 픽셀 확산 학습 파이프라인에 비교적 저렴하게 얹을 수 있는 선택지다. 다만 병목 폭과 RG 스케일, 적용 구간이 학습 길이에 따라 최적값이 달라진다는 점(200 에폭과 600 에폭에서 최적 d가 다름)은 실무에서 직접 스윕해야 할 항목이다.

저자들이 밝힌 한계는 명확하다. 실험이 이미지 생성에 국한되어 있고, RG를 더 큰 모델 스케일이나 비디오 생성, 3D·오디오 같은 다른 모달리티로 확장하는 일은 향후 과제로 남겨 두었다.