Chapter 20: DDPO — 디퓨전의 정책 그래디언트

디퓨전 모델에 DPO를 옮긴 Diffusion-DPO는 쌍비교 데이터 — "이 이미지가 저 이미지보다 낫다"는 인간 레이블 — 가 있으면 보상 모델 없이도 디퓨전 모델을 정렬(사람이 바라는 방식으로 출력을 맞추기)할 수 있다.

하지만 현실에서는 쌍비교 데이터가 없고, 대신 보상 모델이 있는 경우가 훨씬 많다.

이런 상황에서 디퓨전 모델을 개선하고 싶다면? DPO는 쓸 수 없다 — 쌍이 없으니까. 그렇다면 전통적인 RL 방식, 즉 보상 모델 + 정책 그래디언트로 돌아가야 한다. LLM에서는 PPO(한 번에 너무 멀리 가지 않게 묶어 둔 정책 그래디언트)가 이 역할을 했다. 디퓨전에서도 같은 일을 할 수 있을까?

의문: 디퓨전에서 정책 그래디언트가 가능한가?

정책 그래디언트(답을 내는 모델, 곧 정책 $\pi_\theta$의 확률을 보상 쪽으로 미는 방법. 아래첨자 $\theta$ 는 모델의 파라미터)는 $\nabla J = \mathbb{E}[R \cdot \nabla \log \pi_\theta(y)]$ 였다. 필요한 것은 두 가지 — 내가 뽑은 출력의 로그확률과 그 그래디언트. 그런데 디퓨전에서는 같은 이미지에 이르는 디노이징 경로가 무수히 많아, 이미지 전체의 로그확률 $\log p(x_0)$는 모든 경로에 대한 적분으로만 정의되고 직접 구할 수 없다. 그렇다면 디퓨전에서 정책 그래디언트는 불가능한가?

여러 스텝: 디퓨전은 시간축의 자기회귀다

디퓨전 모델은 이미지 한 장을 그리려고 같은 네트워크를 수십 번 부른다. 언어모델이 토큰마다 한 번씩 부르는 것은 이해가 된다 — 앞 토큰이 정해져야 다음 토큰을 고를 수 있으니까. 그런데 이미지는 픽셀이 한꺼번에 나온다. 네트워크가 노이즈에서 이미지로 가는 답을 이미 배웠다면, 왜 한 번에 그리지 않고 스텝을 여러 번 밟는가?

역사: 한 번에 그리려던 GAN, 나눠 그린 디퓨전

2014년 몬트리올의 술집 「레 트루아 브라쇠르(Les 3 Brasseurs)」에서, 이언 굿펠로(Ian Goodfellow)는 사진을 저절로 만드는 컴퓨터를 궁리하던 친구들에게 신경망 둘을 서로 겨루게 해 보자고 했다. 그날 밤 집에 돌아가 새벽까지 코드를 짰고, 첫 시도에 돌아갔다고 전해진다. 이것이 GAN(생성적 적대 신경망)이다. 노이즈를 이미지로 한 번에 바꾸는 생성기를, 진짜와 가짜를 가려내는 판별자와 겨루게 해서 가르친다.

판별자를 속이기만 하면 되니 생성기는 데이터 전체를 덮을 이유가 없다. 굿펠로의 첫 논문도 이 약점을 적어 두었다. 판별자를 함께 고치지 않고 생성기만 너무 오래 학습시키면, 생성기가 서로 다른 노이즈를 같은 그림으로 몰아 버려 다양성을 잃는다. 논문은 이것을 「헬베티카 시나리오」라 불렀다. 뒤에 모드 붕괴(확률이 몰린 봉우리 몇 개로만 쏠림)라는 이름으로 굳은 이 쏠림과 학습 불안정이 GAN을 내내 따라다녔다. 2018년에는 멀리 떨어진 영역끼리 맞추라고 셀프 어텐션을 넣은 SAGAN(장(Han Zhang)과 동료들)도 나왔지만, 생성은 끝까지 한 번이었다.

디퓨전은 다른 쪽에서 왔다. 2015년 열역학에서 착안해 나오고 2020년 DDPM이 이미지 품질로 널리 알린, 데이터에 노이즈를 조금씩 넣는 과정을 거꾸로 되감는 모델이다. GAN이 스텝을 늘려서 디퓨전이 된 것이 아니다. 디퓨전이 발견한 것은, 스텝을 여러 개로 쪼개면 겨루기 없이 평범한 회귀만으로 생성이 된다는 사실이다.

한 번에 회귀하면 평균이 나온다

회귀가 무엇을 배우는지부터 보자. 디퓨전 네트워크는 노이즈 섞인 이미지 $x_t$ 를 받아 깨끗한 이미지 $x_0$ 를 맞히도록 제곱 오차로 학습한다(노이즈 $\varepsilon$ 을 맞히는 방식도 식 한 줄로 바꿔 쓸 수 있어 같은 이야기다).

제곱 오차를 최소로 하는 답은 언제나 조건부 평균이다. $\textcolor{#1c9c60}{x_t}$ 가 거의 순수한 노이즈라면 거기서 나올 수 있는 이미지는 무수히 많고, 그 평균은 어느 이미지도 아닌 흐릿한 얼룩이다. 한 번에 노이즈 → 이미지를 회귀하면 바로 그 얼룩이 나온다. 회귀는 봉우리 하나를 고르지 않고 봉우리들의 평균을 낸다.

장난감으로 확인하자. 데이터가 $-1$ 과 $+1$ 두 점뿐이고(반반), 노이즈 크기가 $\textcolor{#993600}{\sigma}$ 라 하자. 봉우리가 둘인 데이터 가운데 가장 간단한 것이라 고른 예다. 봉우리가 하나뿐이면 평균이 곧 그 봉우리여서, 「평균이 어느 봉우리도 아니다」를 시험할 거리가 없다.

같은 관측 $\textcolor{#1c9c60}{x_t} = 0.8$ 에서 노이즈가 크면($\textcolor{#993600}{\sigma} = 2$) 최선의 답은 $\tanh(0.2) \approx 0.197$ — $+1$ 쪽으로 조금 기운 평균이다. 노이즈가 작으면($\textcolor{#993600}{\sigma} = 0.3$) $\tanh(8.9) \approx 1.000$ 으로 한쪽에 붙는다. 노이즈가 줄수록 조건부 평균은 봉우리 하나로 좁혀진다.

움직이는 조건: 언어모델과 나란히

한 스텝에서 네트워크가 보는 조건은 두 가지다. 프롬프트 같은 외부 조건은 샘플링 내내 고정이다. 반면 지금까지 그려 놓은 $\textcolor{#1c9c60}{x_t}$ — 이 장에서 내부 조건이라 부르자 — 는 스텝마다 바뀐다. 한 스텝이 조금 덜 흐린 $\textcolor{#1c9c60}{x_{t-1}}$ 을 내놓으면, 그것이 다음 스텝의 조건이 된다. 조건이 한 번 움직일 때마다 조건부 평균도 움직이고, 흐릿한 평균이 특정한 그림 쪽으로 한 걸음씩 좁혀진다. 그래서 반복이 필요하다.

이것은 언어모델이 방금 쓴 토큰을 다음 토큰의 문맥으로 삼는 것과 같은 구조다. 디퓨전은 시간을 축으로 한 자기회귀다. 다른 것은 확정하는 순서뿐이다.

손가락 개수는 픽셀 하나하나를 따로 봐서는 맞출 수 없다. 「여기에 이미 손가락 네 개가 있다」는 부분적 결정을 읽어야 다섯 번째에서 멈춘다. 한 스텝의 출력은 이 조율이 덜 된 추정이지만, 그것을 다시 입력으로 넣으면 이미 내린 결정이 다음 스텝의 조건이 된다. 결합분포(모든 픽셀이 함께 맞아야 하는 분포)를 한 번에 맞추지 못하니, 조금씩 확정해 가며 맞추는 것이다. 그래도 자주 틀린다는 것은 누구나 아는 대로다.

추정이 바뀌는 것은 정상이다

스텝마다 네트워크의 답 $\mathbb{E}[\textcolor{#1c9c60}{x_0} \mid \textcolor{#1c9c60}{x_t}]$ 가 바뀐다면, 처음 답은 틀렸던 것인가? 내비게이션으로 비유해 보자. 서울에서 부산으로 출발할 때 내비가 「4시간」이라고 한다. 이 4시간은 나에 대한 예측이 아니다. 아직 정해지지 않은 모든 경우 — 막히는 날과 뚫린 날, 휴게소에 들르는 차와 안 들르는 차 — 를 평균 낸 값이다. 30분을 달리면 나는 더 이상 평균적인 차가 아니다. 어느 길로 들어섰는지, 얼마나 막혔는지가 정해졌고, 그 조건 아래의 평균이 새 예측이 된다. 추정이 바뀐 이유는 처음 답이 틀려서가 아니라 조건이 특정되어서다.

그리고 처음 값은 틀린 게 아니라 뭉뚱그려진 값이다. 조건부 평균에는 이런 성질이 있다: 지금의 추정은 앞으로 갱신될 추정들의 평균과 같다. 장난감에서 확인하면, $\textcolor{#993600}{\sigma} = 2$ 에서 본 $\textcolor{#1c9c60}{x_t} = 0.8$ 의 추정 0.197은, 한 스텝 덜 흐린 $\textcolor{#993600}{\sigma} = 1$ 의 관측이 가능한 모든 값에서 새로 낸 추정 $\tanh(\cdot / 1)$ 들을 평균한 값과 같다(수치 적분으로 0.1974). 한 스텝 뒤의 추정 하나하나는 $+1$ 이나 $-1$ 쪽으로 더 벌어지지만, 그 평균은 처음 값을 벗어나지 않는다. 이런 성질을 가진 과정을 마팅게일(martingale)이라 부른다.

아래 그림은 같은 출발점에서 노이즈를 한 단계씩 걷어 내며 추정을 새로 낸 경로 여덟 개다. 경로마다 $+1$ 이나 $-1$ 로 갈라지지만, 이런 경로를 20만 개 그려 평균하면 어느 노이즈 크기에서나 0.197 근처에 머문다.

그렇다면 출발 시점의 방향과 속도로 4시간을 곧장 달리면 어떻게 될까? 부산이 아닌 곳에 도착한다. 첫 추정이 가리키는 곳은 흐릿한 평균이기 때문이다. 큰 스텝 하나로 가는 것이 이와 같다. 이렇게 보면 디노이징 궤적이 휘는 것도 원인이 아니라 결과다 — 조건이 쌓이며 조건부 평균이 이동한 자취가 휜 궤적이다.

ML에서: 셀프 어텐션과 적은 스텝의 모델

내부 조건을 읽는 부품은 셀프 어텐션이다. 요즘 이미지 생성 모델의 몸통은 트랜스포머이거나(디퓨전 트랜스포머), 컨볼루션 층 사이사이에 어텐션 블록을 끼운 U-Net(가운데가 좁아지는 U자 모양 신경망)이다. 어느 쪽이든 지금까지 그린 그림의 각 부분이 서로를 읽는 셀프 어텐션은 빠질 수 없는 부품이다. 프롬프트(외부 조건)는 따로 둔 크로스 어텐션으로 읽거나, 글과 그림을 한데 넣은 어텐션으로 함께 읽는다.

그러면 한 스텝 모델은 불가능한가? 가능은 하다. 다만 일반화라는 목표에서는 거의 늘 비실용적이다. 적은 스텝 모델은 대부분 여러 스텝 모델이 이미 그려 놓은 궤적을 교사로 삼아 시작과 끝을 잇도록 증류(큰 모델의 출력을 작은·빠른 모델이 따라 하게 가르치기)한 것이다. 살리먼스(Tim Salimans)와 호(Jonathan Ho)의 점진적 증류(2022)는 8192스텝 샘플러를 절반씩 줄여 4스텝까지 내렸다. 적대적 디퓨전 증류(ADD, 자우어(Axel Sauer)와 동료들, 2023)는 한 스텝 생성이 되지만, 교사인 SDXL(공개된 이미지 생성 모델)의 수준에는 4스텝에서 닿았다고 보고했다. 같은 논문은 한 스텝 표본의 다양성이 교사보다 낮고, 스텝을 늘리면 다양성이 오른다고도 적었다. 내부 조건을 스텝마다 갱신하며 좁혀 가는 과정을 없앤 것이 아니라, 이미 답을 아는 교사에게서 압축해 받아 온 것이다.

강화학습 쪽에서 보면 이 절의 결론은 한 줄이다. 디퓨전의 생성은 여러 번의 결정이다. 각 스텝이 그때까지의 결과를 조건으로 다음 결과를 고른다 — 행동이 여럿인 결정 과정이다.

문제 1 — 내일 강수량을 숫자 하나로 예보하면

비 올 확률이 30%면 안 올 가능성이 더 크니까 0mm로 예보하죠. 예보는 실제로 일어날 일을 말해야 하잖아요.

민준 학생, 그 예보로 1년을 지내면 하루 평균 제곱 오차가 얼마나 될까요? 비 오는 날과 안 오는 날을 나눠서요.

비 오는 날은 $(10-0)^2 = 100$ 인데 그런 날이 30%니까 30. 안 오는 날은 0. 평균 30이요.

10mm로 예보하면 안 오는 70%의 날에 100씩 잃으니까 70. 3mm면 $0.3 \times 49 + 0.7 \times 9 = 14.7 + 6.3 = 21$. 셋 중에 3mm가 가장 작아.

3mm는 비 오는 날에도 안 오는 날에도 실제로는 안 내리는 양인데요.

그래요. 제곱 오차로 매기면 어느 날에나 맞는 값이 아니라 평균을 고르게 돼요. (나)에서 아침 하늘을 보면 무엇이 달라지죠?

흐린 아침이면 $0.9 \times 10 = 9$mm, 맑은 아침이면 $0.1 \times 10 = 1$mm. 하늘이라는 조건이 붙으니 예보가 한쪽으로 기울어요.

(다)는 $0.25 \times 9 + 0.75 \times 1 = 3$. 전날의 3mm는 틀린 예보가 아니라, 아침에 바뀔 예보들의 평균이었네요.

학기 초 예상 학점이 중간고사 뒤에 바뀌는 거랑 같아요. 처음 예상이 틀렸던 게 아니라 아직 점수가 안 나왔던 거죠.

정리 (가) 0mm: 30, 10mm: 70, 3mm: 21. 제곱 오차의 최선은 가장 그럴듯한 값이 아니라 평균(3mm)이다. (나) 흐린 아침 9mm, 맑은 아침 1mm. (다) $0.25 \times 9 + 0.75 \times 1 = 3$. 조건이 붙을 때마다 예보는 바뀌지만, 바뀔 예보들의 평균은 처음 예보와 같다.

문제 2 — 한쪽이 더 흔한 두 점

민준 학생, 순수한 노이즈 하나만 보고 내는 최선의 답은요?

아까 비 예보랑 같아요. 평균이니까 $0.8 \times (-1) + 0.2 \times 1 = -0.6$. 식에 $\textcolor{#993600}{\sigma} \to \infty$ 를 넣어도 $\tanh(\tfrac12 \ln 0.25) = -0.6$ 이에요.

조교님이 정답 확신이 없으면 부분점수라도 받게 쓰라던 거랑 같아요. 찍어서 틀리면 크게 잃으니까요.

그럼 (나), $\textcolor{#1c9c60}{x_t} = 0.5$ 에 $\textcolor{#993600}{\sigma} = 1$ 이면 부호가 어떻게 될까요?

0.5는 $+1$ 쪽에 더 가까우니까 양수죠.

넣어 봐. $\tanh(0.5 - 0.693) = \tanh(-0.193) \approx -0.19$ 야. 음수야.

관측은 $+1$ 쪽인데 답은 $-1$ 쪽이에요?

관측 0.5가 $+1$ 에서 나왔을 가능성은 $-1$ 에서 나왔을 가능성의 몇 배죠? 그리고 처음부터 두 점은 몇 대 몇이었죠?

정규분포 밀도의 비가 $e^{2 \times 0.5 / 1} = e \approx 2.7$ 배로 $+1$ 쪽이에요. 그런데 처음부터 $-1$ 이 네 배 많았으니까, 2.7배로는 아직 못 뒤집어요. 관측이 미는 쪽과 원래 많던 쪽을 같이 봐야 하는구나.

$\textcolor{#993600}{\sigma} = 0.4$ 면 $\tanh(0.5/0.16 - 0.693) = \tanh(2.43) \approx 0.985$. 노이즈가 작아지니 관측이 이겨서 $+1$ 에 거의 붙어요.

그런데 (가)의 $-0.6$ 은 데이터에 없는 값이잖아요. 그건 모델이 틀린 거 아닌가요? 학습을 더 시키면 $-1$ 이나 $+1$ 을 낼 것 같은데요.

순수한 노이즈 $\textcolor{#1c9c60}{x_T}$ 하나만 보고 $-1$ 인지 $+1$ 인지 가를 정보가 있나요?

없어요. 그럼 $-0.6$ 이 틀린 게 아니라 제곱 오차에서는 그게 최선이에요. 아무리 학습해도 한 번에 회귀하면 두 봉우리의 평균, 흐릿한 답이 나와요.

해석학에서 최소제곱 근사는 가진 정보 안에서 가장 가까운 것을 준다고 배웠잖아요. 정보가 없으면 평균밖에 못 주는 거고요.

그래요. 그러니 조금 덜 흐린 $\textcolor{#1c9c60}{x_t}$ 를 만들어 조건을 바꾸고, 다시 추정하는 거예요.

정리 (가) $\textcolor{#993600}{\sigma} \to \infty$ 에서 추정은 $\tanh(\tfrac12\ln 0.25) = -0.6 = 0.8 \times (-1) + 0.2 \times 1$ — 가장 흔한 점($-1$)도 가운데(0)도 아닌 평균이다. 학습이 덜 된 것이 아니라 정보가 없어서다. (나) $\textcolor{#993600}{\sigma} = 1$: $\tanh(-0.193) \approx -0.19$, $\textcolor{#993600}{\sigma} = 0.4$: $\tanh(2.43) \approx 0.985$. 관측이 $+1$ 쪽이어도 노이즈가 크면 원래 흔한 쪽이 이기고, 노이즈가 작아지면 관측이 이긴다. (다) 한 번의 회귀는 평균만 낼 수 있으므로, 조금 덜 흐린 결과를 다음 스텝의 조건으로 넣어 조건부 평균을 한 봉우리로 좁혀 간다.