이미지를 만드는 길을 곧게 펴면 빨라질까: Rectified Flow 읽기

생성 모델이 느린 이유를 신경망 크기에서만 찾으면 절반을 놓치기 쉽다. 작은 네트워크도 같은 일을 수십 번 반복하면 느리다. 그렇다면 한 번의 계산을 줄이는 대신, 잡음이 이미지로 변하는 경로 자체를 단순하게 만들 수는 없을까. Rectified Flow는 이 질문을 놀라울 만큼 직접적으로 다룬다. 목적지까지 가는 길을 곧게 펴면 적은 발걸음으로도 도착할 수 있다는 생각이다.

물론 이미지 공간에서 직선이라는 말은 사진을 점점 밝게 만드는 효과를 뜻하지 않는다. 이미지 한 장을 거대한 벡터로 보고, 두 벡터 사이를 선형 보간한다는 뜻이다. 이 평범한 연산이 생성 모델의 학습 목적이 되고, 다시 빠른 샘플링의 근거가 된다는 점이 이 논문의 묘미다.

출발점과 도착점만 있으면 만들 수 있는 정답

출발 분포에서 x0를, 도착 분포에서 x1을 뽑는다. 이미지 생성에서는 x0가 가우시안 잡음이고 x1이 실제 이미지가 될 수 있다. 시간 t가 0에서 1로 증가할 때 중간 상태를 다음과 같이 만든다.

xt = (1 - t) × x0 + t × x1

이 직선의 속도는 항상 x1 - x0다. 따라서 신경망 v는 xt와 t를 입력으로 받아 이 속도를 맞히면 된다. 학습은 예측 속도와 목표 속도 사이의 제곱오차를 줄이는 회귀다. 실제 생성에서는 새 잡음에서 시작해 현재 상태 += 시간 간격 × 예측 속도를 반복한다. 원 논문의 알고리즘은 이 회귀와 ODE 풀이를 분명히 구분한다.

여기까지 들으면 의문이 생긴다. 이미 x1을 알아야 직선을 만들 수 있는데, 생성할 때는 x1을 모른다. 학습 중에는 도착점을 이용해서 정답을 만들지만, 학습한 함수는 현재 위치와 시간만 보고 속도를 내놓는다. 정답을 아는 학습 과정과 정답을 찾는 추론 과정의 정보 차이가 핵심이다.

직선으로 가르쳐도 처음부터 직선으로 움직이지 않는다

평면 위에 출발점 두 개와 도착점 두 개를 찍고 서로 교차하게 연결해 보자. 교차 지점에서는 같은 위치와 같은 시간에 서로 다른 속도가 필요할 수 있다. 하지만 현재 상태와 시간만 보는 신경망은 둘을 구별할 수 없다. 제곱오차를 최소화하면 그 조건에서 가능한 속도들의 평균을 예측하게 된다.

이 평균 속도장을 따라 움직이는 경로는 학습에 사용한 개별 직선과 다를 수 있다. 적절한 조건에서 해가 유일한 ODE의 궤적은 같은 시간에 만나서 서로 다른 방향으로 갈라질 수 없기 때문이다. Rectified Flow는 원래의 연결을 그대로 암기하기보다, 분포를 옮길 수 있는 새로운 연결을 만들어 낸다.

그래서 선형 보간으로 학습했다고 한 단계 생성이 자동으로 완성되는 것은 아니다. 입력으로 준 학습 경로가 직선인지, 실제 속도장을 적분해서 얻는 생성 경로가 직선인지 따로 봐야 한다. 이 구분 없이 논문을 읽으면 가장 중요한 절차가 사소한 후처리처럼 보인다.

reflow는 생성 결과로 연결을 다시 짠다

먼저 배운 모델에 출발 샘플을 넣고 끝까지 이동시킨다. 이제 하나의 출발점과 그 모델이 실제로 만든 도착점이라는 새로운 짝이 생긴다. 이 짝들을 다시 직선으로 연결해 새 속도 모델을 학습하는 것이 reflow다. 처음에 무작위로 짝지었던 관계 대신, 기존 모델이 만들어 낸 더 일관된 관계를 사용한다.

직관적으로는 복잡하게 엉킨 배송 동선을 한 번 운행한 뒤, 실제 배달된 출발지와 목적지를 바탕으로 새 경로를 설계하는 과정과 닮았다. 이 비유에서 배송 물품의 신원이 반드시 보존된다고 생각하면 안 된다. 생성 문제의 관심사는 개별 학습 사진의 복사가 아니라 최종 샘플들이 목표 분포를 따르는지에 있다.

원문은 이상적인 조건에서 rectification이 여러 볼록 운송 비용을 증가시키지 않는 성질을 보이고, 반복할수록 경로를 곧게 만드는 효과를 분석한다. 그러나 이것이 일반적인 고차원 문제에서 전역 최적 운송해를 보장한다는 뜻은 아니다. 실제 신경망은 회귀 문제를 정확하게 풀지 못하며, ODE도 유한한 간격으로 계산한다. 정리의 조건과 구현의 근사 사이에는 거리가 있다.

곧은 경로가 왜 빠른지도 조금 더 정확히 말할 필요가 있다. 논문에서 이상적인 직선 경로는 경로를 따라 속도가 일정한 경우다. 이때 한 번의 Euler 업데이트로 끝점을 정확하게 계산할 수 있다. 길의 모양만 곧고 속도가 계속 변한다면, 시작점의 속도만으로 끝점을 맞히기 어렵다.

빠르다는 결과를 읽는 기준

논문은 이미지 생성, 이미지 간 변환, 도메인 적응을 다룬다. 특히 reflow로 얻은 거의 곧은 경로가 아주 적은 단계에서도 좋은 결과를 낼 수 있음을 보여 준다. 중요한 실험 질문은 가장 많은 단계를 썼을 때의 최고 품질뿐 아니라, 단계 수를 줄일 때 품질이 얼마나 유지되는가다. 저자들의 설명 페이지에서도 이 연결을 시각적으로 확인할 수 있다.

한편 작은 단계 수가 곧 작은 전체 비용이라는 보장은 없다. reflow에는 이전 모델로 학습용 짝을 생성하는 비용과 새 모델을 훈련하는 비용이 든다. 생성량이 적은 작업과 매일 많은 샘플을 만드는 서비스는 이 비용을 다르게 평가해야 한다. 추론에서 아낀 시간을 사전 처리와 추가 학습에 투자하는 선택이라고 보는 편이 현실적이다.

또한 이미지 간 변환에서 그럴듯한 도착 이미지를 만드는 것과 출발 이미지의 원하는 의미를 보존하는 것은 별도의 목표다. 얼굴을 다른 도메인으로 옮겼을 때 배경이나 표정이 유지되는지는 운송 비용 하나가 충분히 설명하지 못한다. 픽셀 공간에서 짧은 이동이 사용자가 원하는 변환과 같은지부터 정의해야 한다.

서비스에서는 경로와 지연시간을 함께 본다

실무적으로 재미있는 적용은 단계 수에 따른 품질 곡선을 먼저 그리는 일이다. 같은 초기 잡음에서 한 단계, 몇 단계, 더 촘촘한 적분 결과를 비교한다. 적은 단계에서 구도가 크게 달라진다면 단순한 디테일 손실보다 경로 근사의 문제가 클 수 있다. 반대로 구도는 유지되고 질감만 달라진다면 서비스가 받아들일 수 있는 품질 기준을 더 구체적으로 세울 수 있다.

여기에 네트워크 호출 수와 실제 처리시간을 함께 기록해야 한다. 한 호출의 계산량, 배치 크기, 데이터 이동 비용이 다르기 때문이다. 내가 이 논문에서 가져가는 실무적 해석은 빠른 모델을 찾기 전에 빠르게 계산할 수 있는 학습 목표를 설계해 보자는 것이다. 논문이 특정 제품의 지연시간을 보장하는 것은 아니다.

Rectified Flow를 읽고 나면 생성 모델의 경로는 결과를 설명하는 그림 이상의 의미를 갖는다. 어떤 경로를 학습시키는지가 수치해석의 난도를 정하고, 그 난도가 실제 호출 수를 정한다. 분포를 맞추는 통계적 문제와 계산을 줄이는 공학적 문제가 한 식에서 만나는 순간이 이 연구의 가장 흥미로운 부분이다.

논문 정보: Xingchao Liu, Chengyue Gong, Qiang Liu, Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, 2022 공개, ICLR 2023. 논문 원문 · 공식 구현