[논문리뷰] Back to Basics: Let Denoising Generative Models Denoise

CVPR 2026. [Paper] [Github]
Tianhong Li, Kaiming He
MIT CSAIL
17 Nov 2025

Introduction

Diffusion model이 처음 개발되었을 때, 핵심 아이디어는 denoising, 즉 손상된 이미지에서 깨끗한 이미지를 예측하는 것이었다. 그러나 diffusion model 발전 과정에서 두 가지 중요한 전환점이 발생하면서 깨끗한 이미지를 직접 예측한다는 목표에서 벗어나게 되었다.

1. Noise 자체를 예측하는 $\epsilon$-prediction은 생성 품질을 획기적으로 향상시켰다.
2. Diffusion model은 깨끗한 데이터와 noise를 결합한 velocity를 예측함으로써 ($v$-prediction) flow 기반 방법과 연결되었다.

깨끗한 이미지를 예측하는 $x$-prediction은 예측 loss의 가중치를 적절히 재구성하면 $\epsilon$-prediction, $v$-prediction과 밀접한 관련이 있. 이러한 관계 때문에 네트워크가 직접 예측해야 하는 대상에 대해서는 상대적으로 관심이 적었다.

그러나 깨끗한 이미지와 noise가 포함된 값의 역할은 동일하지 않다. 머신러닝에서는 오랫동안 “고차원 데이터는 저차원 매니폴드 상에 존재한다”는 가설이 제기되어 왔다. 이러한 매니폴드 가정 하에서, 깨끗한 데이터는 저차원 매니폴드 상에 존재하는 것으로 모델링될 수 있는 반면, noise가 포함된 값은 본질적으로 고차원 공간에 분포되어 있다. 따라서 깨끗한 데이터를 예측하는 것은 noise 또는 noise가 포함된 값을 예측하는 것과는 근본적으로 다르다.

저차원 매니폴드가 고차원 관측 공간에 임베딩된 시나리오를 생각해 보자. 이 고차원 공간에서 noise를 예측하려면 높은 처리 용량이 필요하다. 즉, 신경망은 noise에 대한 모든 정보를 보존해야 한다. 반면, 처리 용량이 제한된 신경망이라도 noise를 제거하면서 저차원 정보만 유지하면 되므로 깨끗한 데이터를 예측할 수 있다. Latent space와 같은 저차원 공간을 사용할 경우 noise 예측의 어려움이 완화되지만, 근본적인 해결보다는 오히려 문제를 숨기는 결과를 낳는다. 사전 학습된 latent space에 대한 과도한 의존은 diffusion model이 자체적으로 완결성을 갖추지 못하게 한다.

본 논문에서는 기본 원리로 돌아가 신경망이 깨끗한 이미지를 직접 예측하도록 하였다. 이를 통해 픽셀로 구성된 대규모 이미지 패치에서 작동하는 단순한 ViT가 diffusion 모델링에 효과적일 수 있음을 보여준다. 본 논문의 접근 방식은 자체적으로 완결되어 있으며 사전 학습이나 보조 loss에 의존하지 않는다.

On Prediction Outputs of Diffusion Models

1. Prediction Space and Loss Space

네트워크의 출력은 임의의 space, 즉 $x$-space, $\epsilon$-space, $v$-space에서 정의될 수 있다. 이를 각각 $x$-prediction, $\epsilon$-prediction, $v$-prediction이라 부른다. Loss 역시 어떤 space에서든 정의할 수 있으며, 한 space에서 다른 space로 reparameterization이 가능하다.

3가지 loss와 3가지 네트워크 출력 형태를 종합적으로 고려하면 총 9가지 조합이 가능하다. 어떤 조합을 사용하든 inference-time에 생성을 하려면 항상 네트워크 출력을 $v$-space로 변환하고 샘플링을 위해 ODE를 풀어야 한다. 따라서 9가지 조합 모두 유효하며, 9가지 경우 중 수학적으로 동일한 경우는 없다.

2. Toy Experiment

매니폴드 가정에 따르면, 데이터 $x$는 저차원 매니폴드에 존재하는 경향이 있는 반면, noise $\epsilon$과 velocity $v$는 매니폴드 밖에 존재한다. 따라서 신경망이 깨끗한 데이터 $x$를 직접 예측하도록 하는 것이 더 다루기 쉽다. 저자들은 간단한 실험을 통해 이 가정을 검증하였다.

본 논문에서는 관측된 $D$차원 공간($d < D$)에 묻혀 있는 $d$차원 기저 데이터의 간단한 예를 고려하였다. 이 시나리오는 column-orthogonal projection matrix $P \in \mathbb{R}^{D \times d}$를 사용하여 합성한다. 행렬 $P$는 무작위로 생성되어 고정된다. 관측 데이터는 $x = P \hat{x} \in \mathbb{R}^D$이고, 기저 데이터는 $\hat{x} \in \mathbb{R}^d$이다. 모델은 행렬 $P$를 모르기 때문에 모델 입장에서는 $D$차원 생성 문제이다.

저자들은 hidden dimension이 256차원인 가진 5-layer ReLU MLP를 생성 모델로 학습시켰다. 구체적으로 $d=2$일 때 \(D \in \{2, 8, 16, 512\}\)의 경우를 조사하였다. 위 그림에서 볼 수 있듯이, $D$가 증가함에 따라 $x$-prediction만이 합리적인 결과를 도출할 수 있다. $\epsilon$-prediction과 $v$-prediction의 경우, 모델은 $D=16$에서 어려움을 겪고, 256차원 MLP가 불완전한 $D=512$에서는 심각한 실패를 보인다.

$x$-prediction은 모델이 불완전한 경우에도 잘 작동하며, 256차원 MLP는 필연적으로 512차원 공간의 정보를 버린다. 그러나 실제 데이터는 저차원 $d$차원 공간에 존재하므로, 이상적인 출력이 $d$차원이기 때문에 $x$-prediction은 여전히 ​​좋은 성능을 보일 수 있다.

“Just Image Transformers” for Diffusion

1. Just Image Transformers

ViT의 핵심 아이디어는 패치 기반 Transformer이다. 본 논문의 아키텍처 설계는 이 철학을 따른다.

$H \times W \times 3$ 차원의 이미지 데이터를 생각해 보자. 모든 $x$, $\epsilon$, $v$, $z_t$는 동일한 차원을 공유한다. 주어진 이미지를 겹치지 않는 $p \times p$ 크기의 패치로 나누어 $\frac{HW}{p^2}$ 길이의 시퀀스를 생성한다. 각 패치는 $3p^2$ 차원의 벡터이다. 이 시퀀스는 linear embedding projection을 거쳐 위치 임베딩이 더해지고, Transformer block들의 스택에 의해 매핑된다. 출력 layer는 각 토큰을 다시 $3p^2$ 차원의 패치로 projection하는 linear predictor이다.

일반적인 관행에 따라, 이 아키텍처는 시간 $t$와 주어진 클래스 레이블에 따라 컨디셔닝된다. 컨디셔닝에는 adaLN-Zero를 사용한다. 개념적으로, 이 아키텍처는 픽셀 패치에 직접 적용된 DiT와 같다.

2. What to Predict by the Network?

저자들은 9가지 조합에 대해 토큰마다 768차원인 JiT-B 모델을 학습시켰다. 그리고 다음과 같은 관찰 결과를 도출했다.

$x$-prediction이 매우 중요하며, loss weighting만으로는 충분하지 않다.

JiT-B/16을 사용했을 때, $x$-prediction만 우수한 성능을 보였으며, 세 가지 loss 모두에서 잘 작동했다. 여기서 패치의 크기는 768차원으로, JiT-B의 hidden size인 768과 일치한다. 언뜻 보기에는 충분해 보일 수 있지만, 실제로는 위치 임베딩 처리와 같은 추가적인 용량이 필요할 수 있다. $\epsilon$-prediction과 $v$-prediction의 경우, 모델은 noise가 포함된 값을 분리하고 유지하는 데 필요한 용량이 부족하다.

저자들은 비교를 위해 64$\times$64 해상도의 JiT-B/4를 살펴보았다. 여기서는 모든 경우에서 비교적 우수한 성능을 보였으며, 9가지 조합 간의 정확도 차이는 미미하여 결정적인 차이는 없었다. 각 패치의 차원은 48로, JiT-B의 hidden size인 768보다 훨씬 작다. 이것이 모든 조합이 비교적 잘 작동하는 이유이다.

Noise-level shift만으로는 충분하지 않다.

기존 연구들에서는 noise level을 높이는 것이 고해상도 픽셀 기반 diffusion에 유용하다고 제안했다. $t$를 샘플링하기 위해 logit-normal distribution을 사용하므로, 이 분포의 파라미터 $\mu$를 변경하여 noise level을 조정할 수 있다 ($\mu$를 음의 방향으로 이동시키면 noise level 증가).

모델이 이미 적절한 성능을 보일 때, 적절히 높은 noise는 유익하다. 그러나 noise level만 조정하는 것으로는 $\epsilon$-prediction 또는 $v$-prediction의 문제를 해결할 수 없다. 이러한 예측 실패는 본질적으로 고차원 정보를 전파할 수 없기 때문이다.

Bottleneck은 유익할 수 있다.

네트워크 폭(즉, hidden unit 수)에 따라 용량이 제한될 수 있으므로, 이를 늘리는 것이 자연스러운 해결책으로 여겨진다. 하지만 반대로 네트워크의 차원을 줄이는 bottleneck을 도입하는 것이 유익할 수 있다. 구체적으로, linear patch embedding layer를 한 쌍의 bottleneck layer로 대체하여 low-rank linear laye로 변환했다. 첫 번째 layer는 차원을 $d^\prime$으로 줄이고, 두 번째 layer는 Transformer의 hidden size까지 확장한다. 두 layer 모두 선형이며 low-rank reparameterization 역할을 한다.

위 그림은 JiT-B/16을 사용하여 bottleneck 차원 $d^\prime$에 대한 FID를 나타낸 것이다. Bottleneck 차원을 16차원까지 줄이더라도 심각한 오류가 발생하지 않는다. 오히려 bottleneck 차원을 적용하면 최대 약 1.3 FID까지 품질을 상당히 향상시킬 수 있다. 표현 학습이라는 더 넓은 관점에서 볼 때, bottleneck은 본질적으로 저차원 표현의 학습을 장려하는 역할을 한다.

3. Algorithm

최종 알고리즘은 $x$-prediction과 $v$-loss를 사용한다.

\[\begin{equation} \mathcal{L} = \mathbb{E}_{t, x, \epsilon} \| v_\theta (z_t, t) - v \|^2 \\ \textrm{where} \quad z_t = t \cdot x + (1-t) \cdot \epsilon, \quad v_\theta (z_t, t) = \frac{\textrm{net}_\theta (z_t, t) - z_t}{1 - t}, \quad v = \frac{x - z_t}{1 - t} \end{equation}\]

수치적 안정성을 위해 분모의 $(1-t)$는 0.05로 clipping된다.

4. “Just Advanced” Transformers

범용 Transformer의 강점 중 하나는 특정 task과 분리하여 설계할 경우 다른 응용 분야에서 개발된 아키텍처 개선 사항을 활용할 수 있다는 점이다. 이러한 특성은 task에 구애받지 않는 Transformer를 사용하여 diffusion model을 설계하는 데 있어 중요한 이점이다.

본 논문에서는 언어 모델을 위해 개발된 SwiGLU, RMSNorm, RoPE, qk-norm과 같은 널리 사용되는 범용 개선 사항들을 통합했다. 또한 저자들은 in-context class conditioning를 살펴보았는데, 시퀀스에 하나의 클래스 토큰을 추가하는 기존 ViT와 달리 여러 개의 토큰(기본 32개)을 추가했다.

Experiments

다음은 패치 크기가 따른 성능을 비교한 결과이다.

다음은 모델 크기에 따른 성능을 비교한 결과이다. (ImageNet 256$\times$256, 512$\times$512)

다음은 ImageNet 256$\times$256과 512$\times$512에서 기존 방법들과 비교한 결과이다.