다이아몬드 맵: 생성 모델 보상 정렬의 가치 함수 추정을 여러 표본으로 개선한 방법

핵심 메시지

• 보상 정렬은 좋은 결과를 내지만 원하는 방향으로 몰아가기 어려운 기반 모델을, 특정 취향이나 과제에 맞게 조종하는 문제다.
• 정렬의 핵심은 생성 도중의 중간 상태가 얼마나 유망한지를 재는 가치 함수인데, 중간 상태는 아직 완성된 표본처럼 보이지 않아 값을 직접 정의할 수 없다.
• 기존 플로우 맵 계열은 중간 상태에서 최종 결과를 결정론적으로 하나만 뽑아 가치를 추정하기 때문에 느리거나 부정확했다.
• 다이아몬드 맵은 같은 중간 상태에서 최종 결과를 여러 개 표집해 보상을 평균 내므로 추정이 훨씬 안정적이고, 그만큼 정렬 성능이 실제로 올라간다.
• 미리 학습된 플로우 맵을 파인튜닝하는 방식과, 학습 없이 추론 시점에만 얹는 방식을 함께 제시해 학습 연산과 추론 연산 사이에서 선택할 수 있게 했다.

쉽게 이해하기

Y Combinator가 스타트업 스쿨에서 처음 연 머신러닝 세션에서 더글러스 첸이 다이아몬드 맵이라는 방법을 소개했다. 다루는 문제는 보상 정렬이다. 이미지 생성 모델이 좋은 이미지를 만들어 낸다고 해서, 내가 원하는 특정 화풍이나 조건 쪽으로 그 모델을 몰아갈 수 있는 것은 아니다. 정렬은 기반 모델을 그대로 두고 원하는 방향으로 조종하는 작업이다.

정렬은 보상 유도에 기대고, 보상 유도를 하려면 가치 함수를 추정해야 한다. 가치가 높다는 것은 그 방향으로 갔을 때 내가 원하는 표본이 나올 가능성이 높다는 뜻이다. 프롬프트 정렬이라면 이미지와 프롬프트가 얼마나 맞는지가 가치가 된다. 문제는 생성 과정이 잡음에서 출발해 최종 분포로 가는 동안 중간 단계의 상태들이 아직 어떤 표본처럼도 보이지 않아, 가치 함수가 그 위에 정의되지 않는다는 점이다.

그래서 중간 상태에서 최종 결과까지 건너뛰어 가치를 재는 방법이 필요한데, 플로우 매칭과 플로우 맵이 그 역할을 한다. 다만 이들은 결정론적으로 움직여 중간 상태 하나에서 최종 결과도 하나만 나온다. 실제로는 같은 중간 상태가 아주 다른 여러 이미지로 자라날 수 있으므로, 표본 하나로 낸 추정은 좋을 수가 없다. 발표자가 지적한 기존 방법들의 두 약점이 바로 느리거나, 결정론적이어서 추정이 부실하다는 것이었다.

다이아몬드 맵은 같은 중간 상태에서 최종 표본을 여러 개 뽑아 각각의 보상을 계산하고 이를 합쳐 가치 함수를 추정한다. 발표에서는 정렬 품질 지표로 LPIPS를 써서 일반적인 플로우 매칭보다 좋은 성적을 보였고, 아무런 조건부 학습도 하지 않은 ImageNet 모델을 프롬프트에 맞춰 조종한 예시도 함께 제시했다. 자연 상태에서는 나올 리 없는 조합, 예컨대 사막 한가운데의 등대 같은 프롬프트를 일부러 골라 정렬 효과를 확인했다.

구현은 두 갈래다. 하나는 미리 학습된 플로우 맵을 파인튜닝하는 방식이고, 다른 하나는 학습 없이 추론 시점에만 붙이는 방식이다. 후자는 기존 플로우 맵 위에 그대로 얹어 확률적으로 표집하게 만들 수 있어, 학습 연산 대신 추론 연산을 쓰는 맞바꿈이 된다.

주요 인사이트

• 중간 상태 하나가 여러 결말로 이어질 수 있다는 점을 인정하는 순간, 결정론적 추정은 편향된 추정이 된다. 다이아몬드 맵의 개선은 새로운 모델이 아니라 추정 절차를 확률적으로 바꾼 데서 나온다.
• 학습 없이 추론 시점에만 얹는 방식의 실질적 가치는 규모에서 드러났다. 발표자는 파인튜닝할 연산이 충분치 않아 ImageNet급이 아닌 Flux 같은 대규모 텍스트-이미지 모델에서는 추론 시점 방식으로 실험했고, 그 덕에 더 큰 규모에서 결과를 낼 수 있었다고 밝혔다.
• 두 방식을 함께 쓰는 것은 권하지 않았다. 이미 파인튜닝한 모델이 있다면 그것을 쓰는 편이 낫고, 둘을 겹쳐 봐야 결과는 비슷한데 연산만 더 든다는 설명이다.
• 정렬 문제를 '모델을 다시 학습시키는 일'에서 '생성 궤적을 어떻게 안내할지의 문제'로 옮겨 놓으면, 이미 배포된 생성 모델도 재학습 없이 용도에 맞게 조정할 여지가 생긴다.

자주 묻는 질문

관련 AI 소식