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

Terminology

• x-prediction: 모델이 노이즈가 섞인 입력으로부터 깨끗한 원본 데이터(clean data)를 직접 예측하는 방식.
• epsilon-prediction: 모델이 원본 데이터에 추가된 노이즈(noise) 자체를 예측하는 방식. 현대 diffusion model에서 가장 널리 쓰이는 표준적인 방식.
• v-prediction: 원본 데이터와 노이즈가 결합된 형태인 속도(velocity)를 예측하는 방식.
• Manifold assumption: 고차원 공간의 자연 데이터(이미지 등)가 실제로는 유의미한 정보를 담고 있는 훨씬 낮은 차원의 매니폴드(manifold) 표면에 집중되어 존재한다는 머신러닝의 고전적 가설.
• JiT (Just image Transformers): Tokenizer, latent space, pre-training 등 복잡한 기법 없이, 순수하게 픽셀 패치(raw pixel patch) 위에 기본 ViT(Vision Transformer) 구조만을 적용한 이 논문의 제안 모델.
• Bottleneck linear embedding: 입력 패치를 고차원 hidden dimension으로 바로 매핑하지 않고, 의도적으로 낮은 차원(예: 128-d)을 거쳐 정보의 병목 현상을 유도한 뒤 다시 확장하는 구조.

Purpose of the Paper

• 기존 pixel-space diffusion model들이 겪는 고차원 연산의 한계("차원의 저주")를 극복하고자 함.
• 기존 연구들은 고차원 공간 전체에 분포하는 노이즈(epsilon)나 속도(v)를 예측하기 위해 복잡한 dense convolutions 구조나 latent space(외부 tokenizer 의존)를 강제받았음.
• 이 연구는 diffusion 본연의 개념인 "denoising"(깨끗한 데이터 예측)으로 회귀하여, manifold assumption에 기반해 x-prediction을 사용하면 단순한 architecture로도 고해상도 pixel-space 생성을 효율적으로 해결할 수 있음을 증명함.

Key Contributions

• x-prediction의 재발견과 근본적 차이 규명: 고차원 pixel-space에서 epsilon-prediction이나 v-prediction은 고용량의 네트워크를 요구해 catastrophic failure를 초래하지만, x-prediction은 저차원 manifold 정보만 유지하면 되므로 capacity가 제한된 네트워크에서도 매우 효과적임을 수학적 모델과 실험으로 증명함. 이는 단순한 loss weighting 효과를 넘어서는 본질적인 차이임.
• 순수 Transformer 기반 Pixel Diffusion (JiT) 제안: Tokenizer, latent space, pre-training, adversarial loss, perceptual loss 등 도메인 특화 기법을 완벽히 배제하고, 오직 plain ViT만을 원본 픽셀에 직접 적용하여 SOTA 급 성능을 달성한 혁신적 단순함을 보여줌.
• Bottleneck 구조의 이점 발굴: Manifold learning 원리에 착안하여 입력단에 의도적인 bottleneck을 주어 차원을 크게 줄였을 때(예: 768-d를 16-d로 압축) 오히려 생성 품질이 향상됨을 발견함. 이는 네트워크의 차원을 무조건 키워야 한다는 일반적인 딥러닝 직관을 뒤집는 참신한 결과임.

Experimental Highlights

• ImageNet 256 x 256 및 512 x 512 datasets에서 latent space 없이 픽셀에 직접 적용한 JiT 모델로 우수한 성능 달성 (JiT-G/16 기준 FID 1.82).
• Patch size가 커져 패치당 차원이 모델의 hidden dimension을 초과하는 극한의 고차원 설정(예: 1024 x 1024 해상도, 12288-d 패치)에서도, x-prediction 기반 JiT는 네트워크 크기를 무리하게 키우지 않고 patch size와 노이즈 스케일만 조절하여 안정적인 generation을 수행함.
• Loss space와 Network prediction space 조합을 모두 실험한 결과, 픽셀 차원이 높은 설정에서는 오직 x-prediction 그룹만이 FID 10 내외의 정상적인 결과를 보였으며, 나머지(epsilon, v)는 FID 300 이상의 catastrophic failure를 기록함을 확인하여 가설을 강력하게 검증함.

Limitations and Future Work

• Limitations: 모델의 순수성(minimalism)을 유지하고 baseline을 제시하기 위해, 성능을 즉각적으로 끌어올릴 수 있는 classification loss, perceptual loss, self-supervised pre-training 등의 최신 기법들을 의도적으로 배제함. 또한 가장 거대한 모델인 JiT-G의 경우 학습 후반부에 overfitting이 발생하는 경향이 있음.
• Future Work: 의도적으로 배제했던 extra loss나 pre-training 기법을 결합하면 SOTA 성능의 한계를 한층 더 돌파할 수 있음 (실제로 탐색 실험에서 단순한 classification loss 추가만으로 FID가 4.37에서 4.14로 향상됨을 확인).
• 도메인 확장성 (Broader Applicability): Tokenizer 설계가 극도로 까다로운 다른 자연 데이터 도메인(단백질 구조 예측, 분자 모델링, 기후 데이터 분석 등)에 이 "순수 Diffusion + Transformer" 패러다임을 직접 적용하여 범용적인 파운데이션 모델로 발전시킬 수 있는 잠재력이 매우 큼.

Overall Summary

이 논문은 현대 diffusion model들이 노이즈(epsilon)를 예측하느라 불필요하게 복잡한 architecture와 latent space에 과존하게 된 현실을 비판하며, 원본 깨끗한 이미지(x)를 직접 예측하는 "denoising"의 기본 원리로 돌아갈 것을 제안합니다. 이를 통해 그 어떤 사전 학습이나 복잡한 보조 기법 없이, 순수하고 투박한 ViT 구조(JiT)만으로 고해상도 pixel-space 생성을 성공적으로 수행해냈습니다. 궁극적으로 이 연구는 특정 도메인에 특화된 기술 의존도를 낮추고, 범용적이고 독립적인 "Diffusion + Transformer" 패러다임이 시각 영역을 넘어 다양한 기초 과학 도메인으로 뻗어 나갈 수 있는 강력한 토대를 마련했다는 점에서 큰 의의를 지닙니다.

쉬운 설명

이 논문은 "모래사장(고차원 노이즈 공간)에서 모래알(노이즈)의 위치를 하나하나 전부 다 기억해서 치우는 것" 대신, "그 아래에 파묻힌 보물상자(저차원 원본 데이터 매니폴드)의 생김새만 직접 추측하는 것"이 딥러닝 모델에게 훨씬 쉬운 문제라고 설명합니다. 기존의 최신 AI들은 모래알 전체를 예측하느라 복잡하고 무거운 기술들이 필요했지만, 이 논문은 보물상자의 형태(원본 픽셀)만 직접 예측하도록 목표를 바꿔주었습니다. 그 결과, 아무런 기교 없이 아주 단순한 기본 모델 구조만으로도 깨끗하고 선명한 고화질 사진을 만들어내는 데 성공했습니다.

 

 

 

Abstract

오늘날의 denoising diffusion models는 고전적인 의미에서 denoise를 수행하지 않습니다. 즉, clean images를 직접적으로 predict하지 않습니다. 오히려 neural networks는 noise나 noised quantity를 predict합니다. 본 논문에서는 clean data를 predict하는 것과 noised quantities를 predict하는 것이 근본적으로 다르다고 제안합니다.

manifold assumption에 따르면, natural data는 low-dimensional manifold 상에 존재해야 하지만, noised quantities는 그렇지 않습니다. 이러한 가정을 바탕으로, 우리는 clean data를 직접적으로 predict하는 models를 지지하며, 이는 겉보기에 under-capacity networks처럼 보이는 모델들이 매우 high-dimensional spaces에서 효과적으로 작동할 수 있게 해줍니다.

우리는 pixels 기반의 단순한 large-patch Transformers가 강력한 generative models가 될 수 있음을 보여줍니다. 이 과정에서 tokenizer, pre-training, 그리고 추가적인 loss는 사용하지 않습니다. 우리의 approach는 개념적으로 "Just image Transformers", 즉 우리가 명명한 JiT에 지나지 않습니다.

우리는 high-dimensional noised quantities를 predict할 때 치명적인 실패를 겪을 수 있는 256 및 512 resolutions의 ImageNet 환경에서, 16 및 32의 큰 patch sizes를 가진 JiT를 사용하여 competitive results를 달성했음을 보고합니다.

우리의 networks가 manifold의 기본으로 다시 매핑됨에 따라, 본 연구는 기본으로 돌아가 raw natural data에 대한 Transformer-based diffusion을 위한 self-contained paradigm을 추구합니다.

 

 

1. Introduction

diffusion generative models가 처음 개발되었을 때, 그 핵심 idea는 denoising, 즉 손상된 버전에서 clean image를 predicting하는 것이었습니다. 그러나 diffusion models의 발전 과정에서 나타난 두 가지 중요한 마일스톤은 clean images를 직접 predicting하는 목표에서 벗어나는 결과를 낳았습니다.

첫째, noise 자체를 predicting하는 것($\epsilon$-prediction으로 알려짐)은 generation 품질에 결정적인 차이를 만들어냈고 이러한 models를 크게 대중화시켰습니다. 나중에 diffusion models는 clean data와 noise를 결합한 quantity인 flow velocity(v-prediction)를 predicting함으로써 flow-based methods와 연결되었습니다. 오늘날 실제 환경의 diffusion models는 일반적으로 noise 또는 noised quantity(예: velocity)를 predict합니다.

광범위한 연구들은 prediction loss의 weighting이 적절히 재구성된다면(Sec. 3에 자세히 설명됨), clean image를 predicting하는 것(x-prediction)이 $\epsilon$- 및 v-prediction과 밀접하게 관련되어 있음을 보여주었습니다. 이러한 관계 때문에, network가 할당된 task를 수행할 능력이 있다고 implicitly 가정하면서, network가 직접 무엇을 predict해야 하는지에 대해서는 덜 주목해 왔습니다.

그러나 clean images와 noised quantity(noise 자체 포함)의 역할은 결코 같지 않습니다. machine learning 분야에서는 "(high-dimensional) data는 (대략적으로) low-dimensional manifold에 존재한다"는 가설이 오랫동안 제기되어 왔습니다. 이러한 manifold assumption 하에서, clean data는 low-dimensional manifold에 존재하는 것으로 modeling될 수 있는 반면, noised quantity는 본질적으로 전체 high-dimensional space에 걸쳐 넓게 분포합니다(Fig. 1 참조). 따라서 clean data를 predicting하는 것은 noise나 noised quantity를 predicting하는 것과 근본적으로 다릅니다.

low-dimensional manifold가 high-dimensional observation space에 내장된(embedded) 시나리오를 생각해 보십시오. 이 high-dimensional space에서 noise를 predicting하려면 높은 capacity가 필요합니다. 즉, network는 noise에 대한 모든 information을 보존해야만 합니다. 대조적으로, limited-capacity network는 noise를 filtering하면서 low-dimensional information만 유지하면 되기 때문에 여전히 clean data를 predict할 수 있습니다.

low-dimensional space(예: image latent)가 사용될 때, noise를 predicting하는 어려움은 완화되지만 동시에 이는 해결되기보다는 숨겨질 뿐입니다. pixel이나 다른 high-dimensional spaces를 다룰 때, 기존의 diffusion models는 dimensionality의 저주를 해결하는 데 여전히 어려움을 겪을 수 있습니다. pre-trained latent space에 대한 과도한 의존은 diffusion models가 self-contained되는 것을 방해합니다.

self-contained 원칙을 추구함에 있어, pixel space에서의 diffusion modeling을 발전시키는 데 강한 초점이 맞춰져 왔습니다. 일반적으로 이러한 methods는 dense convolutions나 더 작은 patches를 사용하거나, channels를 늘리거나, 긴 skip connections를 추가하는 등의 방식으로 networks 내의 information bottleneck을 explicitly 혹은 implicitly 피합니다. 우리는 이러한 디자인들이 high-dimensional noised quantities를 predict하려는 요구에서 비롯되었을 수 있다고 제안합니다.

본 논문에서 우리는 첫 번째 원칙으로 돌아가 neural network가 직접 clean image를 predict하도록 합니다. 그렇게 함으로써, 우리는 raw pixels로 구성된 큰 image patches에서 작동하는 일반적인 Vision Transformer (ViT)가 diffusion modeling에 매우 효과적일 수 있음을 보여줍니다.

우리의 approach는 self-contained되어 있으며, 어떠한 pre-training이나 auxiliary loss에도 의존하지 않습니다. 즉, latent tokenizer, adversarial loss, perceptual loss(따라서 pre-trained classifier도 없음), 그리고 representation alignment(따라서 self-supervised pre-training도 없음)가 전혀 필요하지 않습니다. 개념적으로 우리의 model은 우리가 JiT라고 부르는 "Just image Transformers"를 diffusion에 적용한 것에 불과합니다.

우리는 256 및 512의 resolutions에서 각각 16 및 32의 patch sizes를 가진 JiT models를 사용하여 ImageNet dataset에 대한 experiments를 수행했습니다. patches가 매우 high-dimensional(수백 또는 수천)임에도 불구하고, x-prediction을 사용하는 우리의 models는 $\epsilon$- 및 v-prediction이 비참하게 실패하는 곳에서도 쉽게 강력한 결과를 만들어낼 수 있습니다.

추가적인 analysis는 network width가 patch dimension과 일치하거나 초과할 필요가 없음을 보여줍니다. 사실, 놀랍게도 bottleneck 디자인이 오히려 유익할 수 있으며, 이는 고전적인 manifold learning에서의 관찰과 일치합니다.

우리의 노력은 native data에 대한 self-contained "Diffusion + Transformer" 철학을 향한 의미 있는 한 걸음을 나타냅니다. computer vision을 넘어, 이 철학은 tokenizer를 디자인하기 어려울 수 있는 natural data(예: proteins, molecules, 또는 weather)를 다루는 다른 도메인에서도 매우 바람직합니다. domain-specific designs를 최소화함으로써, 우리는 computer vision에서 비롯된 일반적인 "Diffusion + Transformer" paradigm이 더 넓은 applicability를 찾게 되기를 희망합니다.

 

논문 요약 정리 노트: Introduction

1. 문제 제기: $\epsilon$-prediction과 v-prediction의 근본적 한계

• Manifold assumption 위배: clean data는 low-dimensional manifold에 존재하지만, noise나 noised quantity는 전체 high-dimensional space에 넓게 분포함.
• Capacity 문제: high-dimensional space에서 noise를 predict하려면 모든 noise의 information을 보존해야 하므로 막대한 network capacity가 요구됨.
• 기존 연구의 우회: 기존 diffusion models는 이 문제를 해결하기 위해 latent space를 도입하여 본질적인 문제를 숨기거나, 픽셀 단위에서는 dense convolutions, 작은 patches 등 복잡한 architecture를 강제하여 해결하려 했음.

2. 핵심 제안: x-prediction으로의 회귀와 Self-contained 모델링

• x-prediction의 재발견: 모델이 noise가 아닌 clean image를 직접 predict하도록 학습. network는 noise를 filtering하고 low-dimensional information만 유지하면 되므로 limited-capacity network로도 충분함.
• JiT (Just image Transformers) 도입: 어떠한 꼼수 없이 원본 raw pixels에 기반한 large image patches를 일반적인 ViT에 직접 입력.
• 완전한 Self-contained: latent tokenizer, pre-training, auxiliary loss, adversarial loss, perceptual loss 등이 일절 필요 없는 순수한 "Diffusion + Transformer" 구조 제안.

3. 주요 실험 결과 및 발견

• 대형 패치에서의 압도적 성능: ImageNet 256 및 512 resolutions에서 아주 큰 patch sizes(16, 32)를 사용했을 때, $\epsilon$-prediction과 v-prediction은 완전히 실패(catastrophic failure)했으나, x-prediction을 적용한 JiT는 매우 강력한 성능을 보여줌.
• Bottleneck 구조의 유효성: network width가 patch dimension을 반드시 초과할 필요가 없으며, 오히려 bottleneck 디자인이 유익할 수 있음을 확인. 이는 classical manifold learning의 관찰 결과와 완벽히 일치함.
• 확장성 기대: tokenizer 설계가 까다로운 타 도메인(proteins, molecules 등)의 natural data에도 쉽게 적용 가능한 일반화된 paradigm 제시.

---

쉬운 설명 :

지금까지 대부분의 diffusion models는 그림에 끼어있는 '노이즈(noise)' 자체를 예측하도록 학습해 왔습니다. 하지만 노이즈는 규칙 없이 사방으로 퍼져 있는 아주 복잡한 데이터(high-dimensional)이기 때문에, 이를 정확히 예측하려면 모델이 불필요하게 똑똑해야 하고 덩치도 커야 했습니다. 그래서 사람들은 이미지를 압축(latent)해서 이 문제를 요리조리 피해갔죠.

이 논문의 저자들은 발상을 전환합니다. **"진짜 깨끗한 그림(clean data)은 사실 핵심적인 규칙(low-dimensional manifold)만 알면 되잖아? 복잡한 노이즈를 예측하느라 헛고생하지 말고, 그냥 깨끗한 그림을 바로 예측(x-prediction)하게 만들자!"**라고 말입니다.

이렇게 목표를 '진짜 그림 찾기'로 바꾸니까 마법 같은 일이 일어납니다. 복잡하게 이미지를 압축하는 과정(tokenizer)이나 미리 학습된 데이터(pre-training) 없이, 그냥 아주 평범한 Transformer 모델에 원본 픽셀을 큼직큼직하게 썰어서(large patches) 던져주기만 해도 모델이 알아서 그림을 아주 잘 그려낸다는 것입니다. 심지어 기존의 노이즈 예측 방식이 완전히 박살 나는 가혹한 조건에서도 말이죠. 요약하자면, **"복잡한 꼼수 쓰지 말고, 기본으로 돌아가서 진짜 그림만 예측하게 하니 평범한 모델로도 다 되더라"**가 이 섹션의 핵심입니다.

 

Related Work

 

Diffusion Models and Their Predictions. diffusion models에 관한 선구적인 연구는 network가 normal distribution의 parameters(예: mean과 standard deviation)를 predict하는 reversed stochastic process를 학습할 것을 제안했습니다. 도입 후 5년이 지나, 이 method는 Denoising Diffusion Probabilistic Models (DDPM)에 의해 부흥하고 대중화되었습니다. 결정적인 발견은 noise 자체를 prediction target으로 삼는 것(즉, $\epsilon$-prediction)이었습니다.

이후 다양한 prediction targets 간의 관계가 (원래 model distillation의 맥락에서) 연구되었으며, 여기서 v-prediction의 개념도 도입되었습니다. 그들의 연구는 reparameterization에 의해 도입된 weighting effects에 초점을 맞추었습니다.

한편, EDM은 denoiser function을 중심으로 diffusion problem을 재구성했으며, 이는 diffusion models 진화의 주요한 이정표를 구성합니다. 그러나 EDM은 pre-conditioned formulation을 채택하여 network의 직접적인 output이 denoised image가 아니게 만들었습니다. 이 formulation은 low-dimensional 시나리오에서는 선호되지만, 본질적으로 network가 data와 noise가 혼합된 quantity를 output하도록 여전히 요구합니다(appendix에 추가 비교 있음).

Flow Matching models는 diffusion modeling framework 내에서 일종의 v-prediction 형태로 해석될 수 있습니다. 순수한 noise와 달리, $v$는 data와 noise의 결합입니다. flow-based models와 이전 diffusion models 간의 연결 고리는 이미 확립되었습니다. 오늘날 diffusion models와 그에 대응하는 flow-based counterparts는 종종 통합된 framework 하에서 연구됩니다.

Denoising Models. 수십 년 동안 denoising의 개념은 representation learning과 밀접한 관련이 있었습니다. BM3D와 기타 모델들로 대표되는 고전적인 methods는 image denoising을 수행하기 위해 sparsity와 low dimensionality 가정을 활용했습니다.

Denoising Autoencoders (DAEs)는 denoising을 training objective로 사용하는 unsupervised representation learning method로 개발되었습니다. 그들은 manifold assumption (Fig. 1)을 활용하여 low-dimensional data manifold를 근사하는 의미 있는 representations를 학습했습니다. DAEs는 Denoising Score Matching의 한 형태로 볼 수 있으며, 이는 결과적으로 현대의 score-based diffusion models와 밀접한 관련이 있습니다. 그럼에도 불구하고, DAEs가 manifold learning을 위해 clean data를 predict하는 것은 자연스러운 일이지만, score matching에서 score function을 predict하는 것은 사실상 (scaling factor를 제외하고) noise를 predict하는 것, 즉 $\epsilon$-prediction에 해당합니다.

Manifold Learning. manifold learning은 observed data로부터 low-dimensional, nonlinear representations를 학습하는 데 중점을 둔 고전적인 분야입니다. 일반적으로 manifold learning methods는 유용한 information만 통과하도록 장려하는 bottleneck structures를 활용합니다. 여러 연구에서 manifold learning과 generative models 간의 연결 고리를 탐구했습니다. Latent Diffusion Models (LDMs)는 첫 번째 단계에서 autoencoder를 통한 manifold learning으로, 두 번째 단계에서 diffusion으로 간주될 수 있습니다.

Pixel-space Diffusion. 오늘날 latent diffusion이 이 분야의 기본 선택이 되었지만, diffusion models의 개발은 원래 pixel-space formulations로 시작되었습니다. 초기의 pixel-space diffusion models는 일반적으로 dense convolutional networks, 가장 흔하게는 U-Net을 기반으로 했습니다. 이러한 models는 종종 over-complete channel representations(예: 첫 번째 layer에서 $H \times W \times 3$을 $H \times W \times 128$로 변환)와 함께 long-range skip connections를 사용합니다. 이러한 models는 $\epsilon$- 및 v-prediction에서 잘 작동하지만, 그들의 dense convolutional structures는 일반적으로 계산적으로 비쌉니다. 이러한 convolutional models를 high-resolution images에 적용해도 catastrophic degradation이 발생하지 않으며, 이 방향의 연구는 generation quality를 더욱 향상시키기 위해 noise schedules 및/또는 weighting schemes에 초점을 맞추는 것이 일반적이었습니다.

이와 대조적으로, Vision Transformer (ViT)를 pixels에 직접 적용하는 것은 더 도전적인 task를 제시합니다. Standard ViT architectures는 공격적인 patch size(예: $16 \times 16$ pixels)를 채택하여, Transformer의 hidden dimension과 비슷하거나 더 큰 high-dimensional token space를 초래합니다. SiD2와 PixelFlow는 더 작은 patches에서 시작하는 hierarchical designs를 채택하지만, 이러한 models는 "FLOP-heavy"하며 standard Transformers에 내재된 generality와 simplicity를 잃어버립니다.

PixNerd는 Transformer output, noisy input, spatial coordinates의 information을 통합하는 NeRF head를 채택하며, representation alignment의 도움을 받아 training을 진행합니다.

이러한 특수 목적의 designs를 사용하더라도, 이러한 연구들의 architectures는 일반적으로 "L" (Large) 또는 "XL" 크기에서 시작합니다. 실제로 최근 연구는 높은 dimensionality를 위해서는 큰 hidden size가 필수적인 것으로 보인다고 제안합니다.

High-dimensional Diffusion. ViT-style architectures를 사용할 때, 현대의 diffusion models는 pixels든 latents든 상관없이 여전히 high-dimensional input spaces에 의해 도전을 받습니다. 문헌에서 계속해서 보고되듯, pixels 또는 latents의 사용 여부에 관계없이 per-token dimensionality가 증가할 때 ViT-style diffusion models는 빠르고 치명적으로(catastrophically) 성능이 저하됩니다.

우리의 연구와 동시에 진행된 일련의 연구들은 high-dimensional diffusion을 해결하기 위해 self-supervised pre-training에 의존합니다. 이러한 노력들과 대조적으로, 우리는 어떠한 pre-training도 없이 단지 Transformers만 사용하여 high-dimensional diffusion이 달성 가능함을 보여줍니다.

x-prediction. x-prediction의 formulation은 자연스럽고 새로운 것이 아니며, 적어도 원래의 DDPM(그들의 코드 참조)까지 거슬러 올라갈 수 있습니다. 그러나 DDPM은 $\epsilon$-prediction이 실질적으로 더 우수함을 관찰했고, 이것이 나중에 주요 해결책이 되었습니다. 이후 연구들에서 분석은 때때로 x-space에서 수행되는 것을 선호했지만, 실제 prediction은 관습적인 이유로 인해 종종 다른 spaces에서 이루어졌습니다.

diffusion이 다루는 image restoration application에 있어서는, network가 clean data를 predict하는 것이 자연스러우며, 이는 image restoration의 궁극적인 목표이기도 합니다. 우리의 연구와 동시에, 이전 frames를 조건으로 하는 generative world models를 위해 x-prediction의 사용을 지지하는 연구도 있습니다.

우리의 연구는 이 기본적인 개념을 재발명하려는 것이 아니라, 기저에 low-dimensional manifolds를 가진 high-dimensional data의 맥락에서 대체로 간과되었지만 중요한 문제에 주의를 환기시키고자 합니다.

 

논문 요약 정리 노트: Related Work

1. Prediction Targets의 역사적 흐름과 본질적 차이

• 대세가 된 Noise 예측: DDPM 이후 $\epsilon$-prediction이 표준이 되었고, EDM이나 Flow Matching(v-prediction) 역시 본질적으로 noise가 섞인 quantity를 예측하는 방향으로 발전함. Score matching 역시 결국 $\epsilon$-prediction과 궤를 같이함.
• 고전적 Denoising과의 괴리: 과거 Denoising Autoencoders (DAEs)는 manifold assumption을 바탕으로 clean data를 직접 예측하여 low-dimensional manifold를 학습했음. 현재의 LDM은 이를 두 단계(Autoencoder로 압축 후 Diffusion)로 나누어 해결하고 있음.

2. Pixel-space Diffusion과 ViT의 한계

• CNN 기반 모델의 비효율성: U-Net 같은 기존 dense convolutional networks는 over-complete channel representations를 사용하여 $\epsilon$-prediction에는 잘 작동하지만 계산량이 너무 큼.
• ViT 적용 시의 차원의 저주: standard ViT를 pixel에 직접 적용하면 patch size로 인해 token space가 high-dimensional이 되어 성능이 급락(catastrophic degradation)함.
• 기존 연구들의 우회책: 이를 해결하기 위해 최근 연구들(SiD2, PixelFlow, PixNerd 등)은 복잡한 hierarchical designs, 과도한 연산량(FLOP-heavy), 추가적인 head 도입, 또는 L/XL 사이즈 이상의 거대한 모델 크기를 강제해야 했음.

3. 본 논문의 독창적 포지셔닝 (Core Contribution)

• Pre-training 없는 High-dimensional Diffusion 달성: 동시대의 연구들이 차원의 저주를 풀기 위해 self-supervised pre-training에 의존하는 것과 달리, 본 논문은 어떠한 pre-training 없이 순수 Transformers만으로 이를 해결함.
• x-prediction의 재발견: x-prediction(clean data 예측)은 DDPM 시절에도 있었으나 성능 문제로 버려졌던 개념임. 본 연구는 이 잊혀진 개념을 **"high-dimensional data + low-dimensional manifold"**라는 맥락으로 다시 가져와, 기존 방식들이 겪던 ViT의 차원 한계를 근본적으로 돌파함.

---

쉬운 설명 :

이 섹션은 **"남들은 다 어떻게 실패했고, 우리는 왜 꼼수 없이 성공했는가?"**를 설명하는 부분입니다.

최근 AI 트렌드에서 이미지를 생성할 때 노이즈(noise)를 예측하는 방식이 대성공을 거두면서, 다들 당연하게 **"디퓨전 = 노이즈 예측"**이라고 생각하게 되었습니다. 동시에 요즘 대세인 Transformer(ViT) 모델을 원본 이미지 픽셀에 바로 적용해 보려는 시도들이 많았는데, 픽셀 데이터가 너무 방대하고 복잡하다 보니 모델이 버티질 못하고 성능이 처참하게 망가졌습니다. 그래서 다른 연구자들은 모델 크기를 무식하게 키우거나, 복잡한 계층 구조를 만들거나, 미리 다른 데이터로 예습(pre-training)을 시키는 등 온갖 '우회로'와 '꼼수'를 써야만 했습니다.

하지만 이 논문의 저자들은 과거의 지혜(DAE)를 떠올립니다. "원래 옛날에는 딥러닝으로 노이즈를 지울 때 깨끗한 원본 이미지를 바로 예측하게 했잖아?"라는 것이죠. 초창기 디퓨전 모델(DDPM) 시절에 쓰다가 성능이 안 나와서 버려졌던 '원본 이미지 직접 예측(x-prediction)' 방식을 다시 꺼내 들었습니다.

결과는 놀라웠습니다. 노이즈 대신 진짜 그림을 찾으라고 목표만 바꿔주었더니, 다른 연구자들이 쓰던 온갖 꼼수나 사전 학습 없이도 아주 평범하고 심플한 Transformer 모델이 방대한 픽셀 데이터를 훌륭하게 처리해낸 것입니다. 즉, 문제를 복잡하게 풀려던 기존 연구들과 달리, 문제의 본질(예측 대상)을 바꿈으로써 가장 단순하고 우아한 정공법을 찾아냈다는 것이 이 섹션의 핵심입니다.

 

 

 

On Prediction Outputs of Diffusion Models

 

Diffusion models는 $x, \epsilon,$ 또는 $v$의 space에서 공식화될 수 있습니다. space의 선택은 loss가 정의되는 위치뿐만 아니라 network가 예측(predict)하는 대상도 결정합니다. 중요한 점은, loss space와 network output space가 반드시 같을 필요는 없다는 것입니다. 이러한 선택은 결정적인 차이를 만들 수 있습니다.

3.1. Background: Diffusion and Flows

Diffusion models는 ODEs의 관점에서 공식화될 수 있습니다. 우리는 더 간단한 출발점으로서 flow-based paradigm, 즉 v-space에서 공식화를 시작한 다음 다른 spaces를 논의할 것입니다.

data distribution $x \sim p_{data}(x)$와 noise distribution $\epsilon \sim p_{noise}(\epsilon)$ (예: $\epsilon \sim N(0, I)$)을 고려해 봅시다. training 동안, noisy sample $z_t$는 interpolation입니다: $z_t = a_t x + b_t \epsilon$. 여기서 $a_t$와 $b_t$는 시간 $t \in [0, 1]$에서 미리 정의된(pre-defined) noise schedules입니다. 본 논문에서는 linear schedule을 사용합니다: $a_t = t$ 이고 $b_t = 1 - t$. 이는 다음을 제공합니다:

$$z_t = t x + (1 - t) \epsilon, \quad (1)$$

이는 $t=1$일 때 $z_t \sim p_{data}$가 되도록 이끕니다. 우리는 $t$에 대해 logit-normal distribution을 사용합니다. 즉, $logit(t) \sim N(\mu, \sigma^2)$ 입니다.

flow velocity $v$는 $z$의 시간에 대한 도함수(time-derivative)로 정의됩니다. 즉, $v_t = z'_t = a'_t x + b'_t \epsilon$ 입니다. Eq. (1)이 주어지면 다음을 얻습니다:

$$v = x - \epsilon. \quad (2)$$

flow-based methods는 다음과 같이 정의된 loss function을 최소화합니다:

$$L = E_{t,x,\epsilon} \lVert v_\theta(z_t, t) - v \rVert^2, \quad (3)$$

여기서 $v_\theta$는 $\theta$로 parameterize된 function입니다. $v_\theta$는 종종 network의 직접적인 output($v_\theta = net_\theta(z_t, t)$)이 되지만, 우리가 자세히 설명하겠지만 그것의 transform이 될 수도 있습니다.

function $v_\theta$가 주어지면, $z_0 \sim p_{noise}$에서 시작하여 $t = 1$에서 끝나는 $z$에 대한 ordinary differential equation (ODE)를 풀어서 sampling을 수행합니다:

$$d z_t / dt = v_\theta(z_t, t), \quad (4)$$

실제 환경에서, 이 ODE는 numerical solvers를 사용하여 근사적으로 풀 수 있습니다. 기본적으로(By default) 우리는 50-step Heun을 사용합니다.

3.2. Prediction Space and Loss Space

Prediction Space. network의 직접적인 output은 $v, x,$ 또는 $\epsilon$ 중 어떤 space에서도 정의될 수 있습니다. 다음으로, 우리는 그에 따른 transformation을 논의합니다. 본 논문의 맥락에서, 우리는 network $net_\theta$의 직접적인 output이 엄격하게 각각 $x, \epsilon, v$일 때만 그것을 "$x, \epsilon, v$-prediction"이라고 지칭한다는 점에 유의하십시오.

세 개의 unknowns $(x, \epsilon, v)$와 하나의 network output이 주어졌을 때, 세 개의 unknowns를 모두 결정하기 위해서는 두 개의 추가적인 constraints가 필요합니다. 두 constraints는 Eq. (1)과 (2)에 의해 주어집니다.

예를 들어, 직접적인 network output $net_\theta$가 $x$가 되도록 할 때, 우리는 다음 방정식 세트를 풉니다:

$$\begin{cases} x_\theta = net_\theta \\ z_t = t x_\theta + (1 - t) \epsilon_\theta \\ v_\theta = x_\theta - \epsilon_\theta \end{cases} \quad (5)$$

여기서 표기법 $x_\theta, \epsilon_\theta, v_\theta$는 그것들이 모두 $\theta$에 의존하는 predictions임을 시사합니다. 이 방정식 세트를 풀면 다음을 얻습니다: $\epsilon_\theta = (z_t - t x_\theta) / (1 - t)$ 이고 $v_\theta = (x_\theta - z_t) / (1 - t)$. 즉, $\epsilon_\theta$와 $v_\theta$ 모두 $z_t$와 network $x_\theta$로부터 계산될 수 있습니다. 이것들은 Tab. 1의 (a) 열에 요약되어 있습니다.

마찬가지로, 직접적인 network output $net_\theta$가 $\epsilon$ 또는 $v$가 되도록 할 때, 우리는 (Eq. (5)의 첫 번째 방정식을 교체함으로써) 다른 방정식 세트들을 얻습니다. transformations는 Tab. 1의 $\epsilon$-, v-prediction을 위한 (b), (c) 열에 요약되어 있습니다.

이것은 $\{x, \epsilon, v\}$ 중 하나의 quantity가 predict될 때, 다른 두 개가 추론될 수 있음을 보여줍니다. 많은 선행 연구에서의 유도(derivations)는 Tab. 1에 포함된 특수한 경우들입니다.

Loss Space. loss는 종종 하나의 참조 space(예: Eq. (3)의 v-loss)에서 정의되지만, 개념적으로 어떤 space에서도 정의될 수 있습니다. 한 prediction space에서 다른 space로 주어진 reparameterization을 통해, loss가 효과적으로 reweighted된다는 것이 입증되었습니다.

예를 들어, Tab. 1(3)(a)에 있는 x-prediction과 v-loss의 조합을 생각해 봅시다. 우리는 prediction으로 $v_\theta = (x_\theta - z_t) / (1 - t)$를 갖고, target으로 $v = (x - z_t) / (1 - t)$를 갖습니다. Eq. (3)의 v-loss는 다음과 같이 됩니다: $L = E \lVert v_\theta(z_t, t) - v \rVert^2 = E \frac{1}{(1-t)^2} \lVert x_\theta(z_t, t) - x \rVert^2$, 이는 x-loss의 reweighted된 형태입니다. 이와 같은 transformation은 Tab. 1에 나열된 모든 prediction space와 모든 loss space에 대해 수행될 수 있습니다.

종합하면, ${x, \epsilon, v}$에서 정의된 3개의 unweighted losses와 network 직접 output의 3가지 형태를 고려할 때, 9가지 가능한 조합이 있습니다(Tab. 1). 각 조합은 유효한 formulation을 구성하며, 9가지 경우 중 수학적으로 동등한 두 개는 없습니다.

Generator Space. 사용된 조합에 관계없이 inference-time에 generation을 수행하기 위해, 우리는 언제나 network output을 v-space(Tab. 1, row (3))로 변환하고 sampling을 위해 Eq. (4)의 ODE를 풀 수 있습니다. 이처럼, 9가지 조합 모두 합법적인(legitimate) generators입니다.

3.3. Toy Experiment

manifold assumption에 따르면, data $x$는 low-dimensional manifold에 존재하는 경향이 있는 반면(Fig. 1), noise $\epsilon$과 velocity $v$는 off-manifold입니다. network가 직접 clean data $x$를 predict하도록 하는 것이 더 다루기 쉬워야(tractable) 합니다. 우리는 이 섹션의 toy experiment에서 이 가정을 검증합니다.

우리는 관찰된 $D$-dimensional space에 "묻혀 있는(buried)" $d$-dimensional 기저 data의 toy case를 고려합니다 ($d < D$). 우리는 column-orthogonal인 projection matrix $P \in R^{D \times d}$ (즉, $P^\top P = I_{d \times d}$)를 사용하여 이 시나리오를 합성합니다. 이 matrix $P$는 무작위로 생성되고 고정됩니다. 관찰된 data는 $x = P\hat{x} \in R^D$ 이며, 여기서 기저 data는 $\hat{x} \in R^d$ 입니다. matrix $P$는 model에 알려지지 않았으며, 따라서 이는 model 입장에서 $D$-dimensional generation 문제입니다.

우리는 256-dim hidden units를 가진 5-layer ReLU MLP를 generator로 train하고 그 결과를 Fig. 2에 시각화합니다. 우리는 $P$를 사용하여 $D$-dim generated samples를 $d$-dim으로 다시 투영(projecting)하여 이러한 시각화를 얻습니다. 우리는 $d=2$일 때 $D \in {2, 8, 16, 512}$인 경우들을 조사합니다. 우리는 모두 v-loss를 사용하는 $x, \epsilon,$ 또는 v-prediction, 즉 Tab. 1(3)(a-c)를 연구합니다.

Fig. 2는 $D$가 증가할 때 오직 x-prediction만이 합리적인 결과를 만들어낼 수 있음을 보여줍니다. $\epsilon$-/v-prediction의 경우, models는 $D=16$에서 어려움을 겪고, 256-dim MLP가 under-complete 상태가 되는 $D=512$에서 치명적으로(catastrophically) 실패합니다.

특히, x-prediction은 model이 under-complete 상태일 때도 잘 작동할 수 있습니다. 여기서 256-dim MLP는 $D=512$-dim space의 information을 불가피하게 버려야 합니다(discard). 그러나 실제 data는 low-dimensional $d$-dim space에 있기 때문에, 이상적인 output이 암묵적으로 $d$-dim이므로 x-prediction은 여전히 잘 수행될 수 있습니다. 다음으로 보여주겠지만, ImageNet의 실제 data 경우에도 유사한 관찰 결과를 얻습니다.

논문 요약 정리 노트: On Prediction Outputs of Diffusion Models

1. Prediction Space와 Loss Space의 개념적 분리

• 독립성: Diffusion 모델 설계 시, 네트워크가 직접 출력하는 대상(Prediction Space: $x, \epsilon, v$)과 오차를 계산하는 기준(Loss Space: $x, \epsilon, v$)은 반드시 일치할 필요가 없음.
• 수학적 동치와 변환: $z_t = t x + (1 - t) \epsilon$ 및 $v = x - \epsilon$ 관계식에 의해, $x, \epsilon, v$ 중 하나만 예측(predict)하면 나머지 둘은 해석적으로 도출 가능함.
• 9가지 조합: 3개의 Prediction Space와 3개의 Loss Space를 조합하여 총 9가지의 Generator formulation이 가능함. Loss 공간을 바꾸는 것은 결국 특정 Prediction Space 관점에서는 시간 $t$에 따른 loss weighting을 다르게 주는 것과 동일한 효과를 지님.

2. Toy Experiment: 왜 High-dimensional에서 x-prediction인가? (핵심 기여)

• 가정: 실제 data는 low-dimensional manifold ($d$-dim)에 존재하지만, 우리가 관찰하는 공간은 high-dimensional space ($D$-dim)에 묻혀있음 ($d \ll D$).6
• 실험 세팅: $d=2$, $D \in \{2, 8, 16, 512\}$. Generator는 고작 256-dim hidden unit을 가진 5-layer MLP. (즉, $D=512$일 때 모델은 관측 차원보다 작은 under-complete 상태가 됨).
• $\epsilon$- / v-prediction의 실패: $D$가 16만 되어도 어려움을 겪으며, $D=512$에서는 완전히 실패(catastrophically fail)함. 원인은 noise $\epsilon$과 velocity $v$가 off-manifold 속성을 지니기 때문임. 이들은 $D$-dim 공간 전체에 넓게 퍼져 있으므로, 이를 정확히 예측하려면 전체 공간의 정보를 담을 수 있는 엄청난 모델 capacity가 필요함.
• x-prediction의 성공: 모델 capacity가 256-dim에 불과해 $D=512$ 환경에서 불가피하게 정보를 버려야 하는 상황에서도 정상적으로 작동함. 원인은 타겟인 $x$ 자체가 암묵적으로 low-dimensional ($d=2$) manifold에 존재하므로, 제한된 network capacity로도 핵심적인 데이터 구조를 충분히 표현하고 예측할 수 있기 때문임.

---

쉬운 설명 :

이 섹션을 한마디로 요약하면 **"쓸데없이 복잡한 노이즈 맞추기에 모델의 지능을 낭비하지 마라"**입니다.

아주 거대하고 복잡한 모래사장(High-dimensional space) 안에 작은 조개껍데기(Low-dimensional manifold에 있는 Clean data)가 하나 숨겨져 있다고 상상해 보세요.

기존의 $\epsilon$-prediction이나 v-prediction 방식은 모델에게 **"모래알들이 정확히 어떻게 흩뿌려져 있는지(노이즈)를 전부 기억해서 맞춰봐!"**라고 요구하는 것과 같습니다. 모래사장이 작을 때는 웬만큼 맞추겠지만, 모래사장이 엄청나게 커지면 모델의 머리(Capacity)로는 그 수많은 모래알의 위치를 다 담아낼 수 없어 완전히 붕괴해 버립니다. 이것이 장난감 실험(Toy experiment)에서 $D=512$일 때 기존 방식들이 처참하게 실패한 이유입니다.

반면 x-prediction 방식은 모델에게 **"복잡한 모래알은 무시하고, 그 안에 있는 조개껍데기의 모양(Clean data)만 찾아내!"**라고 지시하는 것입니다. 조개껍데기 자체는 모래사장 크기와 상관없이 그 모양이 단순하고 규칙적입니다. 따라서 모델의 머리가 조금 나빠도(Under-complete 상태여도) 모래를 털어내고 조개껍데기의 형태를 유추하는 것은 아주 쉽게 해낼 수 있습니다.

저자는 이 챕터의 수학적 증명과 실험을 통해, 출력 대상(Prediction)을 노이즈가 아닌 진짜 데이터($x$)로 바꾸기만 해도 차원의 저주를 극복하고 모델의 한계를 돌파할 수 있음을 아주 명쾌하게 보여주고 있습니다.

 

 

 

"Just Image Transformers" for Diffusion

앞선 분석을 바탕으로, 우리는 pixels에서 작동하는 일반적인 Vision Transformers (ViT)가 단순히 x-prediction을 사용하는 것만으로 놀라울 정도로 잘 작동할 수 있음을 보여줍니다.

4.1. Just Image Transformers

ViT의 핵심 idea는 "Transformer on Patches (ToP)"입니다. 우리의 architecture design은 이 철학을 따릅니다.

공식적으로, $H \times W \times C$-dim image data ($C=3$)를 고려해 봅시다. 모든 $x, \epsilon, v$ 및 $z_t$는 이 동일한 dimensionality를 공유합니다. image가 주어지면, 우리는 이를 $p \times p$ 크기의 겹치지 않는(non-overlapping) patches로 나누어 $\frac{H}{p} \times \frac{W}{p}$ 길이의 sequence를 만듭니다. 각 patch는 $p \times p \times 3$-dim vector입니다. 이 sequence는 linear embedding projection에 의해 처리되고, positional embedding과 더해진 후, Transformer blocks 스택에 의해 매핑됩니다. output layer는 각 token을 다시 $p \times p \times 3$-dim patch로 투영하는(projects) linear predictor입니다. Fig. 3을 참조하십시오.

일반적인 관행(standard practice)에 따라, architecture는 시간 $t$와 주어진 class label을 조건으로 합니다(conditioned). 우리는 conditioning을 위해 adaLN-Zero를 사용하며 나중에 다른 옵션들을 논의할 것입니다. 개념적으로 이 architecture는 pixels의 patches에 직접 적용된 Diffusion Transformer (DiT)에 해당합니다.

전체 architecture는 우리가 JiT라고 부르는 "Just image Transformers"에 지나지 않습니다. 예를 들어, 우리는 $256 \times 256$ images에서 JiT/16(즉, patch size $p=16$)을 조사하고, $512 \times 512$ images에서 JiT/32 ($p=32$)를 조사합니다. 이러한 설정은 각각 patch당 768 ($16 \times 16 \times 3$) 및 3072 ($32 \times 32 \times 3$)의 dimensionality를 초래합니다. 이렇게 high-dimensional patches는 x-prediction에 의해 처리될 수 있습니다.

4.2. What to Predict by the Network?

우리는 9가지 가능한 loss combinations를 요약했습니다. 우리는 token당 768-dim의 hidden size를 가진 "Base" model (JiT-B)을 train합니다. 우리는 Tab. 2(a)에서 $256 \times 256$ resolution의 JiT-B/16을 연구합니다. 참고(reference)를 위해, 우리는 Tab. 2(b)에서 $64 \times 64$의 JiT-B/4 (즉, $p=4$)를 검토합니다. 두 설정 모두에서 sequence length는 동일합니다 ($16 \times 16$).

우리는 다음과 같은 관찰 결과를 도출합니다:

x-prediction is critical. JiT-B/16을 사용한 Tab. 2(a)에서 오직 x-prediction만이 우수한 성능을 보이며, 세 가지 losses 모두에서 잘 작동합니다. 여기서 patch는 768-d ($16 \times 16 \times 3$)이며, 이는 JiT-B의 768이라는 hidden size와 일치합니다. 이것이 "대략 충분해" 보일 수 있지만, 실제 환경에서 models는 positional embeddings를 처리하는 등 추가적인 capacity를 요구할 수 있습니다. $\epsilon$-/v-prediction의 경우, model은 noised quantities를 분리하고 유지할 충분한 capacity가 없습니다. 이러한 관찰은 toy case (Fig. 2)에서의 관찰과 유사합니다.

비교를 위해, 우리는 $64 \times 64$ resolution에서 JiT-B/4를 검토합니다 (Tab. 2(b)). 여기서는 모든 경우들이 합리적으로 잘 수행됩니다. 9가지 combinations 간의 정확도(accuracy) 격차는 결정적(decisive)이지 않고 미미합니다(marginal). dimensionality는 patch당 48 ($4 \times 4 \times 3$)로, JiT-B의 hidden size인 768을 훨씬 밑돌며, 이는 왜 모든 combinations가 합리적으로 잘 작동하는지를 설명합니다. 우리는 많은 이전 latent diffusion models가 이와 유사하게 작은 input dimensionality를 가지며 따라서 여기서 논의하는 문제에 노출되지 않았음을 지적합니다.

Loss weighting is not sufficient. 우리의 연구가 relevant factors의 combinations를 나열한 첫 번째 연구는 아닙니다. 선행 연구에서는 loss weighting과 network predictions의 combinations를 탐구했습니다. 그들의 experiments는 U-net을 사용하여 low-dimensional CIFAR-10 dataset에서 수행되었습니다. 그들의 관찰 결과는 ImageNet $64 \times 64$에서의 우리 결과와 더 가까웠습니다.

그러나 ImageNet $256 \times 256$에 대한 Tab. 2(a)는 loss weighting이 전부가 아님을 시사합니다. 한편으로, $\epsilon$- 및 v-prediction 모두 Tab. 2(a)에서 loss space와 무관하게 치명적으로(catastrophically) 실패하며, 이는 (논의한 바와 같이) 다른 loss spaces에서 다른 effective weightings에 해당합니다. 다른 한편으로, x-prediction은 세 가지 loss spaces 모두에서 작동합니다: v-loss에 의해 유도된 loss weighting이 선호되지만 결정적인(critical) 것은 아닙니다.

Noise-level shift is not sufficient. 이전 연구들은 pixel-based diffusion의 고해상도(high-resolution) 처리를 위해 noise level을 증가시키는 것이 유용하다고 제안했습니다. 우리는 Tab. 3에서 JiT-B/16으로 이를 검토합니다. $t$를 sampling하기 위해 logit-normal distribution을 사용하므로 (appendix 참조), noise level은 이 distribution의 parameter $\mu$를 변경하여 shift될 수 있습니다. 직관적으로 $\mu$를 음수 방향으로 shift하면 $t$가 작아지고 따라서 noise level이 증가합니다 (Eq. (1)).

Tab. 3은 model이 이미 꽤 잘 수행될 때(여기서는 x-pred), 적절히 높은 noise가 유익하다는 것을 보여주며, 이는 이전 관찰과 일치합니다. 그러나 noise level을 조정하는 것만으로는 $\epsilon$- 또는 v-prediction을 구제(remedy)할 수 없습니다. 그들의 실패는 high-dimensional information을 전파(propagate)하지 못하는 본질적인 무능력에서 비롯됩니다.

참고로(As a side note), Tab. 3에 따라 우리는 ImageNet $256 \times 256$에 대한 다른 experiments에서 $\mu = -0.8$로 설정했습니다.

Increasing hidden units is not necessary. capacity는 network width (즉, hidden units의 수)에 의해 제한될 수 있으므로, 이를 늘리는 것이 자연스러운 idea입니다. 그러나 observed dimension이 매우 높을 때 이 해결책(remedy)은 원칙적이지도 실행 가능하지도 않습니다. 우리는 x-prediction의 경우 이것이 불필요함을 보여줍니다.

다음 섹션의 Tab. 5와 Tab. 6에서, 우리는 비례적으로 큰 patch size인 $p=32$ 또는 $p=64$를 사용하여 resolution 512에서 JiT/32, resolution 1024에서 JiT/64의 결과를 보여줍니다. 이는 patch당 3072-dim (즉, $32 \times 32 \times 3$) 또는 12288-dim에 해당하며, B, L, H models의 hidden size보다 실질적으로 훨씬 큽니다. 그럼에도 불구하고 x-prediction은 잘 작동합니다. 사실, noise를 비례적으로 scaling하는 것(예: resolution 512 및 1024에서 2배 및 4배로; appendix 참조) 외에는 아무런 수정 없이 작동합니다.

이 증거는 다른 많은 neural network applications에서 그러하듯, network design이 observed dimensionality와 크게 분리(decoupled)될 수 있음을 시사합니다. hidden units의 수를 늘리는 것은 유익할 수 있지만 (deep learning에서 널리 관찰되듯), 결정적(decisive)이지는 않습니다.

Bottleneck can be beneficial. 더욱 놀랍게도, 우리는 반대로 network의 dimensionality를 줄이는 bottleneck을 도입하는 것이 유익할 수 있음을 발견했습니다.

구체적으로, 우리는 linear patch embedding layer를 한 쌍의 bottleneck (여전히 linear인) layers로 교체하여 low-rank linear layer로 변환합니다. 첫 번째 layer는 dimension을 $d'$로 줄이고, 두 번째 layer는 이를 Transformer의 hidden size로 확장합니다. 두 layers 모두 linear이며 low-rank reparameterization 역할을 합니다.

Fig. 4는 raw patch당 768-d인 JiT-B/16을 사용하여 FID 대 bottleneck dimension $d'$를 plot합니다. bottleneck dimension을 심지어 16-d만큼 작게 줄여도 치명적인 실패를 일으키지 않습니다. 사실, 넓은 범위(32에서 512)에 걸친 bottleneck dimension은 최대 $\sim 1.3$ FID라는 꽤 큰 차이로 quality를 향상시킬 수 있습니다.

representation learning의 더 넓은 관점에서 볼 때, 이 관찰 결과는 전혀 예상치 못한 것은 아닙니다. bottleneck designs는 종종 본질적으로 low-dimensional representations의 학습을 장려하기 위해 도입됩니다.

4.3. Our Algorithm

우리의 최종 algorithm은 x-prediction과 v-loss를 채택하며, 이는 Tab. 1(3)(a)에 해당합니다. 공식적으로 우리는 다음을 최적화합니다:

$$L = E_{t,x,\epsilon} \lVert v_\theta(z_t, t) - v \rVert^2, \quad (6)$$

여기서: $v_\theta(z_t, t) = (net_\theta(z_t, t) - z_t) / (1 - t)$ 입니다.

Alg. 1은 training step의 pseudo-code를 보여주고, Alg. 2는 sampling step (Euler solver; Heun 또는 다른 solvers로 확장 가능)의 pseudo-code를 보여줍니다. 간결함을 위해 class conditioning과 CFG는 생략되었지만, 둘 다 일반적인 관행(standard practice)을 따릅니다. $1/(1-t)$에서 zero division을 방지하기 위해, 우리는 이 나눗셈을 계산할 때마다 분모를 (기본적으로 0.05로) clip합니다.

4.4. "Just Advanced" Transformers

general-purpose Transformer의 강점은 부분적으로, 그 design이 특정 task와 분리(decoupled)될 때 다른 applications에서 개발된 architectural advances의 혜택을 받을 수 있다는 점에 있습니다. 이 속성은 task-agnostic Transformer로 diffusion을 공식화하는 이점의 뒷받침이 됩니다.

우리는 원래 language models를 위해 개발된 널리 쓰이는 general-purpose improvements(SwiGLU, RMSNorm, RoPE, qk-norm)를 통합합니다. 우리는 또한 in-context class conditioning을 탐구합니다. sequence에 하나의 class token을 추가하는 원래 ViT와 달리, 우리는 여러 개의 이러한 tokens를 추가합니다 (기본적으로 32개; appendix 참조). Tab. 4는 이러한 components의 효과를 보고합니다.

 

논문 요약 정리 노트: "Just Image Transformers" for Diffusion

1. JiT (Just image Transformers) 아키텍처 제안

• 구조적 단순함: 어떠한 계층적(hierarchical) 구조나 복잡한 변형 없이, 원본 픽셀(raw pixels)에 순수 ViT를 직접 적용하는 "Transformer on Patches (ToP)" 철학 채택.
• Large Patch Size의 도입: 고해상도 이미지(256, 512, 1024)를 처리하기 위해 매우 큰 patch size(16, 32, 64)를 사용. 이로 인해 token당 dimensionality가 극단적으로 높아짐 (최대 12288-dim).

2. Prediction Target에 따른 실험적 발견 (가장 중요한 기여)

• High-dimensional 환경에서 x-prediction의 필수성: patch dimension이 모델의 hidden size와 비슷하거나 클 때(예: JiT-B/16, 768-d), $\epsilon$-prediction과 v-prediction은 완전히 붕괴(catastrophic failure)함. 오직 x-prediction만이 정상적으로 작동함.
• Latent Diffusion이 성공했던 이유 규명: patch dimension이 작은 경우(예: JiT-B/4, 48-d)에는 모든 prediction 방식이 잘 작동함. 기존 LDM들은 input dimensionality가 낮았기 때문에 $\epsilon$-prediction으로도 이 문제를 겪지 않았던 것임.
• Loss Weighting & Noise Shift의 한계: loss space를 변경하거나(weighting 효과) noise level을 높이는 것은 보조적인 성능 향상에는 도움이 되나, $\epsilon$-/v-prediction의 본질적인 high-dimensional capacity 부족 문제를 해결(remedy)하지는 못함.

3. Capacity Decoupling과 Bottleneck 효과

• 차원과 Hidden Size의 분리: 12288-dim의 거대한 patch를 입력받더라도, 모델의 hidden size를 무리하게 키울 필요 없이 x-prediction만으로 안정적인 처리가 가능함.
• Bottleneck의 성능 향상 효과 (Counter-intuitive): Linear patch embedding layer를 Low-rank bottleneck layer(예: 768-d $\rightarrow$ 16-d $\rightarrow$ hidden size)로 교체하여 의도적으로 네트워크 내 dimensionality를 줄였을 때, 오히려 FID 성능이 최대 1.3가량 크게 향상됨. 이는 manifold learning의 특성(유용한 low-dimensional 정보만 통과시킴)이 생성 모델에서도 완벽히 통용됨을 증명.

4. 최적의 알고리즘 및 확장성

• 설정: x-prediction을 target으로, v-loss를 목적 함수로 사용. (Tab. 1의 3-a 조합)
• General-purpose 강점: 모델 구조가 특정 task에 종속되지 않은 순수 Transformer이므로, 언어 모델(LLM) 등에서 발전된 최신 기법(SwiGLU, RMSNorm, RoPE, 다중 in-context class tokens)들을 아무런 충돌 없이 바로 통합하여 성능을 극대화할 수 있음.

---

쉬운 설명 :

이 섹션은 **"어떻게 평범한 트랜스포머(ViT)가 무식하게 큰 원본 픽셀 덩어리를 소화해 냈는가?"**에 대한 해답을 줍니다.

고해상도 이미지를 트랜스포머에 넣으려면 이미지를 아주 큼직큼직한 조각(Large patch)으로 썰어 넣어야 합니다. 조각이 크다 보니 조각 하나에 담긴 정보량(차원, Dimension)이 엄청나게 방대해집니다.

여기서 기존 방식대로 모델에게 **"이 방대한 픽셀 덩어리 속에 낀 노이즈를 예측해!"**라고 하면, 모델은 그 무작위하고 복잡한 노이즈의 패턴을 전부 외울 머리(Capacity)가 안 돼서 완전히 고장 나버립니다. (기존 모델들이 이미지를 압축해서 작게 만든 뒤에야 노이즈를 예측했던 이유가 바로 이것입니다.)

하지만 저자들의 방식대로 **"노이즈 말고, 픽셀 덩어리에서 진짜 그림(x-prediction)만 찾아내!"**라고 목표를 바꿔주면 상황이 180도 달라집니다. 노이즈와 달리 진짜 그림은 특정한 규칙(Manifold)이 있어서, 모델의 머리 크기(Hidden size)가 입력되는 픽셀 덩어리 크기보다 훨씬 작아도 아주 여유롭게 그림을 찾아냅니다.

심지어 한술 더 떠서, 모델에 들어가는 정보의 통로를 억지로 확 좁혀버렸더니(Bottleneck 도입) 성능이 오히려 더 좋아졌습니다! 통로가 좁아지니 쓸데없는 노이즈는 알아서 걸러지고 진짜 중요한 그림의 핵심 정보만 쏙쏙 뽑아내게 되었기 때문입니다. 결국 **"예측 대상을 진짜 그림으로 바꾸는 순간, 기존의 차원 제약이나 모델 크기의 한계를 완전히 박살 낼 수 있다"**는 것이 이 섹션의 통쾌한 결론입니다.