확산 모델(Diffusion Model)이란? 파이토치로 구현하는 순방향·역방향 과정과 U-Net

핵심 메시지

• 확산 모델은 이미지에 노이즈를 점진적으로 더해 완전한 잡음으로 만든 뒤, 신경망으로 그 과정을 되돌려(디노이징) 새 데이터를 생성하는 생성 모델이다.
• VAE는 학습이 쉽지만 결과가 흐릿하고 GAN은 고품질이나 학습이 불안정한 반면, 확산 모델은 고품질과 다양성을 함께 얻지만 순차적 역과정 탓에 샘플링이 느리다.
• 구현에는 노이즈를 순차적으로 더하는 스케줄러, 이미지 속 노이즈를 예측하는 U-Net, 현재 시간 단계를 알려주는 위치 임베딩 세 가지가 필요하다.
• 가우시안의 합은 여전히 가우시안이라는 성질 덕분에, 누적 분산 스케줄로 임의의 시간 단계 t의 노이즈 이미지를 순차 반복 없이 한 번에 계산할 수 있다.
• 손실 함수는 결국 예측한 노이즈와 실제 노이즈의 L2 거리로 단순화되며, 저자는 스탠퍼드 카(약 1만6천 장) 데이터로 500에폭 학습해 자동차 이미지를 생성해 보였다.

쉽게 이해하기

확산 모델은 생성 딥러닝의 한 갈래로, 데이터 분포를 학습해 새 데이터를 만들어낸다. GAN·VAE·정규화 흐름 등 기존 생성 모델과 비교하면, VAE는 학습이 쉽지만 출력이 흐릿하고 GAN은 고품질이지만 그래디언트 소실·모드 붕괴 같은 학습 난제를 안고 있다. 확산 모델은 비교적 최근 등장해 고품질이면서도 다양한 샘플을 만들어, DALL·E 2나 이매진 같은 텍스트 기반 이미지 생성에서 큰 성공을 거뒀다.

동작 원리는 입력 이미지를 노이즈만 남을 때까지 점진적으로 망가뜨린 뒤, 신경망으로 그 과정을 되돌리는 것이다. 각 단계가 직전 단계에만 의존하므로 이는 마르코프 연쇄이며, 잠재 상태가 입력과 같은 차원을 유지한다는 특징이 있다. 모델의 과제는 각 이미지에 더해진 노이즈를 예측하는 것으로, 순수한 잡음에서 역과정을 밟으면 새로운 데이터가 생성된다. 대신 순차적 역과정 때문에 GAN·VAE보다 샘플링이 느리다.

순방향 과정은 비교적 단순하다. 분산 스케줄 베타(β)가 각 단계에서 더할 노이즈 양을 정하고, 다음 이미지는 직전 이미지에 √(1-β)를 곱한 값을 평균으로 하는 조건부 가우시안 분포에서 샘플링된다. 저자는 β를 선형으로 최대 0.02까지 키우며, 알파(α=1-β)의 누적곱을 이용하면 특정 시간 단계의 노이즈 이미지를 순차 반복 없이 곧바로 샘플링할 수 있다고 설명한다. 덕분에 학습은 시간 단계 t를 무작위로 뽑아 해당 노이즈 이미지를 모델에 넘기는 방식으로 매끄러워진다.

역방향(디노이징) 과정에는 U-Net을 쓴다. U-Net은 오토인코더처럼 입력을 다운샘플링하며 채널을 늘려 병목까지 압축했다가 다시 업샘플링하는 구조로, 출력이 입력과 같은 형태여서 이미지 생성에 잘 맞는다. 저자는 잔차 연결과 다운·업샘플링 같은 핵심 요소만 담은 '심플 U-Net'을 만들며, 분산을 고정했기에 픽셀당 한 값(노이즈의 평균)만 출력한다. 모델은 모든 단계에서 가중치를 공유하므로, 지금이 몇 번째 시간 단계인지 알려주기 위해 트랜스포머의 사인·코사인 위치 임베딩을 추가 입력으로 넣는다.

손실 함수는 변분 하한에서 출발하지만, 디노이징 스코어 매칭과 연결되며 결국 예측 노이즈와 실제 노이즈의 L2 거리라는 단순한 형태로 정리된다. 저자의 심플 U-Net은 약 6천만 개 파라미터로, 콜랩에서 100에폭을 돌렸을 때는 형태만 잡히다가 개인 GPU에서 500에폭까지 학습하자 뚜렷한 자동차 이미지가 나왔다. 확산 모델은 이미지에 그치지 않고 분자 그래프·오디오 등으로 확장되고 있으며, 확산-GAN 같은 변형도 등장하고 있다.

주요 인사이트

• 확산 모델의 핵심 트릭은 '가우시안의 합은 가우시안'이라는 성질이다. 이 덕분에 누적 분산 스케줄로 임의의 t 단계 노이즈 이미지를 한 번에 만들 수 있어, 순차 반복 없이 학습이 훨씬 매끄러워진다.
• 데이터는 -1~1 범위로 정규화해야 분산 스케줄과 맞물려 최종적으로 등방성 가우시안(평균 0, 고정 분산)에 도달한다. 노이즈를 너무 적게도, 너무 많이도 더하면 안 된다.
• 모델이 이미지 자체가 아니라 '더해진 노이즈'를 예측하도록 하는 것이 디노이징 스코어 매칭이며, 이미지 평균을 직접 예측하는 방식도 실험적으로 작동한다.
• 모든 시간 단계가 같은 가중치를 공유하기 때문에, 사인·코사인 위치 임베딩으로 현재 단계를 알려주지 않으면 모델은 서로 다른 노이즈 강도를 구분하지 못한다.
• 6천만 파라미터의 단순 모델로도 500에폭 학습 시 알아볼 수 있는 자동차 이미지가 나왔다는 점은, 그룹 정규화·어텐션 등 개선을 더하면 고품질 생성이 가능함을 시사한다.

자주 묻는 질문

관련 AI 소식