Evaluating GFlowNet from partial episodes for stable and flexible policy-based training

이 논문은 GFlowNet 의 흐름 균형 원리를 활용하여 부분 에피소드 기반의 평가 균형을 제안함으로써, 정책 기반 학습의 신뢰성을 강화하고 오프라인 데이터 수집 및 매개변수화된 역방향 정책 통합과 같은 유연성을 확보하는 방법을 제시합니다.

핵심

🎨 핵심 비유: "미술관 그림 그리기 프로젝트"

상상해 보세요. 여러분은 거대한 미술관에서 수많은 새로운 그림을 그려야 하는 프로젝트 팀장입니다.

• 목표: 미술관에는 이미 '좋아하는 그림들 (보상, Reward)'이 정해져 있습니다. 우리는 이 좋아하는 그림들과 비슷한 비율로 새로운 그림들을 무작위로 그려내야 합니다. (예: 풍경화가 30%, 초상화가 20% 등)
• 문제: 가능한 그림의 종류가 너무 많아서 (우주만큼 많음), 모든 그림을 다 그려보고 "어떤 게 좋은지" 계산하는 것은 불가능합니다.

이때 GFlowNet은 "한 번에 한 줄씩 그림을 그려가는 (경로, Trajectory)" 방식으로 이 문제를 해결하려 합니다.

🏃‍♂️ 기존 방법의 문제점: "나침반이 흔들리는 등산"

기존의 GFlowNet 훈련 방식은 크게 두 가지 접근법이 있었습니다.

1. 흐름 맞추기 (Value-based): "이 길로 가면 물 (흐름) 이 얼마나 흐를까?"를 계산하며 길을 찾습니다. (안정적이지만 유연성이 떨어짐)
2. 정책 학습 (Policy-based): "내가 그린 그림이 얼마나 좋은지 점수 (평가 함수, V) 를 매겨서, 다음에 더 잘 그릴 수 있도록 배우는 것"입니다. (유연하지만 점수를 매기는 '심사위원'을 믿기 어려움)

여기서 큰 문제가 생깁니다.
정책 학습 방식은 "내가 그린 그림의 점수 (V)"를 정확히 알아야 하는데, 이 점수를 매기는 '심사위원 (평가 함수)'을 신뢰할 수 없게 훈련시키는 경우가 많았습니다. 마치 나침반이 자꾸 흔들리는 등산과 같아서, 팀원들이 길을 잃거나 너무 천천히 진전하는 문제가 있었습니다.

💡 이 논문의 해결책: "Sub-EB (부분 구간 평가 균형)"

이 논문은 "흐름 맞추기"의 원리를 빌려와서 "심사위원 (평가 함수)"을 더 정확하게 훈련시키는 방법을 제안합니다.

1. 새로운 아이디어: "전체 그림이 아니라, '부분 스케치'로 평가하기"

기존에는 그림이 완성될 때까지 기다렸다가 점수를 매겼습니다. 하지만 이 논문은 "그림을 그리는 중간 과정 (부분 에피소드, Partial Episode)"에서도 흐름이 균형 잡혔는지 확인하자고 말합니다.

• 비유: 그림을 그릴 때, "초록색을 칠하는 단계"와 "파란색을 칠하는 단계"가 서로 조화를 이루는지 매 단계마다 확인하는 것입니다.
• 효과: 이렇게 하면 심사위원 (평가 함수) 이 훨씬 더 신뢰할 수 있는 점수를 매겨주게 됩니다. 등산할 때 나침반이 훨씬 안정적으로 작동하는 셈입니다.

2. 두 가지 큰 장점

① 더 유연한 훈련 (Backward Policy)

• 기존: 그림을 그릴 때 '뒤로 돌아가는 길 (Backward Policy)'은 고정된 규칙만 따를 수 있었습니다.
• 이 논문: "뒤로 돌아가는 길"도 함께 배우고 수정할 수 있게 했습니다.
• 비유: 그림을 그릴 때, "어떤 색을 먼저 칠할지" 정하는 규칙뿐만 아니라, "실수했을 때 어떻게 지우고 다시 칠할지"에 대한 규칙도 함께 배우게 되어, 훨씬 더 똑똑하고 유연하게 그림을 그릴 수 있게 됩니다.

② 오프라인 데이터 활용 (Offline Training)

• 기존: 새로운 그림을 그릴 때, 팀원들이 직접 그리는 것 (온라인) 만 쓸 수 있었습니다.
• 이 논문: 과거에 다른 팀이 그려둔 좋은 그림들 (오프라인 데이터) 을 가져와서 훈련할 수 있게 했습니다.
• 비유: 직접 모든 그림을 그릴 필요 없이, 도서관에 있는 명화집 (오프라인 데이터) 을 보고 배워도 된다는 뜻입니다. 이렇게 하면 시간과 비용을 아끼면서도 더 좋은 그림을 그릴 수 있습니다.

📊 실험 결과: "더 빠르고, 더 잘 그리는 팀"

연구팀은 이 방법을 다양한 테스트 (가상의 격자, 분자 설계, 생물학적 서열 등) 에 적용해 보았습니다.

• 결과: 기존 방법들보다 더 빠르게 수렴했고, 더 안정적이었습니다.
• 특히, 분자 구조나 복잡한 네트워크를 설계하는 거대한 작업에서도 이 방법이 가장 좋은 성능을 보여주었습니다.

🚀 요약

이 논문은 **"GFlowNet 이 복잡한 세상을 탐색할 때, 중간 과정의 흐름을 잘 맞춰주면 (Sub-EB), 더 신뢰할 수 있는 점수판 (평가 함수) 을 갖게 되고, 그 결과 더 빠르고 유연하게 최고의 결과물을 만들어낼 수 있다"**는 것을 증명했습니다.

마치 나침반이 고장 난 등산대에게 **정밀한 GPS(부분 구간 흐름 균형)**를 장착해 준 것과 같습니다. 이제 그들은 더 이상 길을 잃지 않고, 목적지에 더 빨리, 더 안전하게 도착할 수 있게 되었습니다.

기술 요약

논문 요약: EVALUATING GFLOWNET FROM PARTIAL EPISODES FOR STABLE AND FLEXIBLE POLICY-BASED TRAINING

이 논문은 생성 흐름 네트워크 (GFlowNet) 의 정책 기반 (Policy-based) 학습의 안정성과 유연성을 크게 향상시키는 새로운 방법론인 **부분 에피소드 평가 균형 (Subtrajectory Evaluation Balance, Sub-EB)**을 제안합니다. 저자는 기존 가치 기반 (Value-based) 방법론과 정책 기반 방법론 간의 이론적 간극을 해소하고, 평가 함수 (Evaluation Function) 를 학습하기 위한 새로운 목적 함수를 도입했습니다.

---

1. 문제 정의 (Problem)

GFlowNet 은 조합 공간 (그래프, 분자, 시퀀스 등) 에서 특정 점수 함수 $R(x)$에 비례하는 확률 분포를 샘플링하기 위한 생성 모델입니다. 학습의 핵심은 순방향 정책 ($\pi_F$) 과 역방향 정책 ($\pi_B$) 간의 불일치를 최소화하는 것입니다.

• 기존 접근법의 한계:
  • 가치 기반 (Value-based): 흐름 균형 (Flow Balance) 조건을 사용하여 학습합니다. (예: Sub-TB, DB). 이는 안정적이지만, 데이터 수집 정책 ($\pi_D$) 을 별도로 설계해야 하며, 역방향 정책 ($\pi_B$) 을 고정해야 하는 경우가 많아 유연성이 떨어질 수 있습니다.
  • 정책 기반 (Policy-based): 강화학습 (RL) 의 Actor-Critic 프레임워크를 차용하여 KL 발산을 직접 최소화하려 합니다. 여기서 'Critic' 역할을 하는 평가 함수 $V(s)$를 학습하는 것이 핵심인데, 기존 방법 (예: $\lambda$-TD) 은 편향 (Bias) 과 분산 (Variance) 의 트레이드오프 문제로 인해 $V(s)$의 신뢰성 있는 추정이 어렵고, 학습이 불안정해지는 문제가 있었습니다.
  • 핵심 과제: 신뢰할 수 있는 평가 함수 $V(s)$를 학습하여 정책 기반 학습의 안정성을 확보하고, 역방향 정책의 파라미터화 및 오프라인 데이터 활용을 가능하게 하는 것입니다.

2. 방법론 (Methodology)

저자는 상태 흐름 함수 (State Flow Function, $F(s)$) 와 평가 함수 ($V(s)$) 사이의 깊은 이론적 연결을 규명했습니다.

• 이론적 연결 (Theoretical Connection):

  • 최적의 흐름 균형 조건 하에서, 임의의 고정된 순방향 정책 $\pi_F$에 대해 상태 흐름 $F(s)$를 만족하는 해는 해당 상태에서의 정확한 KL 발산 (즉, 참된 평가 함수 $V^\dagger(s)$) 과 일치함을 증명했습니다.
  • 이를 바탕으로 부분 에피소드 평가 균형 (Sub-EB, Subtrajectory Evaluation Balance) 조건을 도출했습니다. 이는 부분 에피소드 (Subtrajectory) 단위로 흐름 균형을 요구하며, 학습된 발산의 차이가 실제 발산과 일치하도록 강제합니다.

• Sub-EB 목적 함수 (Sub-EB Objective):

  • $V(s)$를 학습하기 위해 새로운 목적 함수 $L_V(\phi)$를 제안했습니다.
  • 기존 $\lambda$-TD 방법이 단일 단계 (edge-wise) 오차와 특정 시작 시점의 이벤트에만 집중하는 반면, Sub-EB 는 **부분 에피소드 전체 (subtrajectory-wise)**의 불일치를 활용합니다.
  • 이는 학습 데이터의 시작과 끝을 모두 고려하여 더 균형 잡힌 학습을 가능하게 하며, 가중치 설정에 있어 더 유연합니다.

• 확장성 (Extensions):

  • 파라미터화된 역방향 정책: Sub-EB 는 역방향 정책 $\pi_B$를 고정하지 않고, 평가 함수 $V$와 함께jointly(공동) 업데이트할 수 있게 합니다. 이는 기존 $\lambda$-TD 기반 방법의 큰 제약 (2 단계 학습 필요) 을 해소합니다.
  • 오프라인 학습 (Offline Training): Sub-EB 조건을 역방향 관점 ($W$) 으로 확장하여, $\pi_F$와 다른 데이터 수집 정책 ($\pi_D$) 을 사용하여 오프라인 데이터를 활용한 정책 기반 학습을 가능하게 합니다.

3. 주요 기여 (Key Contributions)

1. 이론적 통합: 흐름 균형 조건 (Flow Balance) 과 평가 함수 (Evaluation Function) 간의 수학적 동치 관계를 증명하여, 가치 기반 방법의 원리를 정책 기반 학습에 적용할 수 있는 이론적 토대를 마련했습니다.
2. Sub-EB 목적 함수 제안: 신뢰할 수 있는 평가 함수 $V$를 학습하기 위한 새로운 목적 함수를 도입하여, 정책 기반 GFlowNet 학습의 안정성과 수렴 속도를 획기적으로 개선했습니다.
3. 유연성 증대:
   • 파라미터화된 역방향 정책 ($\pi_B$) 을 단일 단계 학습에 통합 가능하게 함.
   • 오프라인 데이터 수집 정책 ($\pi_D \neq \pi_F$) 을 활용한 정책 기반 학습을 가능하게 함.
4. 광범위한 실험 검증: 합성 데이터 (Hypergrid), 생물학적/분자 시퀀스 설계, 베이지안 네트워크 구조 학습, 분자 그래프 설계 등 다양한 실제 및 시뮬레이션 작업에서 기존 방법 (Sub-TB, RL, Q-Much 등) 대비 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

• Hypergrid (합성 데이터):
  • 256x256 및 3 차원 격자 환경에서 Sub-EB 기반 정책 기반 방법은 기존 RL 방법보다 더 빠른 수렴 속도와 더 높은 안정성을 보였습니다.
  • 파라미터화된 $\pi_B$를 사용할 때 Sub-EB-P 가 모든 방법 중 가장 우수한 성능을 기록했습니다.
• 시퀀스 및 분자 설계 (Sequence & Molecular Design):
  • SIX6, QM9, sEH 등 다양한 데이터셋에서 Sub-EB 는 분포 모델링 정확도 (DTV, DJSD) 와 모드 발견 (Mode Discovery) 능력 모두에서 경쟁력 있는 성능을 보였습니다.
  • 오프라인 기법 (Sub-EB-B) 을 적용했을 때, 고보상 상태 탐색 능력이 향상되었으나 분포 모델링 정확도는 일부 희생되는 트레이드오프가 관찰되었습니다.
• 베이지안 네트워크 (BN) 및 분자 그래프 설계:
  • 대규모 조합 공간 (최대 $10^{35}$개의 DAG) 에서 Sub-EB 는 높은 평균 보상과 빠른 수렴을 달성했습니다.
  • 특히 분자 그래프 설계 (LogP, JNK3) 에서 Sub-EB 는 RL 과 Q-Much 대비 더 높은 보상을 얻으면서도 적절한 다양성을 유지했습니다.

5. 의의 및 결론 (Significance)

이 논문은 GFlowNet 학습의 패러다임을 변화시키는 중요한 기여를 했습니다.

• 안정성 확보: 정책 기반 학습이 가진 불안정성 문제를 해결하고, 이를 가치 기반 방법의 강점 (흐름 균형) 과 결합하여 신뢰할 수 있는 학습 프레임워크를 제시했습니다.
• 유연한 학습 환경: 역방향 정책의 동적 적응과 오프라인 데이터 활용을 가능하게 함으로써, 실제 응용 (데이터 수집 비용이 높은 시나리오 등) 에 GFlowNet 을 적용하는 데 있어 실용적인 장벽을 낮췄습니다.
• 미래 지향성: Sub-EB 는 단순한 정책 기반 방법을 넘어, 더 복잡한 강화학습 알고리즘 (예: TRPO) 과의 통합 및 최적 가중치 설계 등 향후 연구의 방향성을 제시합니다.

결론적으로, 이 연구는 GFlowNet 의 이론적 기반을 강화하고, 실제 복잡한 조합 최적화 문제 해결을 위한 강력하고 유연한 도구로 발전시켰다는 점에서 의의가 큽니다.