기술 요약: OOD(Out-of-Distribution) 발견을 위한 능동적 흐름 확장 (Active Flow Expansion)

문제 정의

플로우(Flow) 및 확산(Diffusion) 모델을 포함한 표준 생성 모델은 일반적으로 **데이터 분포 매칭(Data Distribution Matching)**을 통해 사전 학습됩니다. 이 목적 함수는 모델이 가용한 학습 데이터의 분포($p_{data}$)를 근사하도록 강제하며, 이는 종종 유효한 설계 공간($\Omega^*$)의 극히 일부만을 커버합니다. 이러한 방식은 행동 복제(Behavioral Cloning)나 관찰된 사례를 재현하는 작업에는 효과적이지만, OOD(Out-of-Distribution) 발견에는 근본적인 한계가 있습니다. 과학적 발견(예: 자연계에 존재하지 않는 새로운 분자 또는 단백질 설계)의 목표는 관찰된 데이터 분포 하에서는 확률이 무시할 수 있을 정도로 낮으면서도 기존 사례들과는 임의로 멀리 떨어진 유효한 설계를 샘플링하는 것입니다. 표준 사전 학습 모델은 이 영역에 접근할 수 없는데, 그 이유는 모델이 샘플링할 수 있는 영역인 "생성 가능 집합(Generable Set)"이 관찰된 데이터 분포 내에 엄격하게 포함되어 있기 때문입니다.

방법론: 능동적 흐름 확장 (ACTFLOW)

저자들은 모델의 생성 가능 집합을 더 넓은 유효 설계 공간으로 확장하기 위해 설계된 지속적 사전 학습 프레임워크인 **능동적 흐름 확장(ACTFLOW)**을 제안합니다. 이 방법은 수동적인 데이터 분포 매칭에서 벗어나, 생성 모델을 유효한 설계 공간의 근사치로 취급합니다.

핵심 개념

1. 생성 가능 집합 ($\Omega^\tau_\theta$): 모델의 최종 밀도 $p^\theta_1(x)$가 임계값 $\tau$를 초과하는 설계 공간의 영역으로 정의됩니다. 표준 사전 학습 결과는 $\Omega^\tau_{\theta_{pre}} \subset \Omega^*$이며, $\text{Vol}(\Omega^\tau_{\theta_{pre}}) \ll \text{Vol}(\Omega^*)$를 만족합니다.
2. 생성 가능 집합 확장: 새로운 학습 원리는 파라미터 $\theta$를 학습하여 $\Omega^\tau_\theta \approx \Omega^*$가 되도록 함으로써, 모델의 커버리지를 유효한 OOD 영역까지 효과적으로 확장하는 것을 목표로 합니다.

알고리즘 워크플로우

ACTFLOW는 능동적 탐색과 모델 적응의 반복 루프를 통해 작동합니다:

• 학습된 표현에서의 능동적 탐색: 현재 모델의 지배적인 모드(Mode)에서 샘플링하는 대신, ACTFLOW는 중간 노이즈 수준 $s \in (0,1)$에서의 속도 네트워크(Velocity Network)로부터 추출된 학습된 플로우 표현 $Z_s$에서 능동적 탐색을 수행합니다.
• 불확실성 편향 샘플링 (Uncertainty-Tilted Sampling): 이 방법은 샘플링을 **검증기 불확실성(Verifier Uncertainty)**이 높은 영역으로 유도합니다. 이는 기대 검증기 불확실성 $\sigma_t(\phi_s(x))$을 최대화하는 동시에, KL-다이버전스 항을 통해 현재 모델 밀도 근처에 머물도록 규제하는 최적화 문제(식 9)를 해결합니다. 이는 새로운 유효한 경계의 탐색과 모델의 사전 지식 사이의 균형을 맞춥나다.
• 검증기 피드백: 블랙박스 검증기 $\tilde{v}$ (예: 분자 유효성 체크 도구 또는 단백질 접힘 예측기)가 생성된 샘플을 평가하여 이진 레이블(유효/무효)을 제공합니다.
• 부호가 있는 리플레이 기반 업데이트 (Signed Replay-Based Update): 모델은 부호가 있는 손실 함수를 사용하여 업데이트됩니다. 수용된 샘플(유효)은 양의 그래디언트(모델을 끌어당김)를 기여하고, 거절된 샘플(무효)은 음의 그래디언트(모델을 밀어냄)를 기여합니다. 이는 $L_{signed} = L_{accepted} - \alpha_t L_{rejected}$라는 부호 있는 리플레이 목적 함수로 정식화되며, 이는 일종의 무효 영역에 대한 데이터 언러닝(Unlearning)으로 볼 수 있습니다.

이론적 프레임워크

본 논문은 OOD 플로우 모델링을 로컬-투-글로벌 도달 가능성(Local-to-Global Reachability) 문제로 프레이밍하여 통계적 학습 보증을 확립합니다:

• 검증기 모델링: 검증기는 학습된 표현 공간 상의 로지스틱 밴딧 피드백 메커즘으로 모델링되며, 여기서 유효성은 미지의 잠재 스코어 $g(z)$에 의해 결정됩니다.
• 도달 가능 연산자 (Reachability Operator): 저자들은 알려진 안전한 집합 $S$로부터 특정 오차 범위 내에서 유효성을 인증할 수 있는 점들의 집합을 나타내는 일단계 도달 가능 연산자 $R_\epsilon(S)$를 정의합니다.
• 주요 정리: 잘 지정된 검증기와 근사적 불확실성 샘플링 가정 하에, ACTFLOW는 $T^*$번의 검증된 샘플 이후 모델의 생성 가능 집합이 초기 유효 시드 집합의 $H$-단계 재귀적 확장($R^H_\epsilon(S_0) \subseteq \Omega^\tau_{T^*}$)을 커버할 것임을 보장합니다. 이는 OOD 플로 모델링에서 생성 가능 집합을 확장하는 것에 대한 첫 번째 이론적 보증을 제공합니다.

주요 기여

1. 생성 가능 집합 확장: 데이터 분포 매칭이 아닌 생성 가능 집합 확장에 기반한 OOD 플로우 모델링을 위한 엄격한 수학적 프레임워크 도입.
2. ACTFLOW 알고리즘: 검증기 피드백과 중간 노이즈 표현에서의 능동적 탐색을 사용하는 실용적이고 태스크 불가지론적인(Task-agnostic) 지속적 사전 학습 방법 제시.
3. 이론적 보증: ACTFLOW가 로컬-투-글로벌 과정을 통해 생성 가능 집합을 도달 가능한 유효 영역으로 확장함을 증명함으로써, OOD 플로우 모델링에 대한 최초의 통계적 학습 보증 제공.
4. 실증적 검증: 다양한 도메인(소분자, 약물 유사 분자, 치료용 펩타이드, 단백질 서열)에 대한 포괄적인 평가를 통해 재귀적 자기 생성 베이스라인 대비 우수한 성능 입증.

실험 결과

저자들은 ACTFLOW를 두 가지 널리 사용되는 베이스라인과 비교 평가했습니다:

• REC-NF: 필터링되지 않은 모델 샘플에 대한 지속적 사전 학습.
• REC-F: 검증기로 필터링된 샘플(유효한 샘플만 사용)에 대한 지속적 사전 학습.

성능 지표: 커버리지(유효 클러스터 수), 다양성(Vendi score), 유효성(유효 샘플 비율).

• 유기 소분자 (QM9): ACTFLOW는 유효 커버리지를 88.40 클러스터(REC-F의 45.40 및 REC-NF의 42.80 대비)로 확장하면서도 높은 유효성(95.90%)을 유지했습니다. REC-NF는 심각한 유효성 저하(12.26%)를 겪었습니다.
• 약물 유사 분자 (GEOM-Drugs): ACTFLOW는 144.3개의 유효 클러스터를 달성하여 REC-F(89.33)와 REC-NF(44.67)를 크게 앞질렀습니다. 또한 사전 학습된 베이스라인(56.6%)보다 유효성을 개선했습니다(39.31% 대비).
• 치료용 펩타이드: ACTFLOW는 커버리지를 61에서 358 클러스터로 증가시켰으며 가장 높은 다양성(58.87 Vendi score)을 달성했습니다. 반면, REC-NF는 유효성이 0%로 붕괴되었습니다.
• 단백질 서열: ACTFLOW는 커버리지를 66.50에서 102.75 클러스터로, 다양성을 42.14(Vendi)로 증가시키는 동시에 유효성을 83.74%로 개선했습니다. REC-F는 더 높은 유효성(88.12%)을 보였으나 커버리지와 다양성 측면에서 붕로(Collapse)되었습니다.

모든 태스크에서 ACTFLOW는 커버리지, 다양성, 유효성 간의 트레이드오프에서 베이스라인을 크게 능가하며, 재귀적 베이스라인이 도달하지 못한 새로운 유효 영역으로의 확장에 성공했습니다.

의의 및 주장

본 논문은 OOD 생성 모델링을 위한 최초의 이론적 및 실무적 접근 방식을 제공한다고 주장합니다. "분포 매칭"에서 "생성 가능 집합 확장"으로 패러다임을 전환함으로써, ACTFLOW는 사전 학습된 모델이 표준적인 샘플링이나 단순한 미세 조정으로는 접근할 수 없는 유효한 설계 공간에 접근할 수 있도록 합니다. 저자들은 이 방법이 안전한 탐색 이론(Safe Exploration Theory)과 생성 모델링 사이의 간극을 메우며, 새로운 자연계 설계를 발견하기 위한 증명 가능한 보증을 제공한다고 강조합니다.

실제적인 발견 측면에서 연구의 범위는 겸손하게 설정되어 있습니다: 본 프레임워크가 이전에는 접근 불가능했던 유효 영역을 향한 태스크 불가지론적 확장을 허용하지만, 이러한 탐색이 특정 실제 응용 분야에서 구체적인 과학적 발견을 낳을 수 있는지 여부를 평가하기 위해서는 향후 연구가 필요하다고 언급합니다. 주요 기여는 이러한 발견을 가능하게 하는 알고리즘적 및 이론적 토대를 마련하는 데 있습니다.