픽셀 확산 트랜스포머의 예측 방식과 고주파 결손을 함께 고친다

무엇인가

이 논문은 잠재 공간 디퓨전의 손실 압축 병목과 별도 오토인코더 문제를 피하려는 픽셀 공간 디퓨전, 그중에서도 JiT 계열의 단순한 픽셀 디퓨전 트랜스포머 설계를 다시 본다. JiT는 16×16처럼 공격적인 패치화로 토큰 수를 줄여 원시 픽셀에서 직접 생성하지만, 각 토큰이 고차원 픽셀 블록을 모델링해야 한다. 그래서 논문은 두 질문을 던진다. 고차원 패치에는 어떤 예측 파라미터화가 적합한가, 그리고 큰 패치를 쓰는 단순 트랜스포머가 미세한 이미지 디테일을 충실히 모델링할 수 있는가. 기존 JiT는 깨끗한 이미지 예측인 x0 예측을 택했고, AsymFlow는 PCA 상위 r개 방향에서만 속도 u를 예측하고 나머지에서 x0를 예측하는 이진 컷오프를 썼지만 r을 데이터셋마다 비싼 탐색으로 정해야 했다. k-Diff의 전역 스칼라 보간은 방향별 차이를 담지 못한다.

어떻게 동작하나

MFlow는 각 데이터 주성분 방향에 0과 1 사이의 연속 이득 m_i를 부여한다. 대각 연산자 M = U diag(m_1,...,m_D) U^T를 두고, 목표를 u_G = x0 - Mε로 정의한다. 성분별로는 u_G,i = x0,i - m_i ε_i이며, m_i=1이면 속도 예측, m_i=0이면 x0 예측이 된다. 따라서 MFlow는 x0 예측과 u 예측의 양 끝점은 물론 AsymFlow의 이진 마스크까지 포괄한다. 이득은 정규화 분산 기반 m_i = λ_i/(λ_i+c)로 초기화하고, 기본적으로 방향별 로짓 φ_i에 대해 m_i = sigmoid(φ_i)로 학습한다. 전방 보간 x_t = t x0 + (1-t)ε와 u_G 목표를 연립하면 각 고유방향에서 x0와 ε을 정확히 복원할 수 있고, 행렬식이 -(1-t+m_i t)라서 모든 m_i∈[0,1]에서 0이 아니므로 복원이 항상 가능하다. 네트워크는 u_G를 출력하지만 이를 전체 속도 u로 변환한 뒤 표준 속도 공간 플로우 매칭 손실로 학습한다. M은 전방 과정이나 수송 목표 u를 바꾸지 않고, 네트워크 회귀 목표의 좌표별 조건화만 바꾼다. 샘플링은 복원된 u로 표준 Euler/Heun을 돌린다.

무엇과 다른가

장난감 실험에서는 2차원 나선을 D차원으로 임베딩해 D=2, 8, 16, 512에서 비교했다. 전체 u 예측은 D가 커질수록 나빠졌지만 x0 예측과 r=d인 AsymFlow는 안정적이었고, 고정 이득 MFlow와 학습 이득 MFlow 모두 D=512에서도 나선을 복원했다. 이는 연속 이득이 고차원에서도 다룰 수 있고, 방향별 배분을 휴리스틱 없이 종단간 학습할 수 있음을 보여준다. 본 실험은 ImageNet-1K 256×256 클래스 조건부 생성에서 FID@50K, UCF-101 16프레임 128×128 비디오에서 FVD@10K, 약 9억 장 사전학습과 3,350만 장 SFT를 쓴 텍스트-이미지 생성에서 GenEval·DPG-Bench·PRISM-Bench로 평가했다. ImageNet은 16×16 패치로 256토큰에 32개 문맥 토큰을 더했고, 비디오는 4×8×8 패치로 1,024토큰, 텍스트-이미지는 16×16 패치와 256p·512p 혼합 해상도, LPIPS 가중치 0.1을 사용했다.

어떻게 쓰나

ImageNet B급에서 LPIPS 없이 기본 JiT-B 설정에서는 AsymFlow가 JiT와 MFlow보다 좋았다. 그러나 PixelGen을 따라 LPIPS 손실을 추가하자 전반적인 성능이 올랐고, 이 설정에서는 MFlow가 가장 좋았다. H급에서도 LPIPS를 붙여 비교했으며, 저자들은 AsymFlow의 근소한 우위가 ImageNet에 과도하게 튜닝된 결과일 수 있다고 본다. MFlow는 파라미터화 탐색 없이 기존 픽셀 디퓨전 트랜스포머들과 경쟁력 있는 수준을 보였다. 절제 실험에서 u_G 손실 대신 u 손실을 쓰는 것이 중요했고, u_G 손실은 FID를 3.66에서 4.88로 악화시켰다. M 초기화를 asym8에서 norm_var로 바꾸면 FID가 3.84에서 3.80으로, M 학습률을 10배로 높이면 3.80에서 3.74로 개선됐다. 또한 AsymFlow에서 PCA[9:16]가 PCA[1:8]보다 나은 경우도 있어, 상위 분산 방향만 고르는 이진 컷오프가 임의적임을 보여준다(FID 3.54 대 3.59). 비디오에서는 B급 UCF-101에서 MFlow가 WAN2.1 잠재 디퓨전, 픽셀 기반 JiT, AsymFlow를 앞섰고, L급에서도 JiT와 AsymFlow를 이겼으며 소량의 사후 Framewise-LPIPS 학습으로 최고 성능을 냈다. 텍스트-이미지는 동일한 1.1B MMDiT-L/16 백본에서 JiT, AsymFlow(r=8), MFlow가 비슷한 수준이었지만 MFlow가 모든 벤치마크에서 가장 좋았고, 13B 총 파라미터·2.9B 활성 파라미터의 MMDiT-XXL/16 MoE로 확장했을 때 비교 모델 중 최고 성능을 냈다.

전제와 한계

논문은 파라미터화와 별개로 JiT 계열 픽셀 디퓨전 트랜스포머가 실제 이미지 대비 고주파가 지속적으로 부족하다는 점을 발견한다. 방사 평균 파워 스펙트럼(RAPS)으로 보면 생성 이미지가 0.5–0.7 cycles/pixel 대역의 방사 주파수를 결여했는데, 이는 16–32배 패치화와 언패치화 사이에 패치 내부 국소 연산이 없는 JiT 백본 특유의 문제로 보인다. DiT의 VAE 재구성/생성, REPA-E, 픽셀 레벨 경로가 있는 PixelDiT는 같은 문제를 보이지 않았다. ConvRefiner는 동결한 디퓨전 백본 위에 가벼운 잔차 컨볼루션 블록을 사후 학습해 이 결손을 보정한다. 네트워크의 x 예측 위에 잔차를 더하고 지각·적대 손실을 적용하며, 다른 파라미터화에서는 예측을 먼저 x로 변환한 뒤 잔차를 적용한다. PatchGAN은 거의 깨끗한 이미지에만 GAN 손실을 준다. ImageNet에서는 600에폭 학습한 동결 베이스 위에 30에폭, 텍스트-이미지에서는 5k 스텝만 추가 학습해 0.5–0.7 대역을 메우면서 FID를 유지했다. 절제 실험에서는 블록 4개와 GAN 가중치 0.02를 선택했고, DiP 패치 디테일러는 고주파를 살리지만 JiT의 FID를 3.42에서 3.90, 4.46으로 악화시켰다.

학습된 M-커브 분석은 파라미터화 선호가 데이터와 용량에 따라 체계적으로 달라짐을 보여준다. ImageNet에서 LPIPS를 추가하면 선두 몇 방향을 제외한 대부분이 x0 예측 쪽으로 이동하면서 FID도 좋아졌다. 비디오에서는 유의미한 이득을 가진 방향 수가 더 줄어 도메인별 차이를 보였다. x0 예측을 쓰는 RAEv2 잠재에서는 이득이 거의 0으로 붕괴해 사실상 순수 x0 파라미터화로 돌아갔고, 소수 선두 방향에 u 예측을 강제해도 FID 변화는 미미했다. 반대로 텍스트-이미지에서 모델 용량을 키우면 더 많은 방향이 u 예측으로 이동했다. 종합하면 모델 용량은 u 예측을 얼마나 수용할 수 있는지를, 데이터 표현의 통계는 u 예측이 어디에 유용한지를 결정한다. 고정된 수의 u 예측 방향을 쓸 때도 배분은 데이터 의존적이어야 하며, 고분산 PCA 방향이 자연스러운 선택이 된다.

개발자 관점에서 이 논문은 픽셀 공간 디퓨전에서 예측 타깃을 하나로 고정하지 말고 데이터 주성분별로 학습하라는 설계 지침을 준다. 특히 큰 패치를 쓰는 단순 ViT 백본에서 x0 예측과 u 예측의 조건화 차이가 성능에 영향을 줄 수 있고, 모델 용량을 키울수록 u 예측 선호가 커진다면 MFlow식 방향별 이득이 수동 랭크 탐색을 줄일 수 있다. 다만 MFlow도 ImageNet B급 LPIPS 없는 설정에서는 AsymFlow에 근소하게 뒤졌고, 저자들은 AsymFlow가 ImageNet에 과도하게 튜닝됐을 가능성을 언급한다. M 이득은 sigmoid 폭주를 막기 위해 [0.05,0.95]로 클램프했고, u_G 손실이 아니라 u 손실로 학습해야 한다. ConvRefiner는 백본을 동결한 사후 적대 학습이 필요하며, 고주파 스펙트럼과 FID는 회복하지만 텍스트-이미지 벤치마크 점수를 끌어올리지는 않았다. 논문은 별도 한계 절을 두지 않지만, 고주파 결손이 JiT 백본 특유라는 점과 ConvRefiner가 소량의 추가 학습을 전제한다는 점은 실무 적용 전 확인해야 할 전제다.