기술 요약: 프록시를 넘어선 오프라인 GFlowNet 훈련을 위한 궤도 증류 기반 안내

문제 제기

생성 흐름 네트워크 (GFlowNets) 는 보상과 비례하는 구성 객체를 샘플링하기 위한 강력한 모델입니다. 그러나 분자 발견, 생물학적 시퀀스 설계, 추천 시스템과 같은 많은 실제 응용 분야에서는 높은 실험 비용이나 인간 평가의 필요성으로 인해 보상을 조회하기 위해 환경과 능동적으로 상호작용하는 것이 불가능합니다. 결과적으로 GFlowNet 은 정적 오프라인 데이터셋을 사용하여 훈련되어야 합니다.

현재의 오프라인 훈련 패러다임은 일반적으로 데이터셋에서 보상 신호를 근사화하는 프록시 보상 모델을 훈련하고, GFlowNet 이 정책 최적화를 위해 이 모델을 조회하도록 의존합니다. 이 접근 방식은 다음과 같은 중대한 도전 과제에 직면해 있습니다:

1. 데이터 부족과 비용: 신뢰할 수 있는 프록시 모델을 구축하려면 종종 크고 다양한 데이터셋과 상당한 도메인 전문 지식이 필요합니다.
2. 오류 전파: 분포 외 (OOD) 샘플이나 프록시가 부정확하게 추정한 상태에 대한 보상 조회는 그래디언트를 통해 오류를 전파하여 학습된 정책을 훼손할 수 있습니다.
3. 프록시 없는 방법의 한계: 기존 프록시 없는 접근 방식 (예: RO-GFlowNets, COFlowNet) 은 종종 정책을 데이터셋과 정렬시키기 위해 거친 제약을 부과합니다. 이러한 제약은 일반화를 제한하고, 효과적인 탐색을 억제하며, 데이터 품질에 대한 민감도를 높일 수 있습니다.

해결된 핵심 질문은 다음과 같습니다: 종단 보상이 신뢰할 수 있는 프록시 모델을 구축하기에 충분하지 않을 때, 어떻게 오프라인 궤도를 효과적으로 활용하여 GFlowNet 을 훈련할 수 있는가?

방법론: 궤도 증류 기반 GFlowNet (TD-GFN)

저자들은 TD-GFN을 제안했습니다. 이는 프록시 보상 모델에 의존하지 않고 환경의 방향성 비순환 그래프 (DAG) 에서 풍부한 구조적 정보를 추출하여 정책 훈련을 안내하는 새로운 프록시 없는 프레임워크입니다. 방법론은 세 가지 주요 단계로 구성됩니다:

1. 역강화학습 (IRL) 을 통한 엣지 보상 추출

TD-GFN 은 종단 보상을 예측하는 대신, IRL 을 활용하여 오프라인 궤도에서 전이 수준의 엣지 보상($R_E$)을 증류합니다.

• 재균형: 오프라인 데이터가 전문가 행동을 대표하지 않을 수 있는 문제를 해결하기 위해, 저자들은 재균형 전략을 도입합니다. 궤도는 종단 보상에 비례하는 확률 ($P(\tau) \propto R(s_T)$) 로 샘플링되어 가짜 전문가 데이터셋을 생성합니다.
• 적대적 훈련: 수정된 GAIL(Generative Adversarial Imitation Learning) 프레임워크를 사용하여 판별기 $D_\phi$가 재균형된 데이터셋의 전이와 정책 $\pi_\psi$가 생성한 전이를 구별합니다.
• 보상 유도: 판별기에서 편향 없는 엣지 수준 보상 함수가 추출됩니다:
  $$R_E(s, s') = \log D_\phi(s, s') - \log(1 - D_\phi(s, s'))$$
  이론적으로 이 보상은 특정 전이가 효과적인 정책을 학습하는 데 얼마나 중요한지를 정량화하는 로그 조건부 유입 점수 ($\log P_B(s|s')$) 를 근사합니다. 중요한 점은 이러한 보상이 종단 예측이 아닌 밀집된 중간 단계 안내를 제공한다는 것입니다.

2. 보상 안내 DAG 가지치기

불완전한 IRL 추정치로 인한 오류 전파를 방지하고 견고성을 확보하기 위해, TD-GFN 은 엣지 보상을 사용하여 정책 그래디언트를 직접 형성하지 않습니다. 대신 이를 사용하여 환경 DAG 를 가지치기합니다.

• 임계값 설정: 보상 분포의 평균에서 $K$배의 표준 편차를 뺀 통계적 임계값 미만의 보상을 가진 엣지는 제거됩니다.
• 이론적 보장: 저자들은 이 가지치기 전략이 높은 엣지 점수를 가진 경로를 통해 도달 가능한 모든 "고-병목" 모드 (종단 상태) 를 보존하면서 저-유틸리티 전이를 제거함을 증명합니다. 이는 고가치 궤도를 우선시하고 IRL 근사 오류로부터 정책을 보호하는 컴팩트하고 표현력 있는 행동 공간을 생성합니다.

3. 우선순위 역방향 샘플링

훈련은 가지치기된 서브그래프에서 우선순위 역방향 샘플링 메커니즘을 사용하여 수행됩니다.

• 샘플링 전략: 종단 상태는 데이터셋과 가지치기된 그래프의 교집합에서 실제 보상 (ground-truth rewards) 에 비례하는 확률로 샘플링됩니다.
• 역방향 궤도 구성: 각 샘플링된 종단에서, 학습된 엣지 보상으로 정의된 학습된 역방향 정책 $P_B$를 사용하여 궤도가 루트 방향으로 역방향으로 구성됩니다:
  $$P_B(s_t | s_{t+1}) = \frac{\exp\{R_E(s_t, s_{t+1})\}}{\sum \exp\{R_E(s, s_{t+1})\}}$$
• 그래디언트 업데이트: GFlowNet 정책은 이러한 샘플링된 궤도와 실제 종단 보상을 사용하여 최적화됩니다. 그래디언트 업데이트를 위해 기록된 종단 보상에만 의존함으로써, 이 방법은 프록시 기반 방법 고유의 오류 전파 문제를 피합니다.

주요 기여

1. 새로운 프록시 없는 패러다임: TD-GFN 은 프록시 모델과 OOD 보상 조회에 대한 의존성을 제거합니다. 재균형 데이터셋에서 전이 수준 엣지 보상을 증류하기 위해 IRL 을 활용하여 정책 훈련과 일반화를 위한 구조화된 안내를 제공합니다.
2. 견고한 안내 메커니즘: 이 프레임워크는 DAG 가지치기와 우선순위 역방향 샘플링이라는 두 가지 간접 안내 메커니즘을 도입합니다. 이 설계는 그래디언트 업데이트가 실제 보상에만 의존하도록 보장하여 오류 전파를 방지하면서도 효율적인 탐색을 위한 밀집된 구조적 안내를 활용합니다.
3. 최첨단 성능: 광범위한 실험을 통해 TD-GFN 이 프록시 기반 방법 (Proxy-GFN), 오프라인 RL 알고리즘 (CQL, IQL), 그리고 다른 프록시 없는 GFlowNet 변형 (COFlowNet, Dataset-GFN) 을 포함한 다양한 베이스라인에 비해 우수한 수렴 속도와 샘플 품질을 달성함을 입증했습니다.

실험 결과

저자들은 TD-GFN 을 세 가지 벤치마크에서 평가했습니다:

• 하이퍼그리드 (Synthetic): 84 차원 그리드에서 희소하고 다중 모달 보상에 대해, TD-GFN 은 5,000 회 미만의 상태 방문으로 16 개의 모든 보상 모드를 식별했습니다 (베이스라인 대비 6 배 개선). 또한 가장 우수한 경쟁자 (COFlowNet) 보다 두 배 빠르게 가장 낮은 경험적 L1 오차로 수렴했습니다. 노이즈가 있는 데이터 (Mixed dataset) 와 데이터 부족 (데이터의 10%) 하에서도 강력한 견고성을 보여주었습니다.
• 생물학적 시퀀스 설계 (항미생물 펩타이드): TD-GFN 은 상위 100 개 시퀀스 중에서 가장 높은 보상과 가장 큰 다양성을 가진 시퀀스를 생성하여 프록시 기반 방법 및 기타 오프라인 방법 모두를 능가했습니다.
• 분자 설계: TD-GFN 은 베이스라인보다 더 많은 고보상 분자와 다양한 모드를 발견했으며, 수렴에 필요한 궤도 수가 현저히 적었습니다 (Proxy-GFN 의 43.7 × 10^4 대비 2.75 × 10^4). 실제 보상을 직접 사용하는 Oracle-GFN 과 고보상 분자 식별 측면에서 동등한 성능을 달성했습니다.

의의 및 주장

이 논문은 TD-GFN 이 오프라인 GFlowNet 훈련을 위한 보다 견고하고 안정적이며 효율적인 패러다임을 확립한다고 주장합니다. IRL 을 통해 오프라인 궤도에 내재된 구조적 정보를 활용함으로써, 이 방법은 제한된 오프라인 데이터셋과 생성 모델링의 복잡한 탐색 요구 사항 사이의 간극을 메웁니다.

저자들은 TD-GFN 이 다음과 같음을 강조합니다:

• 잠재력 개방: 능동적 보상 조회가 불가능한 복잡한 분야에서 GFlowNet 의 배포를 가능하게 합니다.
• 오류 전파 방지: 안내 메커니즘 (엣지 보상) 과 그래디언트 업데이트 메커니즘 (종단 보상) 을 분리함으로써 프록시 모델의 함정을 피합니다.
• 탐색 강화: 밀집된 구조적 안내를 통해 모델이 훈련 데이터에서 희소하게 표현되거나 관찰되지 않더라도 고보상 영역을 효과적으로 탐색할 수 있게 합니다.

이 연구는 프록시 기반 방법에서 종종 간과되는 궤도 수준의 정보가 효과적인 오프라인 GFlowNet 훈련의 초석임을 시사하며, 신약 개발, 생물학적 설계, 조합 최적화와 같은 실제 응용 분야를 위한 확장 가능한 솔루션을 제공합니다.