Your GFlowNet Secretly Learns an Optimal Transport Plan

기술 요약: 당신의 GFlowNet은 최적 운송 계획을 비밀리에 학습한다

문제 정의
생성적 흐름 네트워크(Generative Flow Networks, GFlowNets)는 유향 그래프에서의 확률적 궤적을 통해 구조화된 이산 객체를 샘플링하도록 설계되었습니다. 초기에는 비순환 그래프를 위해 개발되었으나, 최근 연구는 비순환(non-acyclic) 환경으로 이를 확장했습니다. 이러한 환경에서는 샘플링 효율성을 위해 기대 궤적 길이를 최소화하는 것이 매우 중요하며, 이는 "최소 흐름(minimum-flow)" GFlowNet의 정식화로 이어집니다. 이와 동시에, 최적 운송(Optimal Transport, OT)은 일반적으로 선형 계획법(Kantorovich 정식화)을 통해 해결되는 운송 비용 최소화라는 원칙적인 프레임워크를 통해 확률 분포를 비교하는 방법을 제공합니다. 그러나 대규모 그래프에서 OT를 해결하는 것은 정확한 솔버를 사용하는 방식으로는 계산적으로 불가능한 경우가 많습니다.

본 논문은 최소 흐름 정식화된 비순환 GFlowNet이 본질적으로 최적 운송 문제를 해결하는지 조사함으로써 이 두 분야 사이의 간극을 다룹니다.

방법론
저자들은 비순환 GFlowNet과 그래프 상의 이산 최적 운송 사이의 이론적 동등성을 확립합니다. 핵심 방법론은 다음과 같습니다:

1. 이론적 재정식화: 저자들은 비순환 GFlowNet을 위한 최소 흐름 최적화 문제에서 시작합니다. 이는 상세 균형(detailed balance) 및 보상 매칭(reward matching) 제약 조건 하에서 총 흐름(기대 궤적 길이와 동일)을 최소화합니다. 특정 소스 분포 $L(u)$에 맞추기 위해 초기 엣지-흐름 분포(소스 상태 $s_0$에서 나가는 흐름)를 고정함으로써, GFlowNet 목적 함수를 선형 계획법(LP) 문제로 변환합니다.
2. 동등성 증명: 저자들은 이 제약된 LP가 Kantorovich 최적 운송 문제와 수학적으로 동일함을 입증합니다. 이 정식화에서:
   • 소스 분포는 $U$ 집합에 정의된 $L(u)$입니다.
   • 타겟 분포는 $X$ 집합에 정의된 보상 분포 $R(x)$입니다.
   • 운송 비용 $d(u, x)$는 그래프 내 $u$와 $x$ 사이의 최단 경로 거리로 정의됩니다.
   • 최적의 GFlowNet 정책은 총 운송 비용을 최소화하는 최적 결합(optimal coupling) $\Pi$를 암시적으로 학습합니다.
3. 학습 알고리즘: 정확한 LP가 불가능한 대규모 그래프에서 이를 해결하기 위해, 저자들은 GFlowNet의 신경망 파라미터화를 사용합니다. 이들은 흐름 정규화를 포함한 궤적 균형(Trajectory Balance, TB) 목적 함수를 활용합니다. 주요 수정 사항은 타겟 분포 $R(x)$에 대한 표준 보상 매칭과 더불어, 고정된 초기 분포 $L(u)$를 강제하기 위해 손실 함수에 "leward"(왼쪽 보상) 항을 포함하는 것입니다.

주요 기여

1. 이론적 동등성: 본 논문은 최소 흐름 GFlowNet에서 초기 흐름 분포를 고정하는 것이 그 목적 함수를 그래프 유도 최단 경로 비용을 가진 Kantoricity OT 문제로 정확히 환원시킨다는 것을 증از합니다. 결과적으로, 최적의 GFlowNet 순방향 정책은 최적 운송 계획을 인코딩하며, 이로부터 궤적을 샘플링하는 것은 최적 결합을 복구합니다.
2. Beckmann 문제와의 연결: 저자들은 이 정식화가 GFlowNet을 고전적인 네트워크 흐름 이론과 연결하는 그래프 기반의 이산 Beckmann 문제와 동등함을 보여줍니다.
3. 확장 가능한 OT 솔버: 본 연구는 GFlowNet 학습 프레임워크가 정확한 OT 솔버로는 다루기 힘든 대규모 그래프 OT 문제의 근사 솔루션을 구하는 데 적용될 수 있음을 보여줍니다. 신경망과 확률적 궤적 샘플링을 사용함으로써, 이 방법은 조합론적 공간으로 확장 가능합니다.
4. 최단 경로 이론의 복구: 부수적으로, 이 구성은 비순환 GFlowNet을 위해 이전에 확립된 최단 경로 이론을 복구하며, 최적의 흐름이 최단 경로의 부분 그래프 상에서만 지지됨을 보여줍니다.

실험 결과
저자들은 두 가지 환경에서 이론과 방법을 검증했습니다:

• 하이퍼그리드(Hypergrid) 환경: 작은 하이퍼그리드($H=10, 15, 20$)에서 학습된 GFlowNet 정책은 고전적인 LP 솔버(POT)를 통해 얻은 정확한 해와 일치하는 운송 비용($OT^\star$)을 달-성합니다. 이 방법은 샘플링에 편향을 주지 않으면서 다양한 소스 분포(ball 및 moon 형태)에 대해 최적의 결합을 성공적으로 복구합니다.
• 순열(Permutation) 환경: 대칭군 $S_n$의 케일리 그래프(Cayley graph)에서, 방법은 정확한 해를 구할 수 있는 $n=4$ 및 $n=8$에 대해 최적에 가까운 해를 복구합니다. 정확한 OT가 불가능한 $n=20$의 경우에도, 고정점 확률에서 낮은 오차를 보이는 합리적인 근사치를 생성합니다.
• 절제 연구(Ablation): 정규화 계수 $\lambda$에 대한 실험은 궤적 길이(샘플링 효율성)와 샘플링 최적성(운송 계획의 정확도) 사이의 트레이드오프를 확인해 줍니다. 낮은 $\lambda$ 값은 더 정확한 운송 계획을 생성하지만 더 긴 궤적을 유발합니다.

의의 및 주장
본 논문은 "비밀스러운" 연결고리를 발견했다고 주장합니다: 즉, 흐름을 최소화하도록 정식화된 GFlowNet은 암묵적으로 최적 운송 문제를 풀고 있다는 것입니다. 이 연구의 의의는 다음과 같습니다:

• 통합: GFlowNet 학습과 최적 운송을 통합하는 새로운 이론적 관점을 제공하며, GFlowNet이 운송 맵을 학습하는 방법으로 간주될 수 있음을 시사합니다.
• 실질적 응용: 신경망과 샘플링을 통해 확장 가능한 GFlowNet 학습 알고리즘을 사용하여, 정확한 솔버로는 계산적으로 불가능한 대규모 그래프 OT 문제에 적용할 수 있게 합니다.
• 프레임워크 확장: GFlowNet 프레임워크를 단순한 샘플링 및 생성 작업을 넘어 계산적 최적 운송 영역으로 확장하여, 질량의 이동 위치뿐만 아니라 근사 최적 경로 비용을 가진 실행 가능한 국소적 이동을 통해 질량이 어떻게 이동하는지를 결정하는 확률적 정책을 제공합니다.

저자들은 향후 함의에 대해 신중한 태도를 유지하며, 이 방법이 계산적으로 효율적이긴 하지만 최적성과 효율성의 균형을 맞추기 위해 정규화 계수 $\lambda$의 세심한 튜닝이 필요하다고 언급합니다. 또한 향-후 연구에서 이 방법들을 비정규화된 분포 간의 운송 맵을 학습하는 데 적용하는 것을 탐구할 수 있다고 제안합니다.