확산 모델(DDPM) 구현: U-Net 설계와 타임스텝 포지셔널 인코딩, 역과정 샘플링까지

핵심 메시지

• 확산 모델 구현에는 이전 이미지를 직접 예측하는 방법, 원본을 한 번에 예측하는 방법, 섞인 노이즈를 예측하는 방법이 있고 초기 확산 모델은 세 번째를 쓴다.
• 신경망은 이미지 분할용으로 나온 U-Net을 변형해 쓴다. 다운샘플링과 업샘플링 사이를 스킵 커넥션으로 이어 위치 정보를 되살리는 것이 핵심이다.
• 일반 신경망과 결정적으로 다른 점은 입력에 타임스텝 정보가 함께 들어간다는 것이다. 정수인 타임스텝을 트랜스포머의 포지셔널 인코딩 방식으로 벡터화한다.
• 학습 데이터를 만들 때는 노이즈를 한 단계씩 쌓을 필요가 없다. 누적곱으로 정의된 계수를 쓰면 원본에서 원하는 단계의 노이즈 이미지를 한 번에 계산할 수 있다.
• 역과정에서 노이즈를 빼면서 다시 노이즈를 더하는 이유는, 신경망 출력이 이전 이미지 자체가 아니라 그 평균이고 거기에 분산을 더해 표본을 뽑기 때문이다.

쉽게 이해하기

확산 모델의 이론을 다룬 앞 회차에 이어, 이번에는 코드로 직접 구현하는 과정을 다룬 스터디 발표다. 목표는 명확하다. 무작위 노이즈를 넣으면 노이즈가 조금 제거된 결과를 내놓고, 그것을 다시 넣기를 반복해 결국 손글씨 숫자 이미지를 만들어 내는 신경망을 만드는 것이다.

먼저 구현 방식의 갈래가 정리됐다. 첫째는 노이즈가 섞인 현재 이미지에서 한 단계 앞 이미지의 평균을 예측하고 분산과 함께 표본을 뽑는 방식이고, 둘째는 중간 단계를 건너뛰고 원본 이미지를 곧바로 예측해 단계 수를 줄이는 방식이다. 셋째는 두 단계 사이에 더해진 노이즈 자체를 예측하는 방식으로, 이번에 구현하는 초기 확산 모델이 이 방식을 택하며 손실 함수는 예측한 노이즈와 실제로 넣은 노이즈의 제곱 오차다. 발표자는 이 모델이 초기 버전이라 단계 수가 수백에서 수천에 이르고 품질에도 한계가 있으며, 후속 모델들이 그 지점을 개선해 왔다고 덧붙였다.

신경망 구조로는 의료 영상의 영역 구분에 널리 쓰이던 U-Net을 변형해 쓴다. 합성곱과 맥스 풀링으로 크기를 줄이며 병목까지 내려간 뒤 업샘플링으로 다시 올라오는데, 핵심은 내려가는 길에서 만든 특징 맵을 올라오는 길의 같은 높이에 이어 붙이는 스킵 커넥션이다. 원 논문 구조에서는 합성곱을 거치며 크기가 줄어 이어 붙일 두 특징 맵의 크기가 맞지 않는데, 발표자도 이 지점의 해법을 코드에서 찾지 못했다고 솔직히 밝혔다. 다만 요즘 구현은 대부분 패딩을 줘 크기를 그대로 유지하므로 문제가 되지 않고, 이번 예제도 28×28 흑백 이미지를 유지한 채 채널 수만 늘려 간다.

가장 까다로운 부분은 타임스텝을 넣는 방법이다. 신경망은 이미지 형태를 받는데 타임스텝은 100이나 200 같은 정수라 형식이 맞지 않는다. 해결책은 트랜스포머에서 쓰던 포지셔널 인코딩으로, 짝수 위치에는 사인 홀수 위치에는 코사인을 적용해 정수 하나를 벡터로 바꾼다. 이 벡터의 차원을 16으로 둔 것은 이미지 크기와 아무 상관이 없는 임의의 설정값이며, 발표 중에도 같은 질문이 나와 의미가 없다는 답이 이어졌다. 각 합성곱 블록 안에 둔 작은 다층 퍼셉트론이 이 벡터를 블록의 채널 수에 맞게 변환하고 형태를 바꿔 더하면, 브로드캐스팅이 가로세로 전체에 같은 값을 채워 넣는다.

학습 데이터를 만드는 정방향 과정에서는 각 단계에서 더할 노이즈 크기를 미리 정해 두는 베타 스케줄링이 나온다. 다만 실제로는 단계마다 노이즈를 하나씩 쌓지 않고, 계수의 누적곱을 써서 원본과 목표 단계만으로 그 단계의 노이즈 이미지를 한 번에 계산한다. 생성은 반대 방향이다. 무작위 노이즈에서 시작해 매 단계 현재 이미지와 단계 번호를 넣어 섞인 노이즈를 예측하고, 그것을 빼서 평균을 구한 뒤 표준편차를 더해 표본을 뽑는다. 노이즈를 제거해야 하는데 왜 다시 더하느냐는 의문에 대해 발표자는 신경망이 내놓는 것이 이전 이미지 자체가 아니라 그 평균이며, 생성 모델은 언제나 평균과 분산을 예측한 뒤 표본을 뽑는 구조라고 정리했다.

주요 인사이트

• 구현 방식이 세 갈래로 나뉜다는 점을 먼저 짚고 들어가면, 확산 모델 코드를 읽을 때 신경망 출력이 무엇을 뜻하는지부터 확인하게 된다. 같은 확산 모델이라도 출력이 이미지일 수도 노이즈일 수도 있다.
• 타임스텝 주입은 논문 수식에는 잘 드러나지 않지만 구현에서는 가장 손이 많이 가는 부분이다. 정수 하나를 벡터로 바꾸고, 다시 각 블록의 채널 수에 맞춰 변환한 뒤, 브로드캐스팅으로 특징 맵에 더하는 세 단계를 거친다.
• 포지셔널 인코딩 차원을 16으로 둔 것이 아무 근거 없는 설정값이라는 확인은 실전에서 유용하다. 코드의 숫자를 모두 의미 있는 값으로 읽으려다 시간을 낭비하는 경우가 많다.
• 누적곱 계수로 정방향 과정을 한 번에 계산하는 요령은 단순한 최적화가 아니라 학습을 가능하게 만드는 조건이다. 단계마다 반복문을 돌면 학습 데이터를 만드는 비용만으로 시간이 다 간다.
• 역과정에서 다시 노이즈를 더하는 이유를 '평균과 분산을 예측한 뒤 표본을 뽑는다'로 정리하면, 확산 모델이 변분 오토인코더와 같은 계열이라는 점이 자연스럽게 드러난다.

자주 묻는 질문

관련 AI 소식