기술 요약: 슈퍼그룹 상대 정책 최적화를 통한 생체분자 유용성-다양성 frontier 확장

1. 문제 제기

생체분자 생성은 특정 화학적 또는 생물학적 목표를 만족하는 후보 물질 (소분자 또는 단백질) 을 생성하는 것을 목표로 합니다. 강화 학습 (RL) 은 사전 훈련된 생성기를 후속 훈련하여 유용성 (예: 약물 유사성, 도킹 점수, 안정성) 을 극대화하는 표준 프레임워크이지만, 유용성만 최적화하면 종종 **모드 붕괴 (mode collapse)**가 발생하여 생성기가 고득점 후보의 좁은 계열에 확률 질량을 집중시킵니다. 반면, 다양성 유지는 어렵습니다. 왜냐하면 다양성은 샘플 집합에 정의된 **집합 수준 속성 (set-level property)**인 반면, 표준 정책 최적화는 개별 롤아웃을 업데이트하기 때문입니다.

기존의 다양성 증진 RL 방법은 역사적 샘플에 대한 신규성 페널티나 메모리 기반 필터와 같은 간접적 대리 변수에 의존하는 경우가 많습니다. 이러한 접근법은 유용한 고밀도 모드를 과도하게 페널티하거나 분포적 드리프트를 유발할 수 있습니다. 본 논문이 다루는 핵심 과제는 다음과 같습니다: 개별 생성된 후보에 유용한 크레딧을 부여하면서도, 어떻게 집합 수준의 다양성을 1 순위 목표로 직접 최적화할 수 있을까요?

2. 방법론: 슈퍼그룹 상대 정책 최적화 (SGRPO)

저자들은 롤아웃 수준 유용성과 집합 수준 다양성 간의 균형을 직접 최적화하는 유연한 프레임워크인 **슈퍼그룹 상대 정책 최적화 (SGRPO)**를 제안합니다. 이는 그룹 상대 정책 최적화 (GRPO) 에서 영감을 받았습니다.

핵심 메커니즘

SGRPO 는 동일한 조건 하에서 여러 후보 집합 (그룹) 을 샘플링하고 이를 비교하여 다양성 보상을 재분배함으로써 작동합니다. 이 과정은 네 가지 주요 단계를 포함합니다:

1. 동일 조건 슈퍼그룹 샘플링:
   주어진 조건 $C$(예: 특정 단백질 포켓 또는 조건 없음) 에 대해 정책은 각각 $K$개의 독립적으로 생성된 롤아웃을 포함하는 $M$개의 그룹을 샘플링합니다. 이 컬렉션 $S(C) = \{G_1, \dots, G_M\}$은 $N=MK$개의 후보로 구성된 "슈퍼그룹"을 형성합니다. 비교는 이 슈퍼그룹 내에서만 제한되어, 정책의 품질을 서로 다른 조건의 고유한 난이도와 혼동하지 않도록 합니다.

2. 그룹 수준 다양성 점수화:
   각 그룹 $G_m$은 사용자가 지정한 지표 (예: 분자의 경우 내부 Tanimoto 거리, 단백질의 경우 Levenshtein 유사도) 를 기반으로 다양성 점수 $R_m = D(G_m)$를 부여받습니다. 그룹 상대 다양성 신호는 슈퍼그룹 내 다른 그룹에 대한 leave-one-out 비교를 통해 계산됩니다:
   $$A_m^{grp} = R_m - \frac{1}{M-1}\sum_{h \neq m} R_h$$
   양의 $A_m^{grp}$는 해당 그룹이 다른 대안들보다 더 다양함을 나타냅니다.

3. 집합 인식 재분배 (Leave-One-Out 기여도):
   집합 수준 신호와 개별 정책 업데이트 간의 간극을 메우기 위해, SGRPO 는 그룹의 다양성에 대한 개별 롤아웃의 기여도에 따라 그룹 다양성 보상을 개별 롤아웃에 재분배합니다.

   • 각 롤아웃 $x_{m,i}$에 대해 leave-one-out 기여도가 계산됩니다: $c_{m,i} = D(G_m) - D(G_m \setminus \{x_{m,i}\})$.
   • 이러한 기여도는 표준화되어 부호 인식 (sign-aware) 소프트맥스 가중치를 형성하는 데 사용됩니다.
   • 그룹이 매우 다양할 경우 ($A_m^{grp} > 0$), 양의 신호는 높은 다양성 기여도를 가진 후보에 집중됩니다. 그룹이 덜 다양할 경우 ($A_m^{grp} < 0$), 음의 신호는 낮은 기여도를 가진 후보에 집중됩니다.
   • 이를 통해 평균 그룹 보상은 유지되면서 다양성을 주도하는 특정 분자에 크레딧이 부여됩니다.

4. 슈퍼그룹 상대 정책 업데이트:
   롤아웃의 최종 보상은 개별 유용성 $r_{m,i}$와 재분배된 다양성 보상 $\tilde{R}_{m,i}$를 결합합니다:
   $$\bar{r}_{m,i} = (1-\lambda)r_{m,i} + \lambda \tilde{R}_{m,i}$$
   여기서 $\lambda$는 유용성 - 다양성 균형을 조절합니다. 이점 함수 (advantage function) 는 슈퍼그룹 내 모든 롤아웃에 대한 leave-one-out 기준선으로 계산되며, KL 정규화가 적용된 클리핑된 PPO 스타일 목적 함수를 사용하여 정책이 업데이트됩니다.

3. 주요 기여

• 직접적 집합 수준 최적화: SGRPO 는 간접적 신규성 대리 변수나 역사 의존적 페널티에 의존하는 대신, 생체분자 생성에서 집합 수준 다양성을 주요 목표로 직접 최적화하는 최초의 프레임워크입니다.
• 분리된 설계: 이 방법은 특정 생성기 아키텍처 (자기회귀 또는 확산), 유용성 보상, 또는 다양성 지표와 분리되어 있어 다양한 GRPO 스타일 최적화기 (예: 자기회귀 모델용 표준 GRPO, 이산 확산용 Coupled-GRPO) 로 인스턴스화할 수 있습니다.
• 크레딧 할당 메커니즘: leave-one-out 다양성 기여도 메커니즘의 도입은 집합 다양성을 실제로 지원하는 개별 후보에 크레딧을 할당할 수 있게 하여, 개별 업데이트로 집합 속성을 최적화하는 문제를 해결합니다.
• 그룹 크기에 대한 강건성: 이 방법은 작은 그룹 크기에서도 효과적으로 작동하여 대규모 생체분자 설계에 계산적으로 실행 가능합니다.

4. 실험 결과

저자들은 세 가지 다른 생체분자 생성 설정에서 SGRPO 를 평가했습니다:

1. De Novo 소분자 설계: GenMol(이산 확산) 사용.
2. 포켓 기반 소분자 설계: GenMol-P(포켓 조건부 확산) 사용.
3. De Novo 단백질 설계: ProGen2(자기회귀) 사용.

성능 지표:
평가는 **유용성 - 다양성 파레토 프론티어 (Utility-Diversity Pareto Frontier)**에 초점을 맞추었으며, 다음 지표로 측정되었습니다:

• 초체적 (Hypervolume, HV): 비우세 운영 점들에 의해 지배되는 면적.
• 이상점까지의 거리 (Distance to Ideal Point, DIP): 이론적 최상 유용성 및 다양성까지의 거리.
• R2 지표: 평균 가중 최악의 경우 부족분.

결과:

• 프론티어 확장: 모든 작업 및 디코딩 스윕에서 SGRPO 는 사전 훈련된 생성기, 표준 GRPO, 메모리 지원 GRPO 에 비해 달성 가능한 유용성 - 다양성 파레토 프론티어를 일관되게 확장했습니다.
• 다양성 보존: GRPO 등 기준선 모델이 심각한 다양성 붕괴를 겪었던 고유용성 영역에서 SGRPO 는 현저히 높은 다양성을 유지했습니다. 예를 들어, 포켓 기반 설계에서 SGRPO 는 운영점 궤적을 바깥쪽으로 이동시켜 동등한 유용성 수준에서 훨씬 더 높은 다양성을 유지했습니다.
• Ablation 연구: 다양성 보상을 제거하면 성능이 표준 GRPO 수준으로 돌아갔습니다. leave-one-out 크레딧 할당을 제거하면 성능이 약화되어, 다양성 신호를 특정 기여자에게 재분배하는 것의 중요성이 확인되었습니다.
• 분포 역학: 시퀀스 분포 시각화는 GRPO 와 메모리 지원 GRPO 가 좁은 영역으로 수축하거나 드리프트하는 반면, SGRPO 는 여러 클러스터를 탐색하고 보존하여 시퀀스 공간의 더 넓은 범위를 커버함을 보여주었습니다.

5. 중요성 및 주장

본 논문은 SGRPO 가 생체분자 생성 모델에 사용 가능한 운영점을 확장하는 광범위하게 적용 가능한 후속 훈련 원칙을 제공한다고 주장합니다. 생성된 다양한 집합을 직접 보상함으로써, 이 방법은 유용성 - 다양성 프론티어를 바깥쪽으로 밀어내어 모델이 커버리지를 희생하지 않고도 고정된 다양성 수준에서 더 나은 유용성을 달성할 수 있게 합니다 (그 반대의 경우도 마찬가지).

저자들은 이 접근법이 현재 RL 기반 생체분자 설계의 근본적인 한계, 즉 정책 최적화 과정에서 다양성을 1 순위 목표로 취급하지 못하는 문제를 해결한다고 강조합니다. 결과는 집합 수준 다양성을 직접 최적화하는 것이 소분자 및 단백질 생성 작업 모두에서 모드 붕괴를 완화하는 실용적이고 강건한 전략임을 시사합니다. 추가 연구와 다양한 도메인에서의 인스턴스화를 촉진하기 위해 코드가 공개되었습니다.