GFlowNet Training by Policy Gradients

당신이 완벽한 레시피, 가장 효율적인 배송 경로, 혹은 새로운 약물 분자를 찾아야 하는 세상을 상상해 보십시오. 하지만 가능한 조합의 수가 너무 방대하여 이를 하나씩 모두 확인하려면 우주의 나이보다 더 긴 시간이 걸릴 수도 있습니다. 이것이 바로 생물학에서 공학에 이르기까지 모든 분야를 괴롭히는 "조합 폭발(combinatorial explosion)"이라는 문제입니다. 이를 해결하기 위해 과학자들은 **GFlowNet(Generative Flow Network)**이라는 영리한 도구를 사용합니다. GFlowNet을 딱딱한 규칙 책이 아니라, 마법 같은 물의 흐름 시스템이라고 생각해보십시오. 이것은 강물이 협곡을 깎아 길을 만드는 것처럼, 복잡한 객체들을 단계별로 구축합니다. 목표는 "물"(또는 확률)이 흘러가서, 가장 아름답고 보상이 높은 골짜기(최적의 솔루션)에 그 골짜기가 마땅히 누려야 할 만큼 정확하게 도달하도록 만드는 것입니다.

전통적으로 이 물 시스템이 올바르게 흐르도록 가르치는 것은 거대하고 보이지 않는 저울의 균형을 맞추는 것과 같았습니다. "가치 기반(value-based)" 접근법으로 알려진 기존 방식은 모든 지점에서의 수위가 특정 방정식과 일치하는지 확인하는 데 집중합니다. 이는 마치 배관공이 누수가 없는지 확인하기 위해 모든 파이프의 압력을 끊임없이 측정하는 것과 비슷합니다. 이 방식도 작동은 하지만, 특히 고립된 고보상 정점들이 존재하여 도달하기 어려운 지형에서는 느리고 서투를 수 있습니다. 연구진은 다음과 같은 질문을 던졌습니다. 모든 지점의 압력을 체크하는 대신, 단순히 경로가 지나가는 길에 보상을 주는 방식으로 이 물 시스템을 가르칠 수는 없을까? 그들은 이 네트워크를 훈련시키는 새로운 방법을 제안했는데, 이는 수학자가 방정식을 푸는 방식이라기보다 비디오 게임 캐릭터가 점수를 모으며 미로를 달리는 법을 배우는 방식에 더 가깝습니다.

흐름을 훈련하는 새로운 방법

저자인 Puhua Niu와 그의 팀은 GFlowNet의 훈련 초점을 "수학을 확인하는 것"에서 "보상을 따르는 것"으로 전환하는 신선한 훈련법을 개발했습니다. 기존의 사고방식에서는 네트워크가 전체 지도에서 물의 흐름을 균형 있게 유지하도록 훈련했는데, 이는 전통적인 강화 학습(RL)이 모든 상태의 가치를 추정하던 방식과 유사합니다. 그러나 새로운 접근 방식은 훈련 과정을 직접적인 정책 경사(policy gradient) 문제로 취급합니다.

이를 이해하기 위해, 당신이 개에게 공 가져오기 훈련을 시키고 있다고 상상해 보십시오.

• 기존 방식 (가치 기반): 당신은 마당의 모든 가능한 지점에 서서 그 지점이 가진 "공 가져오기의 가치"를 정확히 계산합니다. 그런 다음 수학적으로 모든 위치에서 완벽하게 맞아떨어지도록 개의 행동을 조정합니다. 매우 정밀하지만, 모든 풀잎의 가치를 계산하는 데 엄청난 정신적 에너지가 필요합니다.
• 새로운 방식 (정책 기반): 당신은 그저 개가 공을 향해 달려갈 때 "착하다!"라고 말하고, 잘못된 방향으로 갈 때 "안 돼!"라고 말하면 됩니다. 당신은 마당의 모든 지점의 가치를 알 필요가 없습니다. 그저 경로를 따라 얻는 보상을 바탕으로 개의 달리기 스타일을 조정할 뿐입니다.

이 논문은 현재 네트워크가 사용하는 전략(또는 정책)에 의존하는 특별한 종류의 "보상"을 도입합니다. 이를 통해 그들은 GFlowNet의 복잡한 흐름 균형 방정식과 현대 AI의 더 단순하고 직접적인 "보상과 처벌" 학습 스타일 사이의 간극을 메웁니다. 연구진은 이 방법이 네트워크가 훨씬 더 빠르고 견고하게 학습할 수 있게 하며, 특히 "보물"(고보상 솔루션)이 찾기 힘든 고립된 곳에 숨겨져 있을 때 더욱 그렇다는 것을 발견했습니다.

그들이 발견한 것과 피한 것

연구진은 격자 시뮬레이션(거대한 체스판 같은), 생물학적 서열 설계(DNA 가닥 같은), 분자 구조 생성(새로운 약물 같은)을 포함한 여러 가지 도전 과제에 대해 이 새로운 "보상 기반" 훈련을 테스트했습니다.

이 시뮬레이션에서 그들의 새로운 방법인 RL-G(스마트한 가이드 사용 시)와 RL-T(변화를 안전하게 유지하기 위한 "신뢰 영역" 사용 시)는 기존 방식들을 일관되게 능가했습니다.

• 속도: 새로운 방법들은 훨씬 더 빠르게 수렴(솔루션 발견)했습니다. 256x256 격자 실험에서, 새로운 방법들은 기존의 "궤적 균형(Trajectory Balance, TB)" 방식보다 더 적은 단계만에 낮은 오차율에 도달했습니다.
• 정확도: 최종 결과 또한 종종 더 정확했습니다. 예를 들어, 256x256 격자에서 그들의 최선책인 RL-G는 약 0.439의 총 변동(Total Variation) 오차를 달ian 반면, 그다음으로 우수한 전통적 방식인 TB-U는 약 0.728이었습니다. 분자 설계 테스트에서도 그들의 방법은 기존 방식보다 더 많은 독특한 "모드(mode)"(다양한 고품질 솔루션)를 찾아냈습니다.

결정적으로, 이 논문은 우리가 항상 복잡한 오프-폴리시 샘플러(Thompson Sampling이나 무작위 혼합 같은 것들)에 의존해야 한다는 생각에 반대합니다. 이러한 방법들은 "탐험(새로운 것을 시도함)"과 "활용(작동하는 것을 사용함)" 사이의 균형을 맞추려 노력하지만, 저자들은 강력한 경사 추정을 결합한 자신들의 정책 기반 접근 방식을 사용함으로써, 복잡한 외부 기술 없이도 네트워크가 자연스럽게 최적의 경로를 찾을 수 있음을 보여줍니다. 그들은 단순히 이렇게 제안만 한 것이 아니라, 여러 데이터셋에 걸쳐 이를 측정하여 새로운 전략이 이러한 네트워크를 훈련시키는 더 안정적이고 효율적인 방법을 제공한다는 것을 입증했습니다.

"신뢰 영역"과 "가이드"

네트워크가 혼란에 빠지거나 나쁜 습관을 들이지 않도록, 저자들은 두 가지 특별한 재료를 추가했습니다:

1. 신뢰 영역 (RL-T): 당신이 개에게 달리기 훈련을 시킨다고 상상해 보십시오. 만약 너무 빨리 달리라고 하면 개가 넘어질 수 있습니다. "신뢰 영역"은 개의 달리기 스타일이 한 번에 너무 많이 변하지 않도록 제한하는 목줄과 같습니다. 이는 학습을 안정적으로 유지하고 네트워크가 무모하고 잘못된 추측을 하는 것을 방지합니다. 논문은 이것이 학습을 더 부드럽고 신뢰할 수 있게 만든다는 것을 보여줍니다.
2. 가이드 정책 (RL-G): 때때로 개에게는 약간의 힌트가 필요합니다. 저자들은 네트워크를 막다른 길(낮은 보상 구역)에서 멀어지게 하고 보물을 향해 부드럽게 유도하는 지도 역할을 하는 "가이드" 정책을 도입했습니다. 이는 네트워크가 좋은 솔루션이 근처에 없는 "보상 사막"에 갇히는 것을 방지하는 데 도움을 줍니다.

이것이 왜 중요한가

이 논문은 GFlowNet의 훈련을 직접적인 보상 최적화 문제로 재정의함으로써, 복잡한 객체를 생성하기 위한 더 나은, 더 빠르고, 더 신뢰할 수 있는 AI를 구축할 수 있다고 결론짓습니다. 새로운 약물을 설계하든, 공급망을 최적화하든, 혹은 우주의 구조를 이해하든, 이 방법은 솔루션으로 가는 더 직접적인 경로를 제공합니다. 저자들은 엄격한 수학적 증명과 실제 및 시뮬레이션 데이터에 대한 광범한 실험을 통해 자신들의 결과에 확신을 가지고 있습니다. 그들은 단순히 이것이 작동할 것이라고 추측한 것이 아니라, 이것이 실제로 작동함을 보여줌으로써 AI가 창조하고 발견하는 법을 가르치는 방법에 대한 유망한 새로운 방향을 제시했습니다.