기술 요약: 반사된 슈뢰딩거 브릿지 매칭 (Reflected Schrödinger Bridge Matching)

문제 정의

최근 생성 모델링, 특히 슈뢰딩거 브릿지(Schrödinger bridges, SBs)의 발전은 플로우 매칭(flow matching)에서 영감을 받은 부분적 시뮬레이션 프리(partially simulation-free) 훈련 방식(예: 반복적 마르코프 적합법 또는 IMF)을 통해 고차원 환경에서 최적 운송 계획(optimal transport plans)을 효율적으로 계산할 수 있게 했습니다. 그러나 기존의 SB 프레임워크는 주로 무제약 참조 프로세스(예: 표준 브라운 운동)에 의존합니다. 이는 이미지 픽셀이 $[0, 1]$로 제한되는 것과 같이 데이터 도메인이 자연스럽게 유계(bounded)인 응용 분야에서 중요한 한계점이 됩니다.

반사 역학(특히 반사 브라운 운동, RBM)은 생성된 샘플이 데이터 도메인 내에 엄격히 머물도록 보장하는 메커니즘을 제공하지만, 현재의 반사 슈뢰딩거 브릿지(Reflected Schrödinger Bridges, RSBs) 학습 방법은 계산 비용이 매우 높습니다. Deng 등(2024)의 기존 방식은 다음과 같은 전방-후방 확률 미분 방정식(SDE) 이론에 의존합니다:

1. 모든 훈련 샘플에 대해 전체 경로(entire paths)를 완전히 시뮬레이션해야 함.
2. 다이버전스 항(divergence terms) 및 고차 미분을 포함하는 값비싼 손실 함수를 사용함.
3. 단지 끝점(endpoints)만이 아니라 전체 경로를 캐싱(caching)해야 함.

본 논문은 단위 초입방체(unit hypercube) 상에서 정의된 반사 역학을 위해, IMF(및 그 변형인 $\alpha$-IMF, DSB 매칭)의 효율적인 부분적 시뮬레이션 프리 훈련 프레임워크를 적용하는 과제를 다룹니다.

방법론

저자들은 Diffusion Schrödinger Bridge Matching (DSBM) 프레임워크의 확장인 **Reflected Schrödinger Bridge Matching (RSBM)**을 제안합니다. 핵심 목표는 컴팩트한 도메인 $D$(구체적으로 단위 초입방체) 상에서 두 분포 $\pi_0$와 $\pi_1$ 사이의 슈뢰딩거 브릿지 $P^*$를 반사 참조 프로세스 $Q_r$을 사용하여 근사하는 것입니다.

알고리즘은 두 가지 투영(projection) 단계를 교대로 수행합니다:

1. 상호적 투영(Reciprocal Projection): 참조 프로세스로부터 브릿지 혼합물(bridge mixtures)을 샘플링합니다.
2. 마르코프 투영(Markovian Projection): 현재 경로 측도(path measure)와 마르코프 프로세스 클래스 사이의 쿨백-라이블러(KL) 다이버전스를 최소화하기 위해 신경망을 사용하여 최적의 드리프트 보정 $v^*$를 근사합니다.

반사 설정에서 참조 프로세스는 다음과 같은 반사 SDE로 정의됩니다:
$$dX_t = \mu(t, X_t)dt + \sigma(t)dW_t + n(X_t)dL_t$$
여기서 $n(X_t)$는 안쪽을 향하는 법선 벡터이며, $L_t$는 프로세스가 $D$ 내에 머물도록 하는 로컬 타임(local time)입니다. 마르코프 투영 단계에서는 반사 브릿지의 스코어 함수(score function)인 $\nabla \log Q^r_{1|t}(X_1 | X_t)$를 근사하도록 네트워크 $v_\theta$를 훈련시킵니다.

핵심 기술 구성 요소

전체 경로 시뮬레이션 없이 훈련을 용이하게 하기 위해, 저자들은 세 가지 구체적인 혁신을 도입했습니다:

1. 효율적인 스코어 함수 근사:
   단위 큐브 상의 반사 브라운 운동에 대해, 전이 밀도(transition density)는 반사 연산자 $R$로부터 유도된 합-오브-가우시안(sum-of-Gaussians) 접근법을 사용하여 근사됩니다. 스코어 함수는 밀도를 차원별로 인수분해함으로써 성분별로 계산되며, 이를 통해 전체 반사 집합의 지수적 복잡성을 피합니다. 이를 통해 복잡한 PDE를 풀지 않고도 손실 함수에 필요한 Doob h-변환 항을 계산할 수 있습니다.

2. 효율적인 브릿지 샘플링 (알고리즘 2):
   반사 브릿지 분포 $Q^r_{t|0,1}$로부터 샘플링하는 것은 까다롭습니다. 저자들은 다음과 같은 샘플링 알고리즘을 제안합니다:

   • 타겟 $x_1$의 반사 연산자 하의 역상(pre-images) 집합으로부터, 반사되지 않은 가우시안 전이 밀도에 의해 가중치가 부여된 "가상" 끝점 $Y_1$을 추출합니다.
   • 시작점 $x_0$와 가상 끝점 $Y_1$이 조건인 표준 브라운 브릿지 중간점 $Y_t$를 샘플링합니다.
   • 최종 샘플을 얻기 위해 결과 $R(Y_t)$를 반사합니다.
     이 방법은 브릿지 샘플링 단계에서 SDE 경로를 경로별로 시뮬레이션하는 것을 방지합니다.

3. 양방향 훈련 손실:
   훈련은 $\alpha$-DSBM과 유사하게 시간 방향으로 양방향(forward and backward) 손실 함수를 사용합니다. 손실 함수는 네트워크 출력과 반사 전이 밀도에서 유도된 타겟 스코어 함수 사이의 제곱 오차를 최소화합니다. 알고리즘은 엔드포인트 쌍과 브릿지 중간점만을 필요로 하며, 전체 궤적(trajectories)을 필요로 하지 않고 온라인으로 가중치를 업데이트합니다.

주요 기여

• 알고리즘 확장: 본 논문은 효율적인 IMF/DSB 프레임워크를 반사 역학으로 확장하여, 부분적 시뮬레이션 프리 접근 방식을 사용하여 유계 도메인 상에서 슈뢰딩거 브릿지를 훈련할 수 있게 합니다.
• 다루기 쉬운 구현: 참조 프로세스를 단위 큐브 상의 반사 브라운 운동으로 제한함으로써, 필요한 스코어 함수에 대한 폐쇄형 근사(closed-form approximations)와 효율적인 브릿지 중간점 샘플링 체계를 도출하였으며, 이를 통해 값비싼 전방-후방 SDE 솔버의 필요성을 제거했습니다.
• 성능 효율성: 제안된 방법은 표준 플로우 매칭과 비교했을 때 계산 시간이 거의 늘어나지 않으면서도 높은 효율성을 유지합니다.

실험 결과

저자들은 이미지 간 변환(image-to-image translation) 태스크에 대해 $\alpha$-RSBM을 비반사 $\alpha$-DSBM과 비교 평가했습니다:

• 데이터셋: MNIST $\leftrightarrow$ EMNIST (28$\times$28 그레이스케일) 및 AFHQ (Cat $\leftrightarrow$ Wild, 64$\times$64).
• 지표: 프레셰 인셉션 거리(FID), 학습된 지각적 이미지 패치 유사도(LPIPS), 그리고 시작점과 끝점 사이의 평균 제곱 거리(MSD).

결과:

• 생성 품질: $\alpha$-RSBM은 $\alpha$-DSBM과 대등하거나 약간 개선된 생성 성능(FID 및 LPIPS)을 달 기록했습니다.
• 도메인 준수: $\alpha$-RSBM의 주요 장점은 샘플이 유효한 픽셀 도메인 $[0, 1]$ 내에 머물도록 보장한다는 점입니다. 이와 대조적으로, $\alpha$-DSBM은 상당한 범위를 벗어난 위반(out-of-bounds violations)을 생성했습니다 (예: MNIST 샘플의 100%가 적어도 하나의 범위를 벗어난 픽셀을 가졌으며, 평균 픽록의 46%가 제약 조건을 위반함).
• 유사성: 반사 모델은 대부분의 경우 소스 이미지에 대해 약간 더 높은 유사성(더 낮은 MSD)을 유지했습니다.
• 훈련 시간: 오버헤드는 미미했습니다. MNIST-EMNIST의 경우, 단일 A100에서 RSBM 훈련에는 32시간, DSBM에는 31.5시간이 걸렸습니다. AFHQ의 경우, 4개의 A100에서 RSBM은 109시간, DSB는 108시간이 걸렸습니다.

의의 및 주장

본 논문은 RSBM이 반사 확산 슈뢰딩거 브릿지를 훈련하는 "저렴한 방법"이라고 주장합니다. IMF 프레임워크를 활용함으로써, 이전에 제안된 값비싼 전방-후방 SDE 방법들에 대한 실질적인 대안을 제공합니다.

저자들은 다음과 같은 점을 강조합니다:

1. 도메인 제약 보장: 클리핑(clipping)을 통해 아티팩트를 유발하거나 범위를 벗어난 샘플을 생성하는 일반적인 확산 모델이나 무제약 SB와 달리, RSBM은 데이터 분포의 서포트(support)를 본질적으로 존중합니다.
2. 효율성 유지: 경로 캐싱이나 다이버전스 기반 손실과 같은 무거운 계산 부담 없이 이러한 보장을 달성하며, 고차원 이미지 데이터로 확장 가능합니다.
3. 이론적 특성 보존: 이 방법은 원래의 IMF 알고리즘의 수렴 특성을 유지하여, 학습된 프로세스가 유효한 슈뢰딩거 브릿지임을 보장합니다.

저자들은 한계점도 언급했습니다. 현재 이론은 제로 드리프트(zero-drift) 참조 프로세스(분산 보존 SDE 제외)와 단위 큐브 도메인을 가정합니다. 이들은 더 복잡한 기하학적 구조나 역학이 가능하지만, 이를 위해서는 추가적인 이론적 발전이나 대안적인 수치적 근사(예: 신경망 브릿지 샘플러)가 필요하다고 제언했습니다.