[CVPR 2026] Back to Basics: Let Denoising Generative Models Denoise

안녕하세요. 이번 x-review는 최근 이미지 생성 분야에서 큰 주목을 받고 있는 논문인 “Back to Basics: Let Denoising Generative Models Denoise”입니다. MIT의 Tianhong Li와 Kaiming He가 저술한 논문으로, 방법론적으로는 JiT (Just image Transformers)라고 논문에선 소개합니다. 일단 제목이 진짜 자극적인데요. 기본으로 돌아가자. 또, 그 뒤에 따라오는 denoising하는 생성모델이 직접 denoising하게 하라. 다 읽고나니 제목 진짜 fancy하게 잘 지은 것 같습니다.. DiT같은 모델이 diffusion 설계로 latent embedding/tokenizer 등등을 사용하게 됐었는데, 이 논문은 최대한 단순하게 ViT에 픽셀패치 인풋으로 받고, 픽셀패치 아웃풋으로 내뱉는 아주 기초적인 방식, 즉 pixel space에서 모델 자체가 이미지 생성하게 하는 거는 왜 그 동안 간과했느냐. 이런 latent-free, tokenizer-free한 쌩 ViT기반 생성 방식도 유효하니 이 쪽 연구도 재고하자 가 핵심입니다. 읽다보니 기존 DiT 기반 다 쒯이야! 이런 것까진 아니었습니다.

요즘 제가 JEPA 기반의 world model 계열을 계속해서 봐오면서 latent representation의 구조적 특성이 학습 안정성에 얼마나 직접적인 영향을 미치는지를 많이 느끼고 있었는데, 본 논문도 비슷한 느낌으로 diffusion 기반 이미지 생성 모델에서 가장 기본이 되는 설계 철학이 뭐여야 하냐를 다시 들춰 “네트워크가 무엇을 예측해야 하는가?”라는 근본 질문으로 manifold hypothesis 관점에서 다시 꺼내 든 논문이라고 보시면 될 것 같습니다. 결론은 ‘복잡한 파이프라인을 다 걷어내고 plain ViT + pixel space + x-prediction(노이즈나 flow가 아닌 data 자체(x)를 예측..!)만으로 충분하다’라는 점에서, 최근 점점 복잡해지는 생성 모델 연구 흐름에서 간과하고 있었던 모델 설계와 학습 시 결국 high-dim context를 가진 데이터가 low-dim latent manifolds에 underlying한다는 건 너무 간과하고 있던 문제다. 이걸 어떤 관점에서 기존 접근들의 근본적 문제를 지적하고, 어떻게 해결했는지를 위주로 한번 살펴보겠습니다.

1. Introduction

[Diffusion Model의 prediction target이 왜 ε/v-prediction이 표준이 되었는지..?]

Diffusion model의 원래 목표는 노이즈가 낀 이미지로부터 깨끗한 이미지를 복원하는 denoising이었습니다. 초기 diffusion model은 reversed stochastic process에서 normal distribution의 parameter(mean과 standard deviation)를 네트워크가 직접 예측하는 방식이었는데요. 이후 DDPM(Denoising Diffusion Probabilistic Models)이 등장하면서, 노이즈 자체(ε)를 예측하는 ε-prediction 방식이 성능이 좋아져 diffusion model의 사실상 표준으로 자리 잡았습니다. 이후에는 v-prediction(velocity prediction)도 도입되었고, Flow Matching 모델들 역시 v-prediction과 유사한 구조를 가지며 발전해 왔습니다.

이 과정에서 여러 연구들이 x-prediction, ε-prediction, v-prediction이 적절한 loss 가중치 조정만 있으면 서로 수학적으로 연결된다는 점을 보여줬습니다. EDM(Elucidated Diffusion Models)도 문제를 denoiser 함수 중심으로 재정의했지만, 여전히 네트워크 출력이 data + noise가 섞인 형태여서 완전히 깨끗한 이미지를 직접 예측하지는 않았습니다. 결국 “네트워크가 무엇을 직접 예측해야 하는가?”라는 근본적인 문제는 거의 주목받지 못한 채, 모델이 어떤 타깃이든 예측할 수 있다고 암묵적으로 가정되어 온 것입니다. 초기 DDPM에서 ε-prediction이 더 우수하다고 관찰된 이래로, 일종의 legacy처럼 ε이나 v 예측을 유지해 온 측면이 있습니다.

본 논문은 이 가정에 정면으로 의문을 제기하는 것으로 인트로를 시작했습니다.

[Manifold Hypothesis : x-prediction과 ε-prediction이 근본적으로 다른 이유]

논문의 핵심 논의 자체는 manifold hypothesis에 기반합니다. 머신러닝에서는 오래전부터 고차원 데이터가 실제로는 훨씬 낮은 차원의 구조(manifold) 위에 놓여 있다는 가정을 사용해 왔는데요. 이 가정에 따르면, 자연 이미지 x는 전체 공간 R^D가 아니라 그 안의 저차원 다양체 M 위에 놓입니다. 수학적으로 표현하면:

x ∈ M ⊂ R^D, dim(M) = d ≪ D

여기서 d는 실제 의미 있는 변화 요인의 수에 해당합니다. 이미지는 수만 개의 픽셀로 이루어져 있지만, 실제 변화 요인(pose, illumination, texture 등)은 그보다 훨씬 적으므로, 데이터가 고차원 전체에 퍼지지 않고 저차원 곡면처럼 모여 있다는 뜻입니다.

반면 노이즈 ε은 특정 구조를 따르지 않고 모든 차원에 독립적으로 퍼져 있는 값입니다:

ε ~ N(0, I_D)

각 차원이 독립적이기 때문에 자연 데이터가 가지는 저차원 구조를 전혀 따르지 않습니다. 노이즈는 manifold 위가 아니라 고차원 전체 공간을 균일하게 채웁니다. v = x – ε 역시 노이즈가 섞여 있으므로 마찬가지로 off-manifold입니다.

이 때문에 모델이 x를 직접 예측하는 문제는 저차원 구조만 복원하면 되는 비교적 쉬운 작업인 반면, ε이나 v를 예측하는 문제는 고차원 공간 전체를 다뤄야 하므로 필요한 모델 capacity가 본질적으로 훨씬 큽니다. 즉, x-prediction과 ε/v-prediction은 수학적으로 연결 가능하더라도, 네트워크가 실제로 수행해야 하는 작업의 난이도는 근본적으로 다릅니다.

[기존 접근법들은 Latent space로 본질적인 문제를 좀 숨겼다?]

기존 diffusion model들은 고차원 pixel space에서 noise를 예측하는 어려움을 회피하기 위해 VAE 기반 latent space로 문제를 옮겨왔습니다. LDM(Latent Diffusion Model), Stable Diffusion 등이 대표적인데요. 본 논문은 이 방식이 문제를 해결한 것이 아니라 단지 숨긴 것에 불과하며, 모델이 self-contained하지 못하게 만든다고 지적합니다. VAE tokenizer에 의존하게 되면 tokenizer의 품질이 곧 성능의 상한선이 되고, 모델 자체가 독립적으로 작동할 수 없게 됩니다.

최근에는 pixel space에서 직접 diffusion을 수행하려는 시도들이 있지만, 이들 역시 복잡한 설계를 필요로 했습니다:

• SiD2, PixelFlow: small patch에서 시작하는 hierarchical 구조를 사용하지만, FLOPs가 매우 커지며 Transformer 고유의 generality와 simplicity를 잃는 문제가 있다고 합니다.
• PixNerd: Transformer output + noisy input + spatial coordinates를 결합하는 NeRF-style head를 사용하며, 추가로 representation alignment까지 필요하다고 합니다.

이 모델들은 대부분 L~XL급의 큰 ViT를 필요로 하며, 고차원 pixel space를 다루려면 큰 hidden size가 필수적이라는 것이 기존의 통념이었습니다. 그래서 본 논문에서는 이 모든 복잡성이 ε/v-prediction이라는 잘못된 prediction target 선택에서 비롯된 것이며, x-prediction으로 바꾸면 plain ViT만으로 충분하다는 것을 핵심 주장으로 삼게 됩니다.

2. Prediction Space와 Loss Space 분석

[2.1 Background: Diffusion과 Flows]

본 논문은 이해가 쉬운 flow-based(v-space) 관점에서 formulation을 시작합니다. 데이터 x ~ p_data(x)와 노이즈 ε ~ p_noise(ε)가 있을 때, 시점 t ∈ [0,1]의 noisy sample은 linear interpolation으로 구성됩니다:

본 연구는 선형 노이즈 스케줄 a_t = t, b_t = 1 – t를 사용합니다. 따라서:

여기서 t=1일 때 z_t가 data distribution에 수렴하고, t=0일 때 noise distribution에 수렴합니다. Flow velocity v는 z_t의 시간 미분으로 정의됩니다:

Flow 기반 모델은 다음 loss를 최소화합니다:

sampling은 다음 ODE를 풀어 수행됩니다:

초기 조건은 z_0 ~ p_noise이고, t=1까지 적분하여 데이터를 복원합니다. 본 논문에서는 50-step Heun solver를 사용합니다.

[2.2 Prediction Space × Loss Space: 9가지 조합]

논문은 prediction space(네트워크가 직접 출력하는 대상, 하이라이트된 부분)와 loss space(학습 loss를 계산하는 공간)를 명확히 분리합니다. 미지수는 x, ε, v 3개이고, network로부터의 output은 1개로 주어져서, 3가지 미지수를 결정하려면 2가지의 추가적인 constraints가 필요한데요. 만약 x_θ = netθ 로 정해놓는다면,

이 말인 즉슨, ε, v 는 x, z 로부터 계산이 된다는 의미입니다. 따라서 네트워크가 x, ε, v 중 어떤 것을 직접 출력할 것인지 정해주는 것에 따라 다른 두가지 미지수가 추론될 수 있다는 것이고, 이로써 세 가지 prediction space가 선정의되고(Table 1), 각각에 대해 x-loss, ε-loss, v-loss 세 가지 loss space를 조합하면 방정식을 해결하는 경우의 수가 총 9가지 경우가 됩니다.

이 방정식 얘기의 핵심은 한 prediction space에서 다른 공간으로 reparameterization하면 loss가 단순히 reweighting된 형태로 변환된다는 점입니다.

예를 들어, x-prediction + v-loss 조합에서 네트워크 출력은 vθ = (xθ – z_t) / (1 – t), 타깃은 v = (x – z_t) / (1 – t)이 되므로:

L = E‖vθ – v‖² = E · (1/(1-t)²) · ‖xθ – x‖²

즉 x-loss에 1/(1-t)² 가중치를 붙인 형태가 됩니다. 이와 같은 변환은 어떤 prediction × loss 조합에도 적용 가능하며, 9개 조합은 모두 유효한 formulation이지만 서로 수학적으로 동등한 경우는 없다고 합니다. 어떤 조합이든 inference 시에는 항상 네트워크 출력을 v-space로 변환한 뒤 ODE sampling으로 생성합니다.

[2.3 Toy Experiment: 왜 pixel space에서 ε-prediction이 실패하는가??]

이 논문에서 가장 설득력 있는 부분 중 하나는 Section 3.3의 toy experiment라고 생각합니다. 실험 설정은 다음과 같습니다.

• 기저 차원 d=2인 데이터를 생성합니다.
• 이를 임의의 column-orthogonal projection matrix P ∈ R^(D×d)로 D차원 공간에 매핑하여 x = P·x̂을 만듭니다.
• 모델은 이 projection matrix를 모르는 상태에서 D차원 생성 문제를 학습합니다.
• 모델은 256 hidden units의 5-layer ReLU MLP로 아주 단순하게 설정했습니다.
• D ∈ {2, 8, 16, 512}로 변화시키며, x/ε/v-prediction 세 가지를 모두 v-loss로 학습합니다.

결과는 Fig.2에 시각화되어 있는데, 매우 직관적입니다

• x-prediction: D가 아무리 커져도(D=512에서도) 원래의 2D manifold를 정확히 복원합니다. 256차원 MLP가 512차원 출력을 다루더라도, 실제 예측 대상이 2차원 구조를 가지므로 문제가 되지 않습니다.
• ε-prediction: D=16부터 어려워지기 시작하고, D=512에서는 데이터가 blob 형태로 완전히 붕괴합니다. 차원이 증가할수록 예측 대상인 노이즈가 공간 전체에 퍼지면서, 모델이 예측한 값이 manifold를 벗어나 허공을 헤매게 됩니다.
• v-prediction: ε과 유사하게 고차원에서 거의 collapse합니다. v = x – ε이므로 여전히 고차원 노이즈 성분을 포함하고 있기 때문입니다.

이 실험이 보여주는 핵심은, x-prediction의 경우 모델 출력이 수학적으로 부분 공간으로의 투영(projection onto subspace) 효과를 가진다는 점입니다. 모델이 예측을 수행할 때마다 결과값이 강제로 데이터 manifold 위로 끌어당겨지므로, 고차원에서도 안정적으로 학습됩니다. 반면 ε/v-prediction은 off-manifold인 고차원 전체 공간에서 예측을 수행해야 하므로, space가 클수록 예측해야 할 노이즈의 분산이 계속 커지고 학습 난이도가 급격히 상승합니다.

저는 이 toy experiment가 pixel diffusion이 실패한 이유가 compute 부족이 아니라 prediction target 선택의 문제였다는 논문의 핵심 주장을 가장 직관적으로 증명하는 실험이라고 생각합니다. 실제로 ImageNet에서도 유사한 현상이 관찰됩니다.

3. JiT Architecture

그리하여 저자들은 위에서 얻은 prediction space 와 loss space에 대한 본질적인 고찰에 힘입어 JiT(Just image Transformers) 아키텍쳐를 정의합니다. 이름 그대로 저스트 이미지 트랜스포머입니다. 지금까지의 diffusion 시스템이 가진 복잡한 구성 요소들을 과감히 제거하고, 원본 이미지의 pixel patch를 Transformer로 직접 처리하는 방식을 사용합니다.

• No Tokenizer: 이미지를 별도의 코드로 변환하지 않고,
• No Latent Model: VAE 같은 사전 학습 모델을 사용하지 않고,
• No Upsampler: 저해상도에서 고해상도로 키우는 cascaded generation 과정 없이 한 번에 생성하고
• No Extra Loss: adversarial, perceptual, self-supervised loss 없이 순수한 diffusion loss만 사용하고
• No Pre-training: 별도의 사전 학습 과정 없습니다.

[3.1 Patchify → ViT → Patch Reconstruction]

JiT는 ViT의 핵심 아이디어인 Transformer on Patches (ToP) 의 철학을 그대로 따릅니다. 이미지(H×W×3)를 p×p 단위의 patch로 나누고, 이를 시퀀스로 간주해 Transformer로 처리합니다. 각 패치는 p×p×3 차원의 벡터이며, linear embedding projection을 거치고 positional embedding을 더한 뒤 Transformer block을 통과시킵니다. 마지막에는 다시 p×p×3 patch로 복원하는 linear layer를 사용합니다.

512×512 이미지 기준으로 patch size에 따른 토큰 구성은 다음과 같습니다:

Patch Size	Patch Dim	Token 개수	특징
16×16	768	1024	토큰 수가 많아 연산 부담이 큼. fine-grained 함.
32×32	3072	256	연산 – 패치 차원 간 균형
64×64	12288	64	토큰 수 적어 빠르지만, 토큰 차원이 매우 큼

JiT의 기본 설정은 p=32로, 각 patch는 3072차원(=32×32×3)인 꽤 큰 벡터입니다. 기존 Diffusion Transformer(DiT 등)는 VAE latent(4~8 channels)이나 CNN feature map을 토큰으로 사용하는 반면, JiT는 아예 원시 pixel 패치를 토큰으로 사용합니다. 이러면서 더 이상 tokenizer라는 게 필요가 없어지고, 이미지 그 자체가 모델의 입력 토큰이 되는 구조를 갖게 됩니다. 이게 가능한 이유가 바로 x-prediction 기반의 안정성 덕분이라는 것이 논문의 주장입니다. 모든 x, ε, v, z_t는 동일한 이미지 차원(H×W×3)을 공유합니다.

[3.2 Bottleneck Embedding: 왜 차원 축소가 오히려 도움이 되는가]

앞서 말했던 p=32일 때 patch dim이 3072인데, 이를 그대로 Transformer에 넣으면 메모리와 compute 비용이 매우 큽니다. 그래서 JiT는 patch embedding 단계에서 bottleneck 구조를 사용합니다. Linear patch embedding을 두 개의 linear layer로 구성된 low-rank bottleneck 구조로 바꾸는 것입니다.

3072 (raw patch) → d' (bottleneck) → 768 (transformer hidden dim)

첫 번째 layer에서 차원을 d’로 줄인 뒤, 두 번째 layer에서 Transformer의 hidden size로 다시 확장합니다.

위의 Figure 4의 결과가 Bottleneck layer에 대한 실험인데요. 꽤 인상적이게도 JiT-B/16(patch dim 768)에서 bottleneck dimension d’를 극단적으로 줄여도 성능이 거의 떨어지지 않습니다. d’가 32~512 범위일 때는 오히려 이밎 생성 성능인 FID가 최대 약 1.3 정도 개선됩니다. d’=32면 원래 patch dim 768의 약 4%에 불과한데도 성능이 더 좋아지는 것입니다.

이 결과는 처음의 그 manifold hypothesis와 직접적으로 연결됩니다. Clean image x는 원래 저차원 manifold 위에 있기 때문에 raw pixel 정보 전체를 보존할 필요가 없고, manifold 구조만 잘 추출하면 ViT가 안정적으로 복원할 수 있다는 manifold hypothesis의 실험적 증거라고도 볼 수 있습니다. Representation learning 관점에서 보면 bottleneck은 저차원 비선형 representation 학습을 유도하기 위해 도입하는 것이고, 이는 고전적 manifold learning의 bottleneck autoencoder와 맥이 닿는 현상입니다. 이 실험은 pixel diffusion은 patch dimension 때문에 불가능하다는 기존 통념이 틀렸음을 보여줍니다.

[3.3 Transformer Backbone: “Just Advanced” Transformers]

Patch embedding 이후에는 거의 그대로의 ViT를 사용합니다. 다만 최근 language model에서 검증된 모델 구조적인 개선 기법을 차용하게 됩니다. 아래 것들이 대표적인데요.

• SwiGLU FFN: 기존 GELU 기반 FFN 대신 gated linear unit 사용
• RMSNorm: LayerNorm 대신 Root Mean Square Normalization
• qk-Norm: attention의 query, key에 normalization 적용
• Rotary Positional Embedding (RoPE): 절대 위치 임베딩 대신 회전 기반 상대 위치 임베딩
• AdaLN-Zero: time step t와 class label 조건을 주입하는 방식

아래 Table 4에서 이들에 대한 ablation을 진행했는데요.

SwiGLU, RMSNorm, qk-Norm, RoPE 등 language model 쪽에서 가져온 개선 기법들이 각각 FID 개선에 기여하는 결과를 보였습니다. 또한 in-context class conditioning도 실험했는데, 기존 ViT가 하나의 class token을 추가하던 방식과 달리 다수의 class tokens을 sequence에 추가하는 방식을 사용하게 됩니다. 이를 사용했더니 또 성능이 좋아지는 결과가 나왔는데, 이러한 요소들은 diffusion에 특화된 설계가 아니라 애초에 Transformer 구조 개선의 결과물이기 때문에 JiT의 철학과 잘 맞을 수 있었던 것 같다고 생각이 듭니다. 이 Transformer의 강점 중 하나가, 다들 아시다시피 특정 task에 강하게 결합되지 않으면 다른 분야에서 개발된 architectural advances를 그대로 활용할 수 있다는 점인데요. 고도화된 U-Net 구조나 latent 특화 모듈 없이, 사실상 일반 language/vision model과 동일한 Transformer로 diffusion을 수행한다는 것은 모델 구조가 특정 도메인에 종속되지 않는다는 의미이기도 합니다. 이것이 diffusion을 task-agnostic Transformer로 공식화했을 때의 주요 장점이 되는 것 같다고 논문은 말하네요.

[3.4 학습 알고리즘: x-prediction + v-loss]

JiT의 최종 학습 알고리즘은 Table 1의 (3)(a)에 해당하는 x-prediction + v-loss 조합을 사용합니다. Transformer는 매 스텝마다 noisy image z_t를 받아 clean image patch x_pred를 직접 예측하고, 이로부터 velocity를 계산하는 구조입니다.

최적화 목표는 다음과 같습니다:

즉 예측은 x로, 학습은 v로 하는 구조입니다. 이 조합이 선택된 이유는 x-prediction이 pixel-space에서 manifold 구조 덕분에 안정적이고, v-loss의 1/(1-t)² 가중치가 gradient 균형을 잘 맞춰주고, flow matching ODE sampling과 자연스럽게 연결되기 때문이라고 합니다.

또한 x_pred는 최종 결과물이 아니라 flow 계산을 위한 중간 추정치라는 점이 중요합니다. JiT는 clean 이미지를 한 번에 예측하는 모델이 아니며 여전히 multi-step sampling을 수행합니다. 1/(1-t) 항에서 zero division을 피하기 위해 분모는 기본값 0.05로 clip하여 계산합니다.

Sampling 절차는 다음과 같습니다:

1. 초기 noise 이미지 z₀ 생성
2. patchify → embedding
3. Transformer로 x_pred(t) 예측
4. x_pred(t) → v_pred 계산
5. z 업데이트 (Heun / Euler ODE step)
6. 다시 patchify하여 반복
7. 50 step 정도 수행 → 최종 clean image 도달

[3.5 Noise Schedule]

시간 t의 sampling에는 logit-normal 분포를 사용합니다:

s ~ N(μ, σ²), t = sigmoid(s) = 1 / (1 + e^(-s))

μ는 logit-normal 분포에서 평균 logit을 의미하며, 값을 음수 방향으로 이동시키면 s가 더 작게 sample되어 평균 t가 작아지고, 결과적으로 더 높은 noise level의 샘플이 더 자주 선택됩니다. σ는 분산을 조절하며, 값이 커질수록 t 분포가 넓어져 0 또는 1에 가까운 극단적인 t도 더 자주 나옵니다. 이를 통해 train과 inference 과정에서 원하는 noise level의 평균 및 분산을 유연하게 제어할 수 있습니다.

4. Experiments

[4.1 9가지 조합 비교]

논문은 두 가지 실험 설정에서 9가지 prediction × loss 조합을 비교합니다:

설정 1: JiT-B/16, 256×256 이미지, patch size p=16 (Table 2(a))

patch dim이 768이고 모델의 hidden size도 768입니다. 겉보기에는 capacity가 맞는 것처럼 보이지만, 실제로는 positional embedding 등 추가 정보를 처리할 용량도 필요합니다.

x-prediction만 성능이 나오고 ε-prediction과 v-prediction은 loss space와 상관없이 모두 실패하는 결과가 나왔는데요(FID 300 이상으로 붕괴). ε/v-prediction은 노이즈가 섞인 정보와 원본 정보를 분리해 보존해야 하는데, 고차원 patch에서 이 작업을 수행하기에는 모델 capacity가 모자라 catastrophic failure가 발생합니다. 이는 앞서 toy experiment에서 관찰된 현상과 일치하는 경향성을 보입니다.

x-prediction 내에서는 세 가지 loss space 모두 작동하지만, v-loss가 FID 8.62로 가장 좋은 결과를 보였습니다. 다만 loss space 간 차이는 prediction space 선택에 비해 결정적이지는 않았습니다.

설정 2: JiT-B/4, 64×64 이미지, patch size p=4 (Table 2(b))

patch dim이 48로 hidden size 768보다 훨씬 작습니다. 이 경우에는 9개 조합 모두 무난하게 작동합니다. Capacity 문제가 발생하지 않기 때문입니다. 이전의 많은 latent diffusion 모델들이 입력 차원이 작았기 때문에, 지금까지 이 issue에 노출되지 않았던 것으로 보입니다. 왜 지금까지 pixel diffusion 모델은 latent space에 의존했는가?에 대해서 말하고자 하는 테이블인 것 같습니다.

[4.2 Loss weighting만으로는 불충분 (Table 2)]

이전 연구에서는 loss weighting과 network prediction 조합의 관계를 분석했지만, 대부분 저차원인 CIFAR-10 기반이었습니다. 저차원에서는 조합 간 차이가 크지 않았기에 적절한 loss weighting이면 충분하다는 결론이 나왔었는데요. 하지만 고차원인 ImageNet 256×256에서는 Table 2(a)에서 ε-prediction과 v-prediction이 loss space와 상관없이 모두 실패하므로, 이는 단순한 loss weighting 문제가 아니라 네트워크가 고차원 정보를 전파하지 못한다는 본질적 한계를 의미한다고 합니다.

[4.3 Noise level 조정만으로는 불충분 (Table 3)]

앞선 연구들은 high-resolution pixel-based diffusion에서 noise level을 높이는 것이 유용하다고 제안한 적이 있습니다. Table 3에서 logit-normal 분포의 μ를 조정해 noise level을 변화시킨 실험 결과를 보면 모델이 이미 괜찮게 작동하는 경우(x-prediction)에서는, 적절히 높은 noise가 유용하고 이건 이제 이전 연구들이랑 유사한데, noise level 조정만으로는 ε/v-prediction의 붕괴를 해결할 수 없고, μ를 어떻게 바꿔도 ε/v-prediction은 실패하는 경향이 생깁니다.

이는 고차원 정보를 전파하지 못하는 본질적 한계에서 비롯된 것으로, noise level 조정은 도움은 될 수 있어도 핵심 해결까지는 못한다는 것으로 해석했습니다.

[4.4 Hidden units 증가도 불필요 (Table 5, 6)]

입력데이터의 resolution이 커지면 Network 자체의 width(hidden units 수)를 늘려야 capacity 부족 문제가 해결될 것처럼 보일 수 있지만, 차원이 계속 늘어나면 차원의 저주로 인해 해결이 불가능한 경우도 왕왕있을텐데요. Table 5에서는 JiT/32(512 resolution, patch dim 3072)와 JiT/64(1024 resolution, patch dim 12288)를 평가한 결과를 보여주면서 JiT가 hidden unit까지 늘리지 않았음에도 성능을 잘 방어하는 이점을 보여줍니다.

• JiT/32, 512×512: patch dim이 3072, B/L/H 모델의 hidden size(768/1024/1280)보다 훨씬 큰데 반해 x-prediction이 안정적으로 잘 됐다는 FID 결과.
• JiT/64, 1024×1024: patch dim이 12288로 더욱 극단적인 상황, x-prediction으로 과한 collapse 없이 동작하는 모습. (1024×1024 pixel diffusion이 collapse 없이 돌아간 첫 사례 중 하나일 수 있다고 합니다!)

table 6에서는 또, noise를 resolution에 맞게 비례적으로 scaling하는 정도가 어떤지에 대해 분석한 결과인데요. 논문은 패치 크기를 이미지 해상도에 비례하게 늘려(JiT/16 @256 → JiT/32 @512) 시퀀스 길이(=패치 수)를 동일하게 유지하면, 모델 파라미터·연산량은 거의 같게 유지되면서도 고해상도에서의 FID가 스케일업과 함께 개선되거나 적어도 심각하게 악화되지 않는다는 점을 보였습니다.

결국 핵심은 network design이 observed dimensionality로부터 분리될 수 있으며, hidden units를 늘리는 것은 결정적 요소가 아니고, x-prediction의 경우 모델이 noise나 velocity기반 latent embedding같이 고차원 manifold로부터 복원하는 것이 아니라, 저차원 manifold 구조만 복원하면 되기 때문에 가능한 것으로 이해했습니다.

[4.5 ImageNet 벤치마크 결과]

ImageNet 256×256:

JiT-XL/16은 FID 2.38을 기록하며, 기존 latent diffusion 모델(DiT-XL/2 등)과 autoregressive 모델(MAR 등) 대비 대등하거나 더 나은 성능을 보입니다. FLOPs 관점에서도 latent diffusion 대비 효율적인 부분이 있는데, VAE encode/decode 비용이 아예 없기 때문입니다. VAE 없이 pixel-space에서 이 수준의 성능이 나온 것 자체가 핵심적인 결과라고 볼 수 있습니다.

ImageNet 512×512:

JiT/32 모델이 512 해상도에서도 competitive한 FID를 기록합니다. patch dim이 3072으로 모든 모델 변형의 hidden size보다 훨씬 큰 상황에서도 안정적으로 작동합니다.

ImageNet 1024×1024:

JiT/64로 1024 해상도까지 확장한 결과도 리포팅됩니다. patch dim이 12288이라는 극단적 상황에서도 collapse 없이 동작하는 것은, pixel-space diffusion의 가능성을 보여주는 의미 있는 성과입니다.

[4.6 x-prediction vs v-prediction 학습 안정성 비교 (Figure 7)]

Figure 7에서는 x-prediction과 v-prediction의 학습 곡선을 비교합니다. x-prediction이 v-prediction보다 loss가 더 낮고 안정적인 것을 확인할 수 있고, 실제 복원 결과도 x-prediction이 t가 낮을 때(노이즈가 많을 때) 좀 더 품질이 좋습니다. 이는 높은 노이즈 수준에서 x-prediction이 manifold projection 효과를 통해 더 안정적으로 복원 방향을 잡는다는 해석과 일치한 경향성을 보였습니다.

5. Discussion

[아쉬운 점]

• ImageNet class-conditional generation 스케일에서만 실험이 진행되어 text-to-image(T2I)나 multi-modal conditioning에 관한 실험이 없습니다. Stable Diffusion처럼 text prompt 기반의 자유로운 이미지 생성이 가능한지는 후속 연구에서 다룰 것 같습니다. 또 개인적으로는 robot-domain에 붙여보고싶네요.
• vram의 연산 효율이 LDM 계열보다 떨어질 가능성이 있다고 합니다. Pixel space를 다루기 때문에 고해상도 데이터에서 메모리 사용량이 매우 큰 것 같습니다.
• Sampling 속도 면에서도 50-step ODE solver를 사용하는 만큼, 최근의 few-step diffusion 모델들(consistency models 등)과 비교했을 때 인퍼런스 속도가 느립니다.
• pixel-space 모델이 latent-space 모델 대비 fine한 detail이 어떻게 다른지에 대한 정성적 비교가 더 있었으면 좋았을 것 같습니다. figure 5 에 대해서도 구체적인 언급은 딱히 없어보였습니다.

---

이 논문을 읽으면서 가장 인상깊었던 부분은, 결국 지금까지 pixel diffusion이 어려웠던 이유가 구조적 한계가 아니라 설계적 선택의 문제일 수도 있다는 인사이트였습니다. 관행적으로 그동안 pixel space는 너무 고차원이라 latent space가 필수적이다라는 것이 거의 공리처럼 받아들여져 왔는데, JiT는 prediction target을 x-prediction으로 다시 생각해보고 바꿔봤더니 이 공리가 틀릴 수도 있음을 보였기 때문인데요. Latent space가 필요했던 이유는 pixel space 자체가 다룰 수 없었기 때문이 아니라, pixel space에서 ε/v를 예측하는 것이 본질적으로 어려웠기 때문이라는 관점을 좀 새로 알아갈 수 있었습니다.

“복잡한 파이프라인 없이 기본 설계만으로도 풀 수 있는 문제인지,, 본질적인 문제에도 놓치지말고 집중해보자.”라는 게 최근 연구들이 점점 더 복잡한 multi-stage 학습이나 auxiliary loss를 막 덕지덕지 쌓아가는 흐름에서 봤을 때도 좋은 시사점이 되는 것 같습니다. 물론 이게 또 다른 태스크로 확장이 됐을 때 위와 같은 단순함이 유지될 수 있을지는 아직 미지수인 것 같긴 합니다. 개인적으로는 LeJEPA sigreg와 짝짝꿍이 잘 맞을 수 있는 구조인지 궁금해집니다.