확산 트랜스포머 쉽게 이해하기: ViT와 DiT, MMDiT의 구조와 조건 부여 방식 단계별 정리

핵심 메시지

• 이미지에 트랜스포머를 쓰려면 '연속적인 이미지를 어떻게 토큰으로 자를 것인가'와 '텍스트 프롬프트를 어떻게 반영할 것인가' 두 문제를 풀어야 한다.
• ViT는 픽셀 대신 16×16 패치를 토큰으로 삼아 계산량을 수십억 회에서 6만 5천여 회 수준으로 낮췄다.
• 확산 모델은 원본 픽셀이 아니라 VAE로 압축한 잠재 공간에서 노이즈를 걷어내며, 이때 시간 단계 정보를 모델에 명시적으로 넣어 줘야 한다.
• DiT 논문은 조건 정보를 넣는 네 가지 방법을 비교했고, 그중 특징을 증폭·억제하는 적응형 레이어 정규화(adaLN) 방식이 비용 대비 효과가 좋았다.
• MMDiT는 이미지와 텍스트를 대등하게 다뤄 하나의 시퀀스에서 서로 주의를 주고받게 함으로써, 이미지 편집처럼 맥락에 따라 의미가 달라지는 지시를 처리한다.

쉽게 이해하기

트랜스포머는 원래 기계 번역을 위해 만들어졌고, 그래서 원문을 양방향으로 읽는 인코더와 번역문을 한 토큰씩 만들어 내는 디코더로 나뉜다. 그런데 이미지는 왼쪽에서 오른쪽으로 읽는 대상이 아니라 전체를 한꺼번에 파악하는 대상이다. 그래서 대부분의 비전 트랜스포머는 디코더를 떼고 양방향 어텐션을 쓰는 인코더만 빌려 쓴다.

문제는 계산량이다. 셀프 어텐션은 모든 픽셀이 모든 픽셀을 보게 하므로 256×256짜리 작은 이미지에서도 수십억 회 규모의 연산이 필요하다. ViT는 이미지를 16×16 픽셀 패치로 잘라 격자로 만들고, 이 패치를 단어처럼 다룬다. '이미지는 16×16 단어의 가치가 있다'는 표현대로, 패치 단위로 바꾸는 것만으로 어텐션 연산이 6만 5천여 회 수준까지 줄어든다. 여기에 위치 임베딩을 더하고, 분류에는 BERT에서 빌려 온 CLS 토큰을 붙여 이미지 전체의 요약으로 삼는다.

이미지 생성으로 넘어가면 구조가 한 번 더 바뀐다. 초기 확산 모델의 뼈대였던 합성곱 기반 U-Net을 트랜스포머로 갈아 끼운 것이 DiT다. 다만 원본 픽셀에서 확산을 돌리면 비용이 크기 때문에, 별도의 VAE로 이미지를 작은 3차원 텐서로 압축한 잠재 공간에서 정방향·역방향 확산을 모두 수행하고, 마지막에 디코더로 다시 픽셀로 되돌린다. 모델은 매 단계 부분적으로만 노이즈를 걷어내므로 지금이 몇 번째 단계인지를 알려 주는 시간 정보가 반드시 입력에 포함돼야 한다.

핵심 질문은 '시간과 클래스 같은 조건 정보를 어떻게 넣느냐'다. 조건을 그냥 토큰으로 붙이면 이미지 토큰과 주의력을 두고 경쟁하다 무시될 수 있고, 별도의 교차 어텐션을 두면 확실하지만 연산량이 가장 많이 늘어난다. DiT가 채택한 방식은 2017년 FiLM 연구에서 온 특징 변조다. 조건 정보로부터 크기(감마)와 이동(베타) 값을 계산해 특징 채널별로 강조하거나 억제하는데, 이는 레이어 정규화와 닮아 적응형 레이어 정규화로 불린다. 노이즈가 심한 초반에는 큰 구조에, 후반에는 세부에 집중하도록 자연스럽게 유도된다.

마지막 조각은 잔차 블록 앞에 붙는 스케일 계수다. 학습 초기에 이 값을 0에 가깝게 두면 상위 계층이 사실상 꺼져 네트워크가 얕아지고, 학습이 안정된다. DiT는 이 값도 조건 정보의 함수로 만들어 노이즈 정도에 따라 네트워크 깊이를 부드럽게 조절한다. 여기까지가 클래스 조건 생성이고, 자유로운 텍스트 프롬프트로 넘어가려면 또 한 걸음이 필요하다. 텍스트 토큰을 교차 어텐션으로 넣는 방식(Pixart-alpha)은 텍스트가 이미지의 영향을 받지 않아 편집 작업에 취약한데, MMDiT는 이미지와 텍스트에 각자의 인코더를 주고 하나의 시퀀스로 합쳐 서로를 보게 한 뒤 다시 분리해 각자의 경로로 보낸다.

주요 인사이트

• 패치 분할은 단순한 최적화 기법이 아니라 트랜스포머를 이미지에 쓸 수 있게 만든 전제 조건이다. 계산량을 감당 가능한 수준으로 낮춰야 '모든 요소가 서로를 본다'는 장점을 살릴 수 있다.
• 조건 정보를 토큰으로 붙이는 가장 단순한 방법이 실패하는 이유는 어텐션이 한정된 자원이기 때문이다. 중요한 신호일수록 별도 경로를 마련해 주는 설계가 필요하다.
• 특징 변조 아이디어는 원래 이미지에 대한 질의응답 연구에서 나왔다. 다른 과제에서 검증된 구조가 생성 모델로 옮겨 오는 흐름은 딥러닝 아키텍처 발전의 전형적인 경로다.
• 학습 초기에 네트워크를 얕게 만드는 스케일 계수 기법은 ResNet과 확산 U-Net에서도 쓰이던 안정화 기법으로, 새 구조라고 해서 모든 것이 새로 발명되지는 않는다는 점을 보여 준다.
• 이미지와 텍스트를 대등하게 다루는 MMDiT의 설계는 '박쥐를 오른쪽으로 옮겨라' 같은 지시가 이미지에 따라 전혀 다른 의미가 되는 편집 작업에서 특히 중요하다.

자주 묻는 질문

관련 AI 소식