[Study] DiT: Diffusion에 Transformer를 쓰는 방식

DiT: Diffusion에 Transformer를 쓰는 방식

DiT는 Diffusion Transformer의 줄임말이다. 말 그대로 디퓨전 모델에서 U-Net이 하던 역할을 Transformer로 바꾼 구조다.

Stable Diffusion 1.5, 2.0, SDXL은 기본적으로 U-Net을 쓴다. 노이즈가 섞인 이미지를 넣으면 U-Net이 그 안의 노이즈를 예측하고, 그 예측을 바탕으로 이미지를 조금씩 깨끗하게 만든다.

DiT는 이 부분을 다르게 본다. 이미지나 잠재 공간을 작은 조각으로 나누고, 그 조각들을 Transformer가 처리하게 만든다. 구조만 놓고 보면 Vision Transformer에 디퓨전 학습을 붙인 형태에 가깝다.

이 글에서는 DiT가 왜 나왔고, U-Net과 무엇이 다르며, 실제로 내부에서 어떤 흐름으로 동작하는지 정리한다. 수식보다는 구조와 흐름 위주로 본다.

---

1. 출발점 — U-Net이 하던 일

디퓨전 모델은 노이즈에서 시작해 이미지를 만든다. 완전한 노이즈를 한 번에 이미지로 바꾸는 것이 아니라, 여러 단계에 걸쳐 노이즈를 조금씩 제거한다.

각 단계에서 모델이 해야 하는 일은 단순하다. 현재 이미지에 들어있는 노이즈를 예측하는 것이다.

노이즈가 섞인 이미지  →  모델  →  예측한 노이즈

Stable Diffusion 계열에서는 이 역할을 U-Net이 맡았다. U-Net은 입력과 출력의 크기를 맞추기 좋고, 이미지의 지역적인 패턴을 잘 잡는다. 그래서 디퓨전 모델에서 오랫동안 표준 구조처럼 쓰였다.

DiT의 질문은 여기서 시작한다. 노이즈를 예측하는 네트워크가 꼭 U-Net이어야 할까? Transformer로도 같은 일을 할 수 있지 않을까?

---

2. DiT가 바꾼 부분

DiT가 바꾸는 것은 디퓨전 전체가 아니다. 노이즈를 넣고 빼는 방식, timestep을 쓰는 방식, 조건을 넣는 방식의 큰 틀은 그대로다.

바뀌는 부분은 노이즈를 예측하는 신경망이다.

구조	노이즈 예측 담당	특징
기존 Stable Diffusion	U-Net	이미지 구조에 맞춘 CNN 기반 구조
DiT	Transformer	이미지를 토큰처럼 보고 attention으로 처리

즉, DiT는 디퓨전의 원리를 새로 만든 모델이라기보다, 기존 디퓨전 모델의 백본을 Transformer로 교체한 모델이라고 보는 편이 이해하기 쉽다.

---

3. 잠재 공간에서 동작한다

DiT는 보통 픽셀 이미지를 그대로 다루지 않는다. Stable Diffusion처럼 VAE가 압축한 잠재 공간(latent space)에서 동작한다.

이미지  →  VAE Encoder  →  잠재 공간  →  DiT  →  VAE Decoder  →  이미지

예를 들어 256×256 이미지를 그대로 보면 픽셀 수가 많다. 하지만 VAE를 거치면 훨씬 작은 잠재 텐서로 압축된다. DiT는 이 작은 텐서 위에서 노이즈를 예측한다.

이렇게 하는 이유는 계산량 때문이다. Transformer는 토큰 수가 많아지면 비용이 빠르게 커진다. 그래서 픽셀을 직접 다루기보다, 압축된 잠재 공간을 쓰는 편이 훨씬 현실적이다.

---

4. 이미지를 토큰으로 바꾸는 과정

Transformer는 원래 토큰의 나열을 입력으로 받는다. 문장이라면 단어 토큰을 넣으면 된다. 그런데 이미지는 2차원 텐서다. 그래서 DiT는 이미지를 작은 패치로 나눈다.

잠재 텐서  →  작은 패치들  →  토큰 시퀀스

예를 들어 잠재 텐서가 32×32 크기라면, 이것을 2×2 또는 4×4 같은 작은 조각으로 자른다. 각 조각을 펼쳐서 벡터로 만들고, Transformer가 읽을 수 있는 차원으로 변환한다.

이 과정을 보통 patchify라고 부른다. Vision Transformer에서 이미지를 패치로 나누는 것과 같은 아이디어다.

패치 크기의 의미

패치를 작게 자르면 토큰 수가 많아진다. 더 세밀하게 볼 수 있지만 계산량이 늘어난다. 반대로 패치를 크게 자르면 토큰 수는 줄어들지만 세부 표현은 약해질 수 있다.

패치 크기	특징
작은 패치	세부 정보를 더 잘 다룸. 계산량은 큼.
큰 패치	계산은 빠름. 대신 세밀한 표현이 줄어들 수 있음.

DiT 이름 뒤에 붙는 /2, /4, /8 같은 표기는 이 패치 크기를 뜻한다. 예를 들어 DiT-XL/2는 XL 크기의 모델에 패치 크기 2를 쓴다는 의미다.

---

5. 위치 정보가 필요한 이유

패치로 나눈 뒤에는 각 패치가 어디에 있던 조각인지 알려줘야 한다. Transformer는 기본적으로 입력 토큰의 위치를 스스로 알지 못하기 때문이다.

그래서 각 패치 토큰에 위치 정보를 더한다. 이것이 position embedding이다.

위치 정보가 없으면 왼쪽 위 패치와 오른쪽 아래 패치를 구분하기 어렵다. 이미지 생성에서는 위치가 중요하다. 얼굴의 눈, 코, 입처럼 위치 관계가 맞아야 하는 요소들이 많기 때문이다.

DiT는 패치 토큰에 위치 정보를 더한 뒤, 그 토큰들을 Transformer 블록에 넣는다.

---

6. DiT 블록이 하는 일

DiT의 중심은 Transformer 블록이다. 각 블록은 토큰들 사이의 관계를 보고, 현재 단계에서 어떤 노이즈를 제거해야 하는지 판단한다.

크게 보면 한 블록은 다음 두 부분으로 구성된다.

1. Self-Attention — 각 패치가 다른 패치들을 참고한다. 멀리 떨어진 영역끼리도 바로 관계를 볼 수 있다.
2. Feed-Forward Network — attention 결과를 각 토큰별로 다시 가공한다.

U-Net은 convolution을 통해 주변 영역부터 점점 넓게 본다. 반면 Transformer는 attention을 통해 처음부터 모든 패치를 서로 연결해서 볼 수 있다. 이 점이 DiT의 중요한 차이다.

예를 들어 이미지의 왼쪽에 있는 물체와 오른쪽에 생기는 그림자처럼, 멀리 떨어져 있지만 서로 맞아야 하는 정보가 있다. Transformer는 이런 장거리 관계를 다루기 좋다.

---

7. timestep과 조건은 어떻게 넣나

디퓨전 모델은 현재 몇 번째 노이즈 제거 단계인지 알아야 한다. 초반에는 노이즈가 많고, 후반에는 세부 디테일을 다듬기 때문이다. 이 정보를 timestep이라고 한다.

조건부 생성에서는 여기에 클래스나 텍스트 조건도 들어간다. 예를 들어 ImageNet 클래스 조건부 모델이라면 “고양이”, “자동차” 같은 클래스 정보가 들어간다. 텍스트-이미지 모델이라면 프롬프트 임베딩이 조건이 된다.

DiT는 이 조건 정보를 Transformer 블록에 전달해야 한다. 방법은 여러 가지가 있지만, DiT 논문에서 특히 중요하게 다룬 것은 adaLN-Zero라는 방식이다.

---

8. adaLN-Zero를 쉽게 보면

adaLN-Zero는 이름이 어렵지만, 핵심은 단순하다. 조건 정보를 attention에 직접 붙이는 대신, Transformer 블록의 정규화 과정에 반영한다.

일반적인 LayerNorm은 토큰의 값을 정리해주는 역할을 한다. adaLN은 이 LayerNorm의 조절값을 조건에 따라 바꾼다.

조건 정보  →  LayerNorm 조절값  →  Transformer 블록에 반영

즉, “지금은 몇 번째 timestep인지”, “어떤 클래스를 만들고 있는지”에 따라 블록의 동작을 조금씩 바꾸는 방식이다.

Zero가 붙는 이유

adaLN-Zero에서 Zero는 초기 상태와 관련이 있다. 학습 시작 시에는 각 블록이 거의 아무것도 하지 않는 상태, 즉 입력을 그대로 통과시키는 상태에 가깝게 만든다.

이렇게 하면 깊은 Transformer를 처음부터 무리하게 학습시키는 것보다 안정적이다. 처음에는 조용히 시작하고, 학습이 진행되면서 각 블록이 점점 역할을 갖게 된다.

이 방식은 계산량을 크게 늘리지 않으면서 조건을 잘 넣을 수 있다는 장점이 있다. DiT에서 adaLN-Zero가 중요한 이유도 여기에 있다.

---

9. U-Net과 비교하면 무엇이 다른가

U-Net과 DiT의 차이는 단순히 “CNN이냐 Transformer냐”만은 아니다. 이미지를 바라보는 방식 자체가 다르다.

항목	U-Net	DiT
기본 단위	특징맵	패치 토큰
주요 연산	Convolution	Self-Attention
강점	지역 패턴, 이미지 구조에 강함	멀리 떨어진 영역 관계를 보기 좋음
확장 방식	구조 설계가 비교적 복잡함	depth, width, 토큰 수를 키우는 방식이 명확함

U-Net은 이미지에 잘 맞는 구조다. 작은 데이터나 제한된 계산량에서는 여전히 강하다. 반면 DiT는 모델과 데이터, 계산량을 크게 키울 때 장점이 잘 드러난다.

Transformer 계열 모델은 크기를 키웠을 때 성능이 비교적 예측 가능하게 좋아지는 경향이 있다. DiT 논문에서도 이 점이 중요하게 다뤄졌다.

---

10. DiT가 중요한 이유

DiT의 의미는 “U-Net을 Transformer로 바꿔도 된다”에서 끝나지 않는다. 더 중요한 점은 디퓨전 모델도 Transformer의 scaling 방식을 따라갈 수 있다는 것을 보여줬다는 점이다.

모델을 크게 만들고, 패치를 작게 하고, 학습을 오래 하면 성능이 꾸준히 좋아지는 흐름이 관찰됐다. 이건 대형 모델을 설계할 때 매우 중요하다. 어떤 방향으로 키워야 할지 판단하기 쉬워지기 때문이다.

이후의 여러 이미지 생성 모델들은 DiT의 아이디어를 직접 또는 간접적으로 가져간다. Stable Diffusion 3처럼 Transformer 기반 구조를 쓰는 모델도 이 흐름 위에 있다.

정리하면 DiT는 U-Net을 완전히 대체한다기보다, 대형 디퓨전 모델을 만들 때 Transformer가 더 자연스러운 선택지가 될 수 있음을 보여준 모델이다.

---

11. 생성 과정

학습이 끝난 DiT로 이미지를 만들 때의 흐름은 기존 디퓨전 모델과 크게 다르지 않다.

1. 잠재 노이즈 생성 — 먼저 잠재 공간에서 랜덤 노이즈를 만든다.
2. 패치 토큰 변환 — 현재 잠재 텐서를 작은 패치 토큰들로 바꾼다.
3. DiT로 노이즈 예측 — Transformer가 각 단계에서 제거할 노이즈를 예측한다.
4. 노이즈 제거 반복 — 여러 스텝에 걸쳐 잠재 텐서를 점점 깨끗하게 만든다.
5. VAE 디코딩 — 최종 잠재를 이미지로 복원한다.

겉으로 보면 Stable Diffusion과 비슷하다. 차이는 노이즈 예측을 담당하는 내부 네트워크가 U-Net이 아니라 Transformer라는 점이다.

---

12. 한계와 주의할 점

DiT가 항상 더 좋은 선택이라는 뜻은 아니다. Transformer는 계산량을 많이 쓴다. 특히 토큰 수가 많아지면 attention 비용이 빠르게 늘어난다.

• 계산 비용이 크다. 패치를 작게 자를수록 품질은 좋아질 수 있지만, attention 비용도 커진다.
• 충분한 학습이 필요하다. Transformer는 규모가 커질수록 강하지만, 그만큼 데이터와 학습 시간이 필요하다.
• 작은 모델에서는 U-Net이 더 실용적일 수 있다. 제한된 환경에서는 검증된 U-Net 구조가 여전히 효율적이다.

따라서 DiT는 “무조건 U-Net보다 좋다”라기보다, 대형 모델을 만들고 확장하는 방향에서 강한 구조라고 보는 편이 맞다.

---

13. 정리

DiT는 디퓨전 모델의 노이즈 예측 네트워크를 U-Net에서 Transformer로 바꾼 구조다. 이미지를 잠재 공간에서 작은 패치로 나누고, 그 패치들을 토큰처럼 처리한다.

핵심은 세 가지다.

• 디퓨전의 기본 원리는 그대로 두고, 백본만 Transformer로 바꾼다.
• 잠재 텐서를 패치 토큰으로 바꿔 Transformer가 처리한다.
• 모델 크기와 계산량을 키웠을 때 성능이 잘 따라오는 구조를 보여줬다.

U-Net은 이미지 생성에서 오래 쓰인 안정적인 구조다. DiT는 그 다음 단계로, 더 큰 모델과 더 많은 계산량을 전제로 할 때 Transformer가 디퓨전 모델의 중심 구조가 될 수 있음을 보여준 사례다.