흐름 매칭(Flow Matching)을 위한 온라인 최적 수송(OT) 미니 배치 스케일링

흐름 모델(flow models)을 확장하는 데 있어 핵심적인 과제는 학습 및 통합 과정에서 안정성을 유지하는 것입니다. 흐름 매칭(flow matching)에서는 이러한 안정성이 소스 샘플이 데이터와 어떻게 쌍을 이루는지에 결정적으로 의존합니다. 독립적인 무작위 쌍을 맺는 것은 상당한 확률적 변동성(stochasticity)을 야기하며, 이는 미니 배치 최적 수송(minibatch optimal transport, OT)을 통해 이러한 무작위성을 줄이는 온라인 소스-타겟 결합(online source-target couplings)을 구성하도록 동기를 부여했습니다 (Tong et al., 2024). 하지만 응용 분야가 더 크고 고차원적인 데이터셋으로 확장됨에 따라, 작은 미니 배치는 이러한 이점을 제한할 수 있으며, 매 반복마다 결합을 상당히 큰 미니 배치로 스케일링하는 것은 계산적으로 불가능합니다. 더욱이, OT 결합이 전체 확률적 기울기(full stochastic gradient) 및 통합(integration)에 어떻게 영향을 미치는지에 대한 이론적인 해결책은 여전히 남아 있습니다. 본 연구에서는 OT의 이점을 유지하면서 상당히 큰 소스-타겟 결합을 가능하게 하는 확장 가능한 온라인 엔트로피 OT 근사치를 개발합니다. 우리는 비정규화된(unregularized) 및 엔트로피 OT 모두에 대해 전체 확률적 기울기 분산과 통합 오차의 엄격한 감소를 증명함으로써 이론적인 격차를 해소하고, 이러한 개선 사항이 온라인 결합 미니 배치가 커짐에 따라 어떻게 강화되는지 추가로 특성화합니다. 단일 세포 유전자 보간 흐름(single-cell gene interpolation-flow) 및 이미지 흐름 학습에서의 실험은 이러한 예측을 검증하며, 더 큰 온라인 결합이 일관되게 학습 안정성과 최종 생성 성능을 향상시킨다는 것을 보여줍니다.