영상처리 기술, 그 이야기

Diffusion Transformer(DiT)는 확산 모델(Diffusion Model)의 이미지 생성 과정에 Transformer 구조를 결합한 생성 모델입니다. 기존 확산 모델에서 널리 사용되던 U-Net 대신 Transformer를 핵심 네트워크로 사용하여 이미지 생성과 변환 작업을 수행합니다. 대표적인 연구 모델로 Meta의 DiT가 있으며, 이후 다양한 고해상도 이미지와 영상 생성 모델의 설계에 영향을 주었습니다.

확산 모델은 처음에는 이미지에 노이즈를 점진적으로 추가하고, 학습 과정에서는 이 노이즈를 다시 제거하는 방법을 학습합니다. 실제 생성 단계에서는 무작위 노이즈에서 시작해 여러 번의 denoising 과정을 거쳐 최종 이미지를 만들어냅니다. DiT에서는 이 과정에서 이미지를 작은 패치(patch) 단위의 토큰으로 변환하고 Transformer의 Self-Attention을 이용해 토큰 간 관계를 학습합니다.

DiT의 핵심 구조

일반적인 DiT는 크게 Patchify → Transformer Blocks → Unpatchify 구조로 이해할 수 있습니다.

먼저 입력 이미지 또는 잠재 공간(latent)의 feature를 일정한 크기의 패치로 나눕니다. 각 패치는 Transformer가 처리할 수 있는 토큰으로 변환됩니다. 이후 여러 개의 Transformer Block을 통과하면서 Self-Attention을 통해 이미지 전체 영역의 관계를 학습합니다. 마지막으로 Transformer의 출력을 다시 이미지 형태로 변환하여 노이즈 제거에 필요한 정보를 생성합니다.

특히 확산 과정에서 사용되는 시간 단계(timestep)와 텍스트 프롬프트 등의 조건 정보도 Transformer에 전달됩니다. 따라서 "고양이가 자동차 위에 있다"와 같은 텍스트 조건에 맞춰 어떤 영역을 어떻게 생성할지 결정할 수 있습니다.

이미지 생성과 변환에서의 활용

DiT는 단순한 Text-to-Image뿐만 아니라 Image-to-Image, Inpainting, Outpainting, 스타일 변환, 이미지 편집 등에도 활용할 수 있습니다. 예를 들어 기존 이미지와 텍스트 조건을 입력하면 원본 이미지의 구조를 유지하면서 특정 부분의 스타일이나 내용을 변경할 수 있습니다.

또한 Transformer는 Self-Attention을 통해 이미지의 먼 영역 사이의 관계를 직접 모델링할 수 있다는 장점이 있습니다. 따라서 복잡한 전역 구조나 객체 간 관계를 학습하는 데 유리합니다.

DiT의 장점과 한계

DiT의 가장 큰 장점은 Transformer의 확장성(scaling)입니다. 모델 크기와 학습 데이터, 연산량을 증가시키면서 성능을 비교적 체계적으로 확장할 수 있습니다. 또한 Transformer 기반 AI 모델과 동일한 아키텍처 철학을 적용할 수 있어 멀티모달 모델이나 영상 생성 모델로 확장하기에도 유리합니다.

반면 Self-Attention은 토큰 수가 증가하면 연산량과 메모리 사용량이 크게 증가합니다. 따라서 고해상도 이미지를 직접 처리하기보다는 VAE 등을 이용해 이미지를 latent 공간으로 압축한 뒤 DiT가 latent를 처리하는 방식이 효율적입니다.

Diffusion Transformer는 확산 모델의 생성 방식과 Transformer의 전역적인 특징 학습 능력을 결합한 구조입니다. 기존 U-Net 중심의 확산 모델 구조를 Transformer 기반으로 확장하면서 이미지의 장거리 관계를 효과적으로 학습할 수 있으며, 이미지 생성, 이미지 변환, 인페인팅, 스타일 변환뿐만 아니라 향후 영상 생성과 멀티모달 생성 모델에도 활용될 수 있습니다.

특히 생성형 AI가 고해상도 이미지와 영상, 3D 콘텐츠 등으로 발전하면서 DiT와 같은 Transformer 기반 확산 구조는 중요한 기술적 방향으로 자리 잡고 있습니다.