Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

이 논문은 GFlowNet 의 모드 붕괴 문제를 해결하기 위해 초기 접두사에 대한 밀도 있는 학습 신호를 제공하는 'Rooted absorbed prefix Trajectory Balance (RapTB)'와 다양성을 증진하는 'Submodular Replay (SubM)' 전략을 제안하여 분자 생성 등 다양한 작업에서 최적화 성능과 다양성을 동시에 향상시키는 방법을 제시합니다.

핵심

이 논문은 GFlowNet이라는 AI 기술을 더 똑똑하고 다양하게 만들 수 있는 새로운 방법을 제안합니다. 이를 이해하기 위해 **'창의적인 요리사'**와 **'레시피 훈련'**이라는 비유를 들어 설명해 드리겠습니다.

1. 문제 상황: 똑같은 요리만 반복하는 요리사

우리가 AI(요리사)에게 "맛있는 요리 (보상) 를 만들어줘"라고 시켰다고 상상해 보세요. 기존 방식 (GFlowNet) 은 훌륭한 요리사였지만, 두 가지 큰 실수를 저질렀습니다.

1. 초반 레시피의 실수 (Prefix Collapse): 요리사가 요리를 시작할 때, "소금 한 꼬집"이나 "양파 다지기" 같은 처음 단계에서 실수를 하면, 나중에 맛있는 요리를 만들어도 그 실수를 고쳐주지 못했습니다. 결과적으로 모든 요리가 처음 단계는 똑같고 (예: 모두 양파를 먼저 다짐), 나중에만 조금씩 달라지는 지루한 요리들만 나오게 되었습니다.
2. 너무 짧거나 긴 요리 (Length Bias): 요리사가 "맛있는 요리"를 만들려고 애쓰다가, 요리를 너무 빨리 끝내거나 (너무 짧은 국물), 반대로 끝없이 재료를 추가하다가 (너무 긴 스프) 적당한 길이를 잃어버렸습니다.

이 문제의 원인은 두 가지였습니다.

• 과거의 실수 (Credit Assignment): 요리가 완성된 후 "맛있다"라고 칭찬만 해줘서, "어디서부터 잘못했지?"를 정확히 알 수 없었습니다.
• 편향된 연습 (Replay Bias): 요리사가 잘 만든 요리 (보상 높은 것) 만 반복해서 보고 훈련해서, 새로운 시도를 하지 않고 같은 요리만 계속 만들어내는 '지루한 루프'에 빠졌습니다.

---

2. 해결책: 두 가지 새로운 훈련법 (RapTB & SubM)

저자들은 이 문제를 해결하기 위해 두 가지 혁신적인 방법을 제안했습니다.

방법 1: RapTB (뿌리에서 시작해 끝까지 흡수하는 훈련)

"요리 과정의 모든 단계에 점수를 매겨라!"

기존에는 요리가 완성된 후 한 번에 점수를 주었지만, 이 방법은 요리하는 과정 중 매 단계마다 점수를 줍니다.

• 비유: 요리사가 "양파를 다지는 순간"에도 점수를 받습니다. 만약 나중에 요리를 잘 만들었다면, 그 양파 다지기 단계에도 "잘했다"는 점수를 돌려줍니다.
• 흡수 (Absorbed) 의 의미: 나중에 요리를 잘 만들었다면, 그 좋은 결과가 과거의 모든 단계로 거슬러 올라가서 (흡수되어) 각 단계의 점수를 높여줍니다.
• 효과: 요리사가 처음 단계 (뿌리) 에서부터 실수를 바로잡고, 요리의 길이를 자연스럽게 조절할 수 있게 됩니다.

방법 2: SubM (다양한 레시피를 모으는 스마트한 파일함)

"유명한 요리만 모으지 말고, 다양한 요리를 섞어라!"

기존의 연습 방식은 "가장 맛있는 요리"만 파일함에 넣어두고 반복해서 보게 했습니다. 하지만 이 방법은 파일함 (Replay Buffer) 을 관리할 때 새로운 규칙을 적용합니다.

• 서브모듈러 (Submodular) 선택: 파일함을 채울 때, 단순히 점수가 높은 요리만 고르는 게 아니라, ① 맛있는 요리, ② 다른 재료로 만든 요리 (다양성), ③ 길이가 다양한 요리를 모두 골고루 섞어서 모으는 알고리즘을 사용합니다.
• 비유: 마치 요리 대회 심사위원이 "가장 맛있는 요리 10 개만 뽑는 게 아니라, 국물 요리, 볶음 요리, 찌개 요리 등 다양한 종류를 골고루 뽑아 심사하는 것"과 같습니다.
• 효과: AI 가 특정 요리 (모드) 에만 매몰되지 않고, **새롭고 다양한 요리 (해결책)**를 계속 찾아낼 수 있게 됩니다.

---

3. 실제 성과: 분자 설계와 수학 문제

이 방법을 실제로 테스트해 보니 놀라운 결과가 나왔습니다.

• 약물 개발 (분자 생성): AI 가 새로운 약물을 설계할 때, 유효한 (안전한) 분자를 만들면서도 다양한 구조의 분자를 계속 찾아냈습니다. 기존 방식은 똑같은 분자만 반복해서 만들었으나, 이 방법은 다양성과 품질을 동시에 잡았습니다.
• 수학 문제 (24 만들기): 숫자와 연산자로 24 를 만드는 문제를 풀 때, AI 가 정답을 찾는 비율은 유지하면서, 더 다양한 풀이 방법을 찾아내었습니다. 특히, 너무 짧거나 긴 문장을 만드는 실수가 크게 줄었습니다.

요약

이 논문은 **"AI 가 똑같은 답만 반복하지 않고, 창의적이고 다양한 해답을 찾게 하는 방법"**을 제시합니다.

1. RapTB: "결과만 칭찬하지 말고, 과정의 모든 단계에 점수를 주고, 좋은 결과가 과거로 돌아오게 하여 실수를 고쳐라."
2. SubM: "가장 좋은 답만 모으지 말고, 다양한 답을 섞어서 AI 가 새로운 시도를 하도록 훈련시켜라."

이 두 가지 방법을 합치면, AI 는 더 안정적이고, 더 창의적이며, 더 다양한 결과를 만들어낼 수 있게 됩니다. 마치 한 가지 레시피만 반복하던 요리사가, 다양한 재료와 기술을 익혀 훌륭한 셰프가 되는 것과 같습니다.

기술 요약

1. 문제 정의 (Problem)

생성 흐름 네트워크 (GFlowNet) 는 보상에 비례하는 확률 분포를 학습하여 다양한 고보상 해를 탐색하는 데 유용하지만, 대규모 언어 모델 (LLM) 에 적용할 때 모드 붕괴 (Mode Collapse) 현상이 빈번하게 발생합니다. 저자들은 이를 두 가지 구체적인 실패 모드로 규명했습니다.

1. 접두사 붕괴 (Prefix Collapse): 초기 토큰의 엔트로피가 급격히 떨어지고, 서로 다른 최종 결과물들이 거의 동일한 초기 접두사 (prefix) 를 공유하게 되는 현상입니다. 이는 초기 단계에 대한 신용 할당 (credit assignment) 이 약하기 때문입니다.
2. 길이 편향 (Length Bias): 모델이 시스템적으로 너무 짧거나 너무 긴 시퀀스를 선호하게 되어, 실제 과제의 요구사항과 맞지 않는 길이의 결과를 생성합니다.

이러한 문제의 원인은 다음과 같이 분석되었습니다.

• 약한 신용 할당: 최종 보상 (Terminal Reward) 만을 사용하는 경우, 중간 단계에 대한 피드백이 고분산 (high-variance) 이고 모호하여 초기 결정에 대한 학습 신호가 약해집니다.
• 재플레이 편향 (Replay Bias): 학습 데이터가 탐색 공간의 극히 일부에 국한되고, 이 좁은 부분집합이 반복적으로 강화되면서 분포가 붕괴됩니다. 기존 방법인 서브트래젝토리 밸런스 (SubTB) 는 LLM 의 종료 가능한 접두사 트리 구조에서 과도한 제약 조건을 만들어 종료 확률 (termination probability) 이 왜곡되는 '종료 드리프트 (termination drift)'를 유발합니다.

2. 제안된 방법론 (Methodology)

저자들은 두 가지 상호 보완적인 메커니즘을 제안하여 위 문제들을 해결했습니다.

A. RapTB (Rooted Absorbed Prefix Trajectory Balance)

전체 경로 균형 (Trajectory Balance, TB) 의 안정성을 유지하면서 접두사 수준의 학습 신호를 강화하는 새로운 목적 함수입니다.

• 루팅된 접두사 잔차 (Rooted Prefix Residuals): SubTB 와 달리 모든 임의의 윈도우를 제약하는 것이 아니라, 루트 (시작점) 에서 시작하는 접두사에 대해서만 밀도 높은 감독 신호를 제공합니다. 이는 서로 다른 시작점을 가진 윈도우들이 공유하는 종료 헤드를 통해 발생하는 충돌을 방지합니다.
• 흡수된 접미사 보상 (Absorbed Suffix Rewards): 중간 접두사에 대한 보상을 추정할 때, 해당 접두사 이후에 관찰된 접미사 (suffix) 의 보상을 '흡수'하여 사용합니다.
  • umax: 접미사 중 최대 보상을 하한선으로 사용.
  • usoft: 로그 공간에서 여러 보상을 부드럽게 집계 (Softmax 기반).
  • 이 두 가지를 혼합하여 utgt (목표 보상) 를 생성하고, 이를 통해 분산을 줄이고 초기 단계의 신용 할당을 개선합니다.
• 종료 그라디언트 분리: 보조 학습 branch 에서 종료 확률 (qθ(⊤|s)) 에 대한 그라디언트를 차단하여, 모델이 보상 신호를 맞추기 위해 종료 확률만 임의로 조정하는 것을 방지합니다.

B. SubM (Submodular Replay)

재플레이 버퍼의 분포 편향을 해결하기 위한 경험 재사용 전략입니다.

• 서브모듈러 최적화: 단순히 높은 보상을 가진 샘플만 저장하는 대신, 보상 (Quality), 다양성 (Diversity), 길이 커버리지 (Length Coverage) 를 동시에 고려하는 서브모듈러 함수를 최대화하는 샘플 집합을 선택합니다.
• 다양성 유지: facility-location 커버리지 연산자와 길이 히스토그램 커버리지를 사용하여, 버퍼가 특정 모드나 특정 길이에 편향되지 않고 다양한 해를 포함하도록 유도합니다.

3. 주요 기여 (Key Contributions)

1. LLM-GFlowNet 의 모드 붕괴 현상 규명: 접두사 붕괴와 길이 편향이 고분산 신용 할당과 재플레이에 의한 분포 이동으로 인해 발생함을 실험적으로 증명했습니다.
2. RapTB 제안: TB 를 기반으로 하되, 루팅된 접두사 제약과 흡수된 보상 기법을 도입하여 분산을 줄이고 안정적인 학습을 가능하게 했습니다. 이는 SubTB 의 종료 드리프트 문제를 해결합니다.
3. SubM 제안: 보상, 다양성, 길이를 균형 있게 고려하는 서브모듈러 재플레이 전략을 도입하여 학습 분포의 커버리지를 확장하고 붕괴를 방지했습니다.
4. 광범위한 실험 검증: 분자 생성 (SMILES), 산술 표현식 생성 (Expr24), 문장 생성 (CommonGen) 등 다양한 태스크에서 기존 방법 (TB, SubTB) 대비 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

• 분자 생성 (SMILES):
  • RapTB + SubM 조합은 높은 유효성 (Validity, ~98.8%) 을 유지하면서 보점 (Score) 과 다양성 (Entropy, FPDiv) 을 모두 극대화했습니다.
  • 기존 SubTB는 유효성이 급격히 떨어지고 (32.8%), TB는 다양성이 낮고 짧은 시퀀스에 집중하는 경향을 보였습니다.
  • 접두사 붕괴 분석: RapTB 는 TB 에 비해 깊은 단계까지 접두사 생존율 (Survival) 을 높이고, Top-1 집중도를 낮춰 더 넓은 분기 (branching) 를 유지함을 보였습니다.
• 산술 표현식 (Expr24):
  • 희소 보상 환경에서 RapTB는 TB 의 모드 붕괴를 해결하고, SubTB 의 정확도 저하를 방지하며 높은 정확도 (>99%) 와 커버리지를 동시에 달성했습니다.
  • SubTB는 종료 확률이 극단적으로 낮아지는 '종료 드리프트' 현상을 보였으나, RapTB 는 이를 교정했습니다.
• 문장 생성 (CommonGen):
  • SubTB 는 자연스러운 문장 종료를 억제하여 길이가 과도하게 길어지는 (20 토큰) 반면, RapTB 는 자연스러운 길이 (1112 토큰) 를 유지하며 BLEU 점수를 크게 향상시켰습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 기반 GFlowNet 의 핵심적인 학습 불안정성 (접두사 붕괴, 길이 편향, 종료 드리프트) 을 체계적으로 분석하고 해결책을 제시했습니다.

• 이론적 기여: TB 의 전역적 안정성과 서브트래젝토리 학습의 국소적 신호를 조화시키면서도, LLM 의 종료 가능한 트리 구조에 특화된 제약 조건 (루팅된 접두사, 흡수된 보상) 을 설계했습니다.
• 실용적 기여: SubM 을 통해 재플레이 버퍼의 다양성을 보장함으로써, GFlowNet 이 고보상 영역뿐만 아니라 탐색 공간 전체를 효과적으로 커버할 수 있게 했습니다.
• 미래 전망: 이 연구는 생성형 AI 가 복잡한 보상 함수 하에서도 다양하고 정확한 해를 탐색할 수 있는 견고한 프레임워크를 제공하며, 신약 개발, 재료 과학 등 다양한 과학적 발견 분야에서 GFlowNet 의 적용 가능성을 높였습니다.

요약하자면, RapTB와 SubM은 GFlowNet 이 LLM 환경에서 모드 붕괴 없이 안정적으로 학습하고, 보상에 비례하는 다양하고 고품질의 생성물을 산출할 수 있도록 하는 필수적인 기술적 진보입니다.