Routing by Reaching: Composition of Pre-trained GFlowNets for Multi-Objective Generation

Motivation

• Known: GFlowNet은 보상 함수에 비례하는 분포에서 다양한 후보를 샘플링하는 생성 모델로, 단일 목적(single objective) 설정에서 잘 연구되어 있다. 다중 목적 설정으로 확장하기 위해 preference-conditioned GFlowNet(MOGFN, HN-GFN) 같은 scalarization 기법과 compositional sculpting 같은 logical operator 기반 기법이 제안되어 왔다.
• Gap: 기존 방법들은 목적 조합이 바뀔 때마다 처음부터 재학습(scalarization의 경우 새 preference-conditioned 모델, logical operator의 경우 새 보조 classifier)이 필요해 오버헤드가 크며, scalarization과 logical composition을 하나의 통합된 프레임워크로 다루지 못한다는 두 가지 근본적 한계가 있다.
• Why: 실제 과학적 발견 문제는 목적 집합이 사전에 고정되어 있지 않고 탐색 과정에서 자주 변경되므로, 재학습 없이 사전 학습된 GFlowNet들을 즉시 조합할 수 있는 방법은 실용적 비용을 크게 줄이고 다양한 조합 스킴(선형·비선형)을 단일 프레임워크로 지원할 수 있어 중요하다.
• Approach: 각 사전 학습된 GFlowNet의 상태별 reaching probability를 자연스러운 가중치로 사용하여 여러 모델의 forward policy를 추론 시점에 혼합(mixing)하는 방법을 제안하고, 이 혼합 정책이 유도하는 분포를 이론적으로 분석한다.

Achievement

Figure 3. Distortion factor δ(x) = uM(x)/NM(x) vs. unnormalized target density G(p1(x), . . . , pk(x)) on the 2D grid do

1. 정확한 분포 복원: 선형 스칼라화(linear scalarization)의 경우, 제안된 mixing policy가 목표 분포를 정확히(exactly) 유도함을 이론적으로 증명하였다.
2. 비선형 연산자에 대한 근사 품질 정량화: harmonic mean, contrast 등 비선형 logical operator에 대해서는 distortion factor δ(x)=u_M(x)/N_M(x)를 도입하여 유도 분포와 목표 분포 간 L1 거리를 정량적으로 분석하고, 고밀도(high-density) 영역에서 근사가 정확함을 실험적으로 보였다.
3. 학습 없는(training-free) 성능: 2D grid 합성 도메인과 fragment/atom 기반 분자 생성 실험에서, scalarization의 경우 preference-conditioned 베이스라인 대비 낮은 오차를 목적 수가 늘어나도 안정적으로 유지했고, logical operator의 경우 classifier-guided 방식과 비슷한 성능을 추가 학습 없이 달성하면서 추론 속도가 크게 빨랐다.

How

Figure 3. Distortion factor δ(x) = uM(x)/NM(x) vs. unnormalized target density G(p1(x), . . . , pk(x)) on the 2D grid do

• 각 사전 학습된 GFlowNet에 대해 상태 s의 reaching probability u(s) (s0에서 s로 도달할 확률의 총합)를 정의하고 계산
• 여러 GFlowNet의 forward policy p_F를 각 모델의 reaching probability로 가중하여 혼합된 forward policy를 구성, 이를 통해 하나의 mixing policy가 여러 목적을 동시에 반영하도록 함
• 선형 스칼라화(가중합 보상)의 경우 이 mixing policy가 유도하는 분포가 목표 분포 R(x)^β/Z와 정확히 일치함을 증명
• harmonic mean, contrast와 같은 비선형 logical operator의 경우 목표 분포와 유도 분포 간 L1 거리를 분석하고, distortion factor δ(x)=u_M(x)/N_M(x)를 정의하여 오차의 크기를 정량화
• 2D grid 합성 도메인 및 fragment-based, atom-based 분자 생성(SEH, QED 등) 태스크에서 preference-conditioned 베이스라인(MOGFN, HN-GFN) 및 compositional sculpting과 비교 실험 수행

Originality

• 기존에는 scalarization과 logical operator를 별도로 다루었으나, 본 논문은 reaching probability 기반 forward policy mixing이라는 단일 메커니즘으로 두 방식을 통합적으로 처리
• 사전 학습된 GFlowNet들을 추가 학습(fine-tuning, auxiliary classifier 학습) 없이 추론 시점에서 바로 조합한다는 점에서, GFlowNet 분야에서 model composition을 training-free 방식으로 최초로 제안
• 선형 스칼라화에 대한 정확한 분포 복원 증명과, 비선형 연산자에 대한 distortion factor라는 새로운 근사 오차 척도를 제시하여 이론적 근거를 뒷받침

Limitation & Further Study

• 비선형 logical operator의 경우 정확한 분포 복원이 아니라 근사(approximation)에 그치며, distortion factor가 저밀도(low-density) 영역에서는 오차가 커질 수 있음이 시사됨
• 실험이 합성 2D grid와 분자 생성(fragment/atom 기반)에 한정되어 있어, 더 복잡한 실제 과학적 발견 문제(예: 단백질 설계 등)에서의 확장성 검증이 필요
• 조합하려는 GFlowNet 수가 매우 많아질 때 reaching probability 계산 및 mixing 과정의 계산 비용 증가 여부에 대한 분석이 제한적
• 후속 연구로 더 다양한 비선형 연산자에 대한 이론적 오차 상한 개선, 그리고 대규모 목적 집합에서의 확장성 검증이 필요

Evaluation

총평: 사전 학습된 GFlowNet의 재사용성을 크게 높이는 실용적이고 이론적으로 탄탄한 training-free composition 프레임워크로, scalarization과 logical operator를 통합한 점이 돋보이는 좋은 연구이다. 다만 비선형 연산자에서의 근사 오차와 대규모 확장성에 대한 추가 검증이 향후 과제로 남는다.