잠재 확산 모델 원리 정리: VAE 잠재 공간과 확산 트랜스포머(DiT) 구조를 MIT 강의로 이해하기

핵심 메시지

• 가로 1,000 세로 600에 색 채널이 셋인 이미지를 픽셀 그대로 벡터로 다루면 180만 차원이 되는데, 확산 모델은 샘플을 하나 만들 때 신경망을 50~100번 반복 호출하므로 이 부담이 특히 크다.
• 오토인코더로 차원을 줄이면 압축은 되지만 잠재 공간의 데이터 분포가 들쭉날쭉해져 다음 단계인 확산 모델 학습이 오히려 어려워진다.
• VAE는 인코더를 확률적으로 바꾸고 KL 발산 항을 더해 잠재 분포를 표준 가우시안에 가깝게 만들며, 재구성 손실과의 균형은 가중치 베타로 조절한다.
• 오늘날 공개된 이미지·영상 생성 모델은 사실상 모두 잠재 공간에서 생성한 뒤 디코더로 되돌리는 방식이며, 오토인코더 설계 자체가 별도의 연구 주제가 됐다.
• 확산 트랜스포머는 시간을 사인·코사인 임베딩으로, 프롬프트를 CLIP·T5 같은 사전학습 임베딩으로, 이미지를 패치 시퀀스로 바꿔 하나의 트랜스포머에서 처리한다.

쉽게 이해하기

MIT의 생성 AI 강의 6.S184 네 번째 시간은 앞서 배운 플로 매칭과 확산 모델을 '대규모로' 돌리려면 무엇이 더 필요한지를 다룬다. 강연자는 두 가지를 꼽는다. 하나는 데이터를 압축하는 잠재 공간이고, 다른 하나는 그동안 추상적으로만 가정해 온 벡터장을 실제로 구현할 신경망 구조다.

왜 압축이 필요한지는 숫자로 드러난다. 색 채널 3개에 높이 600, 너비 1,000인 이미지는 벡터로 펼치면 180만 차원이다. GPU 메모리가 감당하기 어렵고, 학습 문제 자체도 어려워지며, 옆 픽셀만 봐도 채워 넣을 수 있는 중복 정보가 많다. 강연자는 이미지 분류와의 차이를 강조한다. 분류는 고차원 입력을 클래스 하나로 줄이지만, 확산 모델의 벡터장은 180만 차원을 다시 180만 차원으로 옮기는 함수인 데다 샘플링 한 번에 50~100회씩 호출된다.

해법은 인코더와 디코더를 둔 오토인코더다. 데이터를 훨씬 낮은 차원의 잠재 공간으로 보내고 되돌리며 평균제곱오차를 줄이도록 학습시키면 압축과 복원은 해결된다. 그러나 문제가 하나 남는다. 이렇게 얻은 잠재 분포가 어떤 모양일지 통제할 수단이 없다는 점이다. 분포가 거칠고 매끄럽지 않으면 그 위에 확산 모델을 얹는 다음 단계가 훨씬 어려워진다.

그래서 강의는 VAE로 넘어간다. 인코더가 결정적 값 대신 평균과 분산을 내놓는 가우시안 분포가 되고, 여기에 KL 발산으로 '인코딩 결과가 표준 가우시안처럼 보이게 하라'는 항을 더한다. 다만 이 항만 최소화하면 데이터를 완전히 무시하는 해가 나오므로, 재구성 손실과 사전분포 손실을 가중치 베타로 더해 균형을 잡는다. 베타를 얼마로 둘지는 이론이 아니라 경험의 영역이고, 강연자는 오토인코더 품질을 확인하는 유일한 방법이 그 위에 확산 모델을 실제로 학습시켜 보는 것이라는 점을 이 분야의 골칫거리로 꼽는다.

마지막 조각은 재파라미터화 트릭이다. 기댓값을 계산하는 분포 자체가 인코더 파라미터에 의존하면 그대로는 유효한 그래디언트를 얻을 수 없어서, 표준 가우시안 노이즈를 뽑아 평균과 분산으로 변환하는 형태로 손실을 다시 쓴다. 강연자는 강화학습의 REINFORCE 알고리즘도 정책이 파라미터에 의존해 같은 성격의 문제를 겪는다고 덧붙인다. 여기까지 오면 학습 데이터를 전부 잠재로 인코딩해 두고, 그 잠재 위에서 확산 모델을 학습한 뒤 샘플을 디코딩하는 잠재 확산 절차가 완성된다.

주요 인사이트

• 확산 모델에서 고차원이 유독 문제인 이유는 두 가지다. 벡터를 벡터로 옮기는 매핑이라 출력 차원이 줄지 않고, 샘플 하나를 만들 때 네트워크를 수십 번 호출한다.
• 오토인코더 품질 평가는 비용이 비대칭이다. 잠재 공간이 '좋은지' 확인하려면 그 위에서 확산 모델을 학습해야 하는데, 그 학습에 수백만 달러가 들 수 있다.
• 영상 생성은 이미지에 프레임 수가 곱해져 텐서가 훨씬 커지므로, 패치를 시간축까지 포함한 3차원으로 나누고 오토인코더 설계의 중요도가 더 올라간다. 강연에서 언급된 메타의 무비 젠은 300억 파라미터 규모에 H100 6,000장으로 학습됐다.
• 논문마다 시간 규약이 다르고(데이터가 시각 1인지 0인지), 이산 시간 형식이나 전방 확산 과정처럼 서술 방식도 갈리지만, 가우시안 설정에서는 대체로 같은 대상을 다르게 표기한 것에 가깝다.

자주 묻는 질문

관련 AI 소식