GFlowState: Visualizing the Training of Generative Flow Networks Beyond the Reward

1. 문제 정의 (Problem)

생성 흐름 네트워크 (Generative Flow Networks, GFlowNets) 는 보상 함수에 비례하여 샘플을 생성하는 확률적 프레임워크로, 분자 및 신약 개발, 신소재 발견 등 과학적 발견 분야에서 강력한 도구로 자리 잡았습니다. 그러나 GFlowNet 의 훈련 과정은 다음과 같은 이유로 해석과 디버깅이 매우 어렵습니다.

• 복잡한 훈련 역학: 기존 머신러닝 모니터링 도구 (TensorBoard 등) 는 손실 (loss) 이나 보상 (reward) 과 같은 집계된 지표를 추적할 수는 있지만, 모델이 샘플 공간을 어떻게 탐색하는지, 샘플 궤적 (trajectories) 을 어떻게 구성하는지, 훈련 중 샘플링 확률이 어떻게 변화하는지에 대한 구조적 통찰력을 제공하지 못합니다.
• 조합적 공간의 규모: GFlowNet 이 다루는 상태 공간 (state space) 은 조합적으로 매우 커서, 수동으로 검사하거나 기존 시각화 도구로는 전체 구조를 파악하는 것이 불가능합니다.
• 모드 붕괴 (Mode Collapse) 및 편향: 모델이 특정 고보상 영역만 반복적으로 샘플링하여 다양성이 떨어지거나, 훈련 초기에 발견된 영역에 갇히는 현상을 식별하고 수정하기 어렵습니다.

이러한 한계를 극복하기 위해 GFlowNet 의 훈련 구조와 역학을 관찰 가능하게 만드는 전용 시각 분석 도구의 필요성이 대두되었습니다.

2. 방법론 (Methodology)

저자들은 GFlowNet 개발자 (Mila 연구소) 와의 긴밀한 협력을 통해 훈련 분석을 위한 핵심 작업을 도출하고, 이를 해결하기 위한 대화형 시각 분석 시스템 GFlowState 를 설계했습니다.

A. 주요 분석 작업 (User Tasks)

GFlowState 는 다음 세 가지 범주의 작업을 지원합니다:

1. 샘플링 궤적 분석 (Traj): 모델이 각 단계에서 내리는 결정과 궤적 흐름을 추적.
2. 생성된 샘플 분석 (Samp): 보상 함수에 비례한 샘플링이 이루어지는지 검증, 고/저 보상 영역 식별, 이상치 탐지.
3. 훈련 역학 분석 (Dyn): 훈련 중 전이 확률의 변화, 샘플 공간 탐색 동향, 모드 붕괴 감지.

B. 시스템 아키텍처 및 4 가지 주요 뷰 (Views)

GFlowState 는 훈련 데이터 (샘플, 궤적, 검증 세트) 를 기반으로 다음 4 가지 상호 연결된 시각화 뷰를 제공합니다:

1. Sample Ranking (샘플 순위):

   • 배임 차트 (Bump chart) 를 사용하여 각 훈련 반복 (iteration) 별 가장 높은/낮은 보상을 가진 객체들의 순위 변화를 시각화합니다.
   • 새로운 고보상 영역 (모드) 의 발견 시점이나 순위의 급격한 변화를 통해 모드 붕괴나 새로운 탐색 성공을 식별합니다.

2. State Projection (상태 투영):

   • 고차원 샘플 데이터를 2 차원 (UMAP 등) 으로 축소하여 육각형 격자 (Hexbin) 로 집계합니다.
   • 메트릭: 보상 - 확률 상관관계, 손실 (loss), 검증 세트 커버리지 (Validation Set Coverage) 등을 색상으로 인코딩하여 모델이 샘플링하지 않은 영역이나 학습이 부족한 영역을 식별합니다.

3. DAG View (방향 비순환 그래프 뷰):

   • 생성된 모든 샘플 궤적을 기반으로 방향 비순환 그래프 (DAG) 를 구성합니다.
   • 전체 그래프를 한 번에 보여주는 대신, 루트 노드에서부터 상호작용적으로 확장 (Interactive Expansion) 하거나 다른 뷰에서 선택된 샘플의 경로를 하이라이트하여 국소적인 하위 그래프를 탐색할 수 있게 합니다.
   • 노드와 엣지는 상태와 전이 빈도/확률 등을 기반으로 시각화됩니다.

4. Transition Heatmap (전이 히트맵):

   • 훈련 기간 동안의 전이 (transition) 확률 변화를 시간 축 (Iteration) 에 따라 시각화합니다.
   • 특정 전이의 확률 변화 추이를 분석하여 모델이 어떤 경로를 선호하게 되었는지, 또는 어떤 경로가 사라졌는지를 파악합니다.

C. 구현 및 통합

• 기술 스택: Python 기반의 Plotly Dash 를 사용하여 구현되었으며, SQL 데이터베이스를 통해 훈련 데이터를 저장하고 관리합니다.
• 확장성: 사용자 정의 환경 (Grid, 분자, 결정 구조 등) 에 쉽게 통합되도록 4 가지 핵심 함수 (데이터베이스 형식, 특징 표현, 단일/다중 상태 시각화) 만 구현하면 됩니다.

3. 주요 기여 (Key Contributions)

1. GFlowNet 훈련 분석 작업 정의: GFlowNet 개발자가 훈련 품질을 평가하고 디버깅하기 위해 필요한 구체적인 분석 작업 (궤적, 샘플, 역학 분석) 을 체계적으로 도출했습니다.
2. GFlowState 시스템 개발: GFlowNet 의 훈련 역학을 시각화하고 상호작용적으로 탐색할 수 있는 최초의 프로토타입 시스템을 제안했습니다.
3. 다양한 도메인에서의 사례 연구: 그리드 환경 (Grid Environment) 과 결정 구조 생성 (Crystal Structure Generation) 두 가지 실제 응용 분야에서 시스템의 유효성을 입증했습니다.

4. 결과 (Results)

두 가지 사례 연구를 통해 GFlowState 의 유용성이 입증되었습니다:

• 그리드 환경 (Grid Environment):

  • 모드 발견 식별: Sample Ranking 뷰를 통해 훈련 중 4,500 번째 반복 부근에서 새로운 고보상 모드가 발견된 시점을 정확히 파악했습니다.
  • 학습 과정 추적: State Projection 과 Transition Heatmap 을 통해 모델이 초기에는 시작점에 가까운 모드만 샘플링하다가, 훈련이 진행됨에 따라 다른 고보상 모드로 탐색을 확장하는 과정을 시각적으로 확인했습니다.
  • DAG 분석: DAG View 를 통해 특정 모드로 가는 전이 확률이 시간이 지남에 따라 감소하는 것을 확인하여, 모델이 단일 모드에 갇히지 않고 다양성을 확보하고 있음을 입증했습니다.

• 결정 환경 (Crystals Environment):

  • 실제 적용: 백금과 팔라듐으로 구성된 결정 구조 생성 과제에서 GFlowState 를 적용했습니다.
  • 전문가 피드백: 머신러닝 및 재료 과학 전문가들은 Sample Ranking 을 통해 훈련 동향을 빠르게 파악하고, DAG View 를 통해 생성된 궤적의 타당성을 검증하는 데 큰 도움을 받았습니다.
  • 하위 환경 의존성 분석: 공간 군 (Space group) 이 생성된 후 조성이 결정되는 과정과 같은 복잡한 의존 관계를 시각적으로 분석할 수 있어, 물리적으로 타당한 결정 생성 모델 개발에 기여했습니다.

5. 의의 및 중요성 (Significance)

• 해석 가능성 (Interpretability) 향상: GFlowNet 의 "블랙박스" 성격을 탈피하여, 모델이 어떻게 샘플 공간을 탐색하고 보상을 학습하는지 구조적으로 관찰할 수 있게 합니다.
• 디버깅 및 최적화 가속: 훈련 실패 원인 (예: 모드 붕괴, 특정 영역의 미탐색) 을 조기에 발견하고 하이퍼파라미터 조정이나 모델 수정에 직접적인 통찰력을 제공합니다.
• 과학적 발견 지원: GFlowNet 이 활용되는 신약 개발, 신소재 설계 등 고위험/고비용의 과학적 발견 분야에서 모델의 신뢰성을 높이고 개발 주기를 단축하는 데 기여합니다.
• 시각 분석의 확장: 기존 머신러닝 시각화 도구가 특정 모델 클래스 (예: 분류, 강화학습) 에 국한되었던 점을 넘어, GFlowNet 과 같은 생성적 흐름 모델에 특화된 시각 분석 패러다임을 제시했습니다.

결론적으로, GFlowState 는 GFlowNet 의 훈련 과정을 투명하게 만들어 개발자들이 더 효과적으로 모델을 이해하고 개선할 수 있도록 돕는 필수적인 도구로 평가됩니다.