[평범한 대학원생이 하는 논문 간단 요약] Back to Basics : Let Denoising Generative Models Denoise (ArXiv 2511)

https://arxiv.org/abs/2511.13720

Back to Basics: Let Denoising Generative Models Denoise

Abstract

Abstract를 가볍게 요약하면 아래와 같다.

• Clean image를 예측하는 것과 noise의 양을 예측하는 것은 다르다 $\rightarrow$ clean image를 예측하는게 manifold assumption에 맞다(?)
• 이에 맞는 Just image Transformer (JiT)를 제안한다.

1. Introduction

 기존의 diffusion model들은 noise prediction ($\epsilon$-prediction)을 하거나 flow velocity prediction ($v$-prediction)을 했다. Clean image prediction ($x$-prediction)은 이 둘과 관련이 있기 때문에, $\epsilon$과 $v$를 reformulation을 통해서 구할 수 있었다. 이를 통해서 기존에는 network가 어떤 것을 prediction하는지에 대해서 덜 관심을 가졌다.

 "(high-dimensional) data lie (roughly) on a low-dimensional manifold"라는 기존 manifold assumption에 따르면, clean data는 low-dimensional manifold에 놓이는 것으로 모델링될 수 있지만 noised quantity는 full high-dimensional space에 분포된다. 그래서 clean image를 예측하는 것은 noise나 noise quantity를 예측하는 것과 다르다.

 기존 diffusion model들은 latent space에서 동작하기 때문에, limited-capacity network이 noise를 예측하는데 어려움이 완화되지만 이는 해결된 것이 아니라 숨겨진 것이다. 그래서 이들은 pixel이나 high-dimensional space에서는 curse of dimensionality를 다루는데 어려움을 겪는다. 

 이를 해결하기 위한 기존 연구들이 존재하지만, 본 논문에서는 raw pixel들로 구성된 larget image patch들에 대해 작동하는 plain ViT가 effective diffusion modeling일 수 있다는 것을 보인다. 본 논문의 방식은 self-contained이며 pre-training과 auxiliary loss에 의존하지 않는다 - no latent tokenizer, no adversarial loss, no perceptual loss, no representation alignment. 본 논문의 model을 Just image Transforme (JiT)라고 부른다.

2. On Prediciton Outputs of Diffusion Models

2.1 Prediction Space and Loss Space

Prediction Space.

 Diffusion model은 크게 $x, \epsilon, v$-prediction이 존재한다. Network이 $x$를 output으로 내놓는다고 하면 다음과 같이 나타낼 수 있다.

즉, $x_\theta$와 $z_t$로부터 $\epsilon_\theta, v_\theta$가 계산될 수 있다.

Loss Space.

위 loss를 $v$-loss라고 한다. 그렇다면 $v$-loss와, $x$-pred를 조합하면 아래 table에서 (a)가 된다.

$v$-loss와 $x$-prediction을 조합하면, (3)은 $\mathcal{L} = \mathbb{E} \Vert v_\theta(z_t, t)-v\Vert^2 = \mathbb{E}\frac{1}{(1-t)^2} \Vert x_\theta(z_t,t) - x \Vert^2$가 된다. 이는 $x$-loss의 reweighted 형태이다. 이를 통해 table 1처럼 9개의 조합을 만들 수 있다. 이 조합들이 모두 수학적으로 동일하지 않다고 한다.

Generator Space.

 Inference시에 모두 v-space로 변환시킬 수 있기 때문에, 모두 다음 ODE를 풂으로써 샘플링할 수 있다.

2.2 Toy Experiment

 Figure 1에서처럼 noise $\epsilon$과 velocity $v$는 off-manifold인 반면 data $x$는 low-dimensional manifold에 존재한다. 그래서 network가 clean data $x$를 직접 prediction하도록 하는 것이 더 tractable하다. 이를 보이기 위해서 본 논문에서는 toy experiment를 진행한다.

본 논문은 256-dim의 hidden unit을 가진 5-layer ReLU MLP로 학습한다.  위 결과를 보면, $x$-prediction만이 D가 증가할 때, 적절한 결과를 생성한다. 이는 true data는 low-dimensional d-dim space에 존재하기 때문이다. $v, \epsilon$-prediction은 전체 dimension이 512면 model이 512차원을 나타내야하는데, 모델의 용량이 256이기 때문에 이를 표현할 수 없다. 그러나 $x$-prediction은 output의 차원인 2차원만을 표현하면 되기 때문에, 잘 표현한다.

3 "Just Image Transformers" for Diffusion

3.1 Just Image Transformers

 JiT의 구조는 기존 plain ViT를 사용한다.

3.2 What to Predict by the Network?

Table 2를 보면, $x$-prediction만 결과가 좋다. 그리고 loss는 크게 중요하지 않다는 것을 보인다. $t$를 샘플링하는 함수로는 logit-normal을 사용한다.

Hidden units의 수는 결정적이지 않고, network의 dimensionality를 줄이는 bottleneck을 도입하는게 효과적일 수 있다.

3.3 Algorithms

3.4 "Just Advanced" Transformers

4. Comparisons