앞에서 막힌 곳은 이미지 전체의 로그확률 logp(x0) 였다. 같은 이미지에 이르는 디노이징 경로가 무수히 많아서 구할 수 없었다. 그런데 정책 그래디언트 식을 다시 보면, 필요한 것은 「내가 뽑은 출력」의 로그확률이다. 디퓨전이 이미지 한 장을 뽑을 때 실제로 한 일은 수십 번의 작은 스텝이었다. 그렇다면 정책 그래디언트에 꼭 필요한 것은 이미지의 확률일까, 내가 실제로 밟은 경로의 확률일까?
역사: 한 번의 행동으로 보던 시도, 여러 번으로 본 DDPO
2023년 초 디퓨전 모델을 보상에 맞추려던 시도들은 생성 전체를 한 번의 행동으로 봤다. 구글 리서치의 한 연구진은 사람 피드백으로 학습한 점수로 생성 이미지마다 가중치(곱하는 수 — 신경망 파라미터가 아니다)를 매기고, 그 가중치를 곱한 디퓨전 학습 손실을 줄였다. 로봇 학습에서 온 보상 가중 회귀(RWR)를 한 번 돌린 셈이다. 문제는 디퓨전의 학습 손실이 이미지 로그확률 그 자체가 아니라 그 아래 한계(ELBO)라는 데 있다. 이 방법은 원래 목표를 어림으로만 좇는다.
같은 해 세르게이 레빈(Sergey Levine) 연구실의 케빈 블랙(Kevin Black)과 동료들은 「Training Diffusion Models with Reinforcement Learning」에서 이 점을 짚었다. 어림이 생긴 까닭은 마지막 이미지 x0 만 보고 디노이징의 순서를 버렸기 때문이다. 디노이징을 여러 스텝에 걸친 결정 과정(MDP, Markov decision process — 상태를 보고 행동을 고르는 일이 이어지는 과정)으로 보면, 스텝마다의 로그확률을 어림 없이 쓸 수 있다. 그들이 붙인 이름이 DDPO다. 같은 보상으로 견준 실험에서 DDPO는 RWR 방식을 여러 번 되풀이한 것보다도 보상을 크게 올렸다.
스텝을 행동으로 보기
DDPO(Denoising Diffusion Policy Optimization)의 관점 전환은 단순하다. 이미지 하나를 한 번의 행동으로 보지 말고, 디노이징 T스텝을 T번의 행동으로 보라. 언어모델에서 응답을 토큰 T개의 연쇄로 봤던 것과 같다.
RL 용어
언어모델
디퓨전 (DDPO)
상태 st
프롬프트 + 지금까지의 토큰
프롬프트 + 현재 노이즈 낀 이미지 xt + 시점 t
행동 at
다음 토큰
한 스텝 덜 노이즈 낀 이미지 xt−1
정책
소프트맥스 분포
가우시안 N(μθ(xt,t),σt2I)
보상
응답이 끝난 뒤 한 번
이미지가 완성된 뒤(x0) 한 번
역방향 한 스텝은 xt−1=μθ(xt,t)+σtε 다(ε 은 앞에서처럼 표준 정규 노이즈). μθ는 모델의 예측이라 그래디언트가 여기로 흐르고, σtε는 파라미터와 무관한 노이즈라 탐색을 제공한다. 그리고 각 스텝이 가우시안이므로 이 경로의 로그확률은 적분 없이 식 하나로(닫힌 형태로) 정확히 계산된다.
언어모델의 정책 그래디언트 식에서 "토큰"을 "디노이징 스텝"으로 바꿨을 뿐이다. 보상이 끝에 한 번 오는 것까지 똑같다.
같은 시작점 xT에서도 스텝마다 다른 노이즈가 섞이면 다른 경로, 다른 이미지, 다른 보상이 나온다. 보상이 좋은 경로에서 μθ가 내린 판단은 강화되고, 나쁜 경로의 판단은 약화된다. 로봇 RL의 가우시안 정책 π(a∣s)=N(μθ(s),σ2)과 똑같은 구조다 — σ가 탐색, μθ가 학습된 판단이다.
문제 3 — 같은 6g, 다른 빵집
두 빵집이 표시 무게 500g 식빵을 판다. 빵집 가의 식빵은 날마다 무게가 표준편차 10g으로 흔들리고, 빵집 나는 2g으로 흔들린다(둘 다 정규분포). 오늘 두 집에서 산 식빵이 모두 506g이었다. (가) 각 빵집에서 506g은 평균에서 표준편차 몇 개만큼 벗어난 값인가? (나) 정규분포 밀도로 견주면 506g이 나올 밀도는 500g이 나올 밀도의 몇 분의 1인가? (다) 「오늘 반죽 기계가 바뀌었나」 의심할 만한 곳은 어느 빵집인가?
김민준
둘 다 6g 더 나갔으니 똑같이 수상하죠.
선생님
민준 학생, 나 빵집은 평소에 2g 안팎으로만 흔들리는 집이에요. 거기서 6g은 평소 흔들림의 몇 배죠? 가 빵집에서는요?
김민준
나는 6/2=3 개, 가는 6/10=0.6 개만큼이네요. 같은 6g이 아니었어요.
이서연
(나)는 밀도가 e−(표준편차개수)2/2 배로 줄어드니까, 가는 e−0.18≈0.84, 나는 e−4.5≈0.011.
김민준
exp(-0.6**2/2), exp(-3**2/2) 돌려 봐도 같아요.
김민준
나 빵집에서는 75배쯤 더 드문 일이에요. 의심할 곳은 나 빵집이네요. 늘 정확하던 집에서 벗어난 6g이 훨씬 큰 신호예요.
정리 (가) 가: 0.6개, 나: 3개. (나) 가: e−0.18≈0.84, 나: e−4.5≈0.011. (다) 나. 같은 편차도 흔들림(σ)이 작은 곳에서는 훨씬 드문 일이고, 그만큼 「무언가 바뀌었다」는 신호가 크다.
문제 4 — 스텝 로그확률의 그래디언트
한 디노이징 스텝에서 xt−1=μθ+σε 로 샘플했다. ∂logp/∂μθ 를 ε과 σ로 쓰고, ε=0.5일 때 σ=0.5와 σ=0.1에서 각각 값을 구하시오.
선생님
노이즈 σ를 0.5에서 0.1로 줄이면, 샘플 하나의 그래디언트는 커질까요 작아질까요?
김민준
빵집 문제만 보면 커질 것 같은데, 여기는 달라요. 벗어난 거리 σε 자체가 σ 에 비례해서 같이 줄잖아요. 표준편차 몇 개만큼 벗어났는지는 ε 그대로니까 그래디언트도 그대로예요.
이서연
식으로 보자. logp=−(x−μ)2/(2σ2)+c 니까 μ로 미분하면 (x−μ)/σ2=σε/σ2=ε/σ.
김민준
ε=0.5면 σ=0.5에서 1, σ=0.1에서 5… 커지네요.
선생님
벗어난 거리는 다섯 배로 줄었는데 왜 커졌을까요?
이서연
분모가 σ2 라서 스물다섯 배 줄었고, 분자는 다섯 배만 줄었어요. 분산이 작으면 같은 편차도 "드문 일"이라, 작은 σ에서 우연히 벗어난 방향이 로그확률에 크게 반영돼요.
김민준
샘플은 덜 흔들려도 그걸 설명하는 기울기는 더 예민해지는 거네요. 빵집 문제의 「몇 개만큼 벗어났나」를 한 번 더 σ로 나눈 셈이에요. 오차 범위가 좁은 실험일수록 조교님이 작은 편차도 크게 감점하던 거랑 같아요.
선생님
그래요. 디노이징 후반 스텝에서 같은 오차에 벌점이 훨씬 커지는 것과 같은 이유예요.
정리∂logp/∂μθ=ε/σ. σ=0.5이면 1, σ=0.1이면 5. 벗어난 거리는 σ 에 비례해 줄지만 분모의 σ2 가 더 빨리 줄어, 노이즈가 작을수록 샘플 하나의 그래디언트가 커지고 흔들림도 커진다.
문제 5 — 노이즈가 없으면?
결정론적 샘플러(DDIM — 노이즈를 섞지 않고 걷도록 고친 샘플러. 노이즈를 얼마나 섞을지 정하는 손잡이 η를 0으로 둔다)로 생성한 경로에 DDPO를 쓸 수 있는가? 문제 4의 식으로 설명하시오.
김민준
σ=0이면 ε/σ가 무한대라 계산이 안 되네요.
김민준
그럼 σ=10−4 같은 아주 작은 값을 넣고 계산하면 되잖아요.
선생님
민준 학생, 문제 4에서 σ를 줄이면 그래디언트가 어떻게 됐죠?
김민준
ε/σ라서… 10−4이면 ε의 만 배예요. 샘플마다 엄청나게 튀겠네요.
이서연
더 근본적인 문제도 있어. η=0이면 경로에 실제로 노이즈가 섞이지 않았으니 ε가 없어. σ를 작게 "가정"해도 그 경로는 그 분포에서 뽑힌 게 아니야. 분포가 한 점에 몰려 있으니 확률밀도 자체가 정의되지 않고.
선생님
그래요. 탐색은요?
김민준
처음 xT 말고는 무작위성이 없으니, 같은 xT에서는 늘 같은 이미지예요. 스텝 단위로 이 경로가 저 경로보다 나았다는 비교를 못 해요.
김민준
실험 조건을 하나도 안 바꿔 보고 민감도 분석을 하겠다는 셈이네요. 조교님이 변화를 줘봐야 뭐가 중요한지 안다고 했는데.
선생님
맞아요. 무작위성 없는 미분방정식으로 샘플링하는 모델은 모두 바로 이 상황이에요. 그래서 같은 분포를 주면서 매 순간 노이즈가 섞이는 방정식으로 바꿔, 노이즈를 다시 넣는 방법이 나왔어요.
정리 쓸 수 없다. σ=0이면 스텝 분포가 한 점이라 로그확률이 정의되지 않고(ε/σ가 한없이 커짐), 스텝 단위의 탐색도 사라진다. 결정론적 샘플러에 RL을 쓰려면 같은 분포를 주면서 노이즈가 섞이는 샘플링으로 바꿔 노이즈를 다시 넣어야 한다.
