CMAB GFN 다중 고득점 서브스페이스 탐색으로 GFlowNet 효율 극대화

초록

본 논문은 GFlowNet의 탐색 효율을 높이기 위해 행동을 조합형 다중 무장 밴딧(CMAB)으로 모델링하고, 저품질 행동을 사전에 차단해 고득점 서브스페이스만을 집중 탐색한다. 제한된 샘플링 예산에서도 높은 보상 후보와 다양성을 동시에 확보한다는 실험 결과를 제시한다.

상세 분석

CMAB‑GFN은 기존 GFlowNet이 전체 상태공간을 거의 균등하게 탐색하면서 저보상 영역에 과도하게 자원을 소모하는 문제를 근본적으로 재구성한다. 핵심 아이디어는 “행동을 기본 팔(arm)으로 보고, K개의 팔을 선택해 하나의 슈퍼 팔(super‑arm)을 구성한다”는 CMAB 프레임워크를 도입하는 것이다. 여기서 기본 팔은 상태에 독립적인 원시 행동 ai (예: 분자 설계에서의 빌딩 블록, RNA 설계에서의 뉴클레오타이드)이며, 슈퍼 팔은 선택된 ai 집합 S 이다. 슈퍼 팔 S 가 선택되면, 모든 상태 s 에서 정책은 (ai∈S) 인 행동만 허용하도록 제한된다. 즉, 저품질 ai 를 사전에 차단함으로써 원래의 거대한 탐색 공간을 고득점 후보가 밀집된 “서브스페이스”로 압축한다.

보상 할당은 전통적인 CMAB의 반쪽 밴딧(semi‑bandit) 형태를 차용한다. 각 라운드 t 에서 제한된 서브스페이스 S 내에서 생성된 후보 x 들의 정규화 보상 r(x) 를 평균내어 해당 후보에 포함된 각 기본 팔 i 에 대한 추정값 Xᵗᵢ 를 계산한다. 정규화는 전역 최소·최대 보상을 실시간으로 업데이트하는 min‑max 스케일링으로 수행한다. 이렇게 얻은 Xᵗᵢ 는 CMAB 알고리즘(예: Upper Confidence Bound, Thompson Sampling 등)의 입력으로 사용되어 다음 라운드에 선택할 슈퍼 팔을 결정한다.

학습 과정은 두 단계로 나뉜다. 첫 번째 단계는 선택된 서브스페이스 S 에 제한된 상태‑전이 그래프에서 GFlowNet을 훈련시켜 흐름 보존(flow‑matching) 손실을 최소화한다. 두 번째 단계는 제한을 해제하고 자유롭게 후보를 샘플링하여 Xᵗᵢ 를 업데이트한다. 제한된 훈련은 네트워크 구조와 파라미터 업데이트 비용을 크게 증가시키지 않으며, 비제한 평가 단계는 멀티스레드 샘플링을 활용해 오버헤드를 최소화한다.

또한, 기본 팔의 수가 적고 시퀀스가 길어지는 경우(예: 비트 시퀀스 생성)에는 “짧은 행동 시퀀스”를 새로운 기본 팔로 정의한다. 길이 t 의 연속 행동을 하나의 팔로 묶어 슈퍼 팔을 구성하면, 깊고 얇은 네트워크 구조가 넓은 탐색 공간으로 변환되어 차단에 따른 솔루션 붕괴 위험을 완화한다.

동적 특성 측면에서, GFlowNet 자체가 학습 과정에서 보상 분포를 변화시키기 때문에 전통적인 정적 CMAB 가정이 깨진다. 논문은 이러한 비정상성을 감안해, 흐름 네트워크가 충분히 다양해질 때까지 Xᵗᵢ 추정이 수렴하도록 설계된 두 단계 샘플링 전략을 제안한다. 실험에서는 분자 설계, RNA 서열 설계, 비트 시퀀스 생성 등 세 가지 벤치마크에서 CMAB‑GFN이 기존 온도 스케일링, Q‑value 통합, 로컬 서치 기반 방법보다 높은 평균 보상과 빠른 수렴 속도를 보였으며, 특히 다양한 고득점 서브스페이스를 순환 탐색함으로써 후보 다양성도 유지함을 확인했다.

요약하면, CMAB‑GFN은 GFlowNet의 탐색을 “행동‑레벨”이 아닌 “서브스페이스‑레벨”로 전환함으로써, 저품질 행동을 효율적으로 차단하고, 고득점 영역을 집중 탐색하면서도 다양한 서브스페이스를 순환해 모드 붕괴를 방지한다. 이는 GFlowNet을 실제 대규모 조합 최적화 문제에 적용할 때 샘플링 효율과 결과 품질을 동시에 향상시킬 수 있는 강력한 설계 패러다임을 제공한다.