플로우 매칭·스코어 매칭 학습 원리 — MIT 6.S184 강의로 보는 확산 모델 훈련 알고리즘

핵심 메시지

• 학습 목표인 '주변 벡터장'은 데이터셋 전체에 대한 고차원 적분이라 직접 계산할 수 없지만, 데이터 한 점에 조건을 건 손실로 바꾸면 계산이 가능해진다.
• 두 손실은 신경망 파라미터와 무관한 상수만큼만 차이 나기 때문에 최소점도 같고 경사도 같아, 조건부 손실만 줄여도 원래 목표를 달성한다.
• 가우시안 경로에 직선 보간을 쓰면 학습은 '노이즈와 데이터를 섞은 점을 넣고 둘의 차이를 맞히기'라는 평균제곱오차 회귀 한 줄로 줄어든다.
• 확산 모델의 스코어 학습도 구조가 같으며, 가우시안 경로에서는 벡터장과 스코어가 서로 변환되므로 신경망을 하나만 학습해도 된다.
• 이 단순한 알고리즘이 메타의 영상 생성 모델과 스테이블 디퓨전 3 같은 실제 대형 생성 모델의 훈련 방식이다.

쉽게 이해하기

MIT 6.S184 세 번째 강의는 앞선 두 강의에서 정의한 확률 경로, 벡터장, 스코어 함수를 실제 학습 알고리즘으로 옮기는 회차다. 강의자는 생성 모델이 결국 초기 분포에서 출발해 미분방정식을 시간 0에서 1까지 시뮬레이션하고 마지막 시점의 결과를 표본으로 삼는 구조라는 점을 다시 확인한 뒤, 그 미분방정식을 이끄는 벡터장을 신경망으로 흉내 내는 것이 훈련의 목표라고 정리한다.

문제는 목표가 되는 '주변 벡터장'을 직접 계산할 수 없다는 데 있다. 이 값은 데이터 분포 전체에 대한 적분으로 정의되기 때문에, 이론적으로는 데이터셋을 전부 훑어 평균을 내야 하고 학습 중 신경망을 한 번 평가할 때마다 수천만 장의 이미지를 순회하는 일이 된다. 강의자는 이 손실을 실제로 계산하는 사람은 없다고 잘라 말한 뒤 대안을 제시한다.

대안은 데이터 한 점에 조건을 건 '조건부 플로우 매칭 손실'이다. 조건부 벡터장은 해석적 공식이 있어 바로 계산할 수 있고, 두 손실은 신경망 파라미터와 무관한 상수만큼만 차이가 난다. 그래서 최소점이 같고 경사도 같으며, 확률적 경사하강법의 학습 궤적까지 동일해진다. 결국 계산 가능한 쪽만 최소화해도 원래 원하던 목표에 도달한다.

구체적인 알고리즘은 놀랄 만큼 단순하다. 시간을 0과 1 사이에서 균일하게 뽑고, 데이터셋에서 한 점을 뽑고, 표준 정규 노이즈를 뽑아 둘을 섞은 중간 지점을 만든 뒤, 그 점을 신경망에 넣어 정해진 가중 평균을 맞히게 한다. 계수를 가장 단순하게 잡으면 노이즈와 데이터를 잇는 직선 위의 한 점을 넣고 두 점의 차이를 예측하는 형태가 되며, 이 경로에는 별도의 이름까지 붙어 있다.

확산 모델 쪽도 같은 논리를 그대로 반복한다. 주변 스코어를 목표로 하는 손실은 계산할 수 없으니 조건부 스코어를 목표로 하는 '디노이징 스코어 매칭' 손실로 바꾸고, 가우시안 경로에서는 이 목표가 '데이터를 오염시킨 노이즈를 맞히기'로 정리된다. 강의자는 노이즈 크기가 아주 작아지는 구간에서 이 손실이 수치적으로 불안정해지며, 논문들이 스코어 신경망 대신 노이즈 예측기로 다시 쓰는 이유도 여기에 있다고 설명한다.

주요 인사이트

• '계산할 수 없는 목표'를 '계산할 수 있는 목표'로 바꾸는 트릭이 이 분야의 핵심이다. 신경망은 한 번도 본 적 없는 주변 벡터장을, 수많은 데이터 점에 대한 조건부 목표를 반복해 맞히는 것만으로 결과적으로 학습한다.
• 생성 모델 훈련이 반드시 복잡할 필요는 없다. 최소최대 최적화가 필요한 적대적 생성망과 달리, 플로우 매칭은 단일 평균제곱오차 회귀여서 구현과 안정성 면에서 유리하다.
• 직선 경로가 선택되는 이유는 이론적 필연이 아니라 경험적 판단에 가깝다. 강의자는 직선이 미분방정식 시뮬레이션 오차를 줄인다는 직관을 소개하면서도, 결국 여러 선택지를 실험해 고르는 문제라고 말한다.
• 가우시안 경로에 한정하면 벡터장과 스코어가 대수적으로 변환되기 때문에 초기 확산 모델 연구들이 스코어 매칭만 다룬 것이 자연스럽다. 다만 이 등가성은 일반적인 경우까지 확장되지 않는다.
• 플로우 매칭 관점의 진짜 확장성은 시작 분포를 자유롭게 고를 수 있다는 데 있다. 노이즈 대신 무음 영상, 저해상도 이미지, 처리 전 세포 데이터처럼 의미 있는 분포에서 출발하는 변환도 같은 틀로 다룰 수 있다.

자주 묻는 질문

관련 AI 소식