Active Flow Matching

이 논문은 VSD 및 CbAS 와 같은 변분 프레임워크와의 통합을 가능하게 하여 온라인 블랙박스 최적화에서 고적합도 영역을 효과적으로 탐색하고 활용하는 새로운 방법론인 'Active Flow Matching(AFM)'을 제안합니다.

핵심

1. 문제 상황: 거대한 미로와 비싼 나침반

생각해 보세요. 여러분은 거대한 미로 속에 있습니다. 이 미로는 단백질이나 분자의 구조를 나타내는데, 가능한 경우의 수가 우주의 별 개수보다도 많습니다 (예: 20 개의 아미노산으로 이루어진 펩타이드는 $20^{20}$가지나 됩니다).

• 목표: 이 미로에서 '보물' (가장 성능이 좋은 분자) 을 찾는 것입니다.
• 문제: 보물을 찾으려면 직접 그 자리에 가서 실험을 해봐야만 점수 (적합도) 를 알 수 있습니다. 하지만 이 실험은 엄청나게 비싸고 시간이 걸립니다.
• 기존 방법의 한계:
  • ** autoregressive (순차적) 모델:** 한 글자씩 순서대로 써가는 방식입니다. 마치 글을 쓸 때 앞 글자를 보고 다음 글자를 정하는 것처럼요. 하지만 단백질은 글자끼리 서로 먼 거리에서도 영향을 주고받는 (상호작용) 경우가 많습니다. 순서대로만 쓰면 이런 복잡한 관계를 놓치기 쉽습니다.
  • 기존의 '비밀 지도' (Implicit Generator): 최신 AI 모델들은 미로의 전체 지도를 그리는 대신, "지금 이 위치에서 다음 단계로 어떻게 움직여야 할지"만 알려줍니다. (예: "지금 이 방에서 동쪽으로 가라"는 식). 하지만 이 모델은 "이 위치가 보물일 확률이 정확히 얼마인가?"라고 숫자로 말해주지 않습니다.
  • 기존 최적화 기술 (VSD, CbAS): 이 기술들은 보물을 찾기 위해 "이곳에 보물이 있을 확률"을 계산해야 하는데, 위 AI 모델은 그 확률 숫자를 알려주지 않아서 함께 쓸 수 없었습니다.

2. 해결책: "활성 흐름 매칭 (AFM)"의 등장

저자들은 이 문제를 해결하기 위해 새로운 전략을 세웠습니다.

핵심 아이디어: "완전한 지도 (확률) 를 알 수 없다면, 현재 위치에서 다음 단계로 가는 '방향'과 '속도'에 집중하자."

비유: 안개 낀 산에서 정상으로 가는 등반

• 기존 방식: 정상 (보물) 에 도달할 확률을 정확히 계산해서 등반 경로를 정하려 했지만, 안개 (복잡한 상호작용) 때문에 확률 계산이 불가능했습니다.
• AFM 의 방식:
  1. 흐름 (Flow) 활용: AI 는 "지금 이 지점에서 정상 쪽으로 가려면 어떻게 움직여야 하는지"를 알려줍니다. (예: "지금 이 바위에서 오른쪽으로 3 걸음, 위로 2 걸음")
  2. 조건부 학습: 우리는 "보물이 있는 지역"으로만 AI 를 유도하고 싶습니다. AFM 은 AI 가 가르쳐 주는 '방향'을 받아서, **"보물이 있을 것 같은 곳으로 갈수록 그 방향을 더 강하게 조정"**합니다.
  3. 결과: AI 는 여전히 "전체 지도"를 말해주지 않아도 되지만, 우리가 원하는 '보물 지역'으로 자연스럽게 이동하도록 유도할 수 있게 됩니다.

3. 어떻게 작동할까? (세 가지 전략)

저자들은 이 흐름을 조정하는 세 가지 방법을 제안했습니다.

1. Forward-KL (가장 성공적인 방법):

   • 비유: "우리가 원하는 보물 지역을 잘 보여주는 지도를 AI 에게 가르치는 것"입니다.
   • AI 가 만들어낸 경로가 우리가 원하는 보물 지역과 얼마나 잘 겹치는지 확인하며, 겹치지 않으면 AI 를 수정합니다. 이 방법이 실험에서 가장 좋은 결과를 냈습니다.

2. Reverse-KL:

   • 비유: "AI 가 만들어낸 경로가 보물 지역을 얼마나 잘 찾았는지 AI 스스로에게 물어보고 수정하는 것"입니다.
   • 때로는 너무 좁은 지역만 찾게 되어 (Exploration 부족), 보물을 놓칠 수도 있습니다.

3. Symmetric-KL:

   • 위 두 가지를 섞은 방법입니다.

4. 실험 결과: 실제로 효과가 있을까?

저자들은 이 기술을 단백질 설계와 약물 개발 시뮬레이션에 적용해 보았습니다.

• 결과: 기존에 가장 잘하던 방법들 (VSD, CbAS 등) 보다 더 적은 실험 횟수로 더 좋은 분자를 찾아냈습니다.
• 특징: 특히 단백질처럼 복잡한 상호작용이 많은 경우, 순서대로만 만드는 기존 방식보다 훨씬 뛰어난 성능을 보였습니다.

5. 요약: 왜 이것이 중요한가?

이 논문은 **"AI 가 완벽한 지도를 주지 않아도, 우리는 그 AI 를 이용해 보물을 찾을 수 있다"**는 것을 증명했습니다.

• 기존: "이곳에 보물이 있을 확률이 50% 입니다." (이 숫자를 계산할 수 없어서 최적화가 어려움)
• AFM: "이곳에서 보물 쪽으로 가려면 이렇게 움직여야 합니다. 그리고 보물이 있을 것 같은 곳으로 갈수록 그 움직임을 더 강조하세요." (이 방향을 이용해 최적화 성공)

이 기술은 비싼 실험을 줄이면서 새로운 약물을 개발하거나, 더 효율적인 단백질을 만드는 데 큰 역할을 할 것으로 기대됩니다. 마치 안개 낀 미로에서 나침반 없이도, 현지인의 안내를 잘 따라가며 보물을 찾아내는 것과 같습니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 고차원 설계 공간의 복잡성: 단백질 설계나 분자 설계와 같은 고차원 목적 함수 (objective landscape) 는 비가산적 (non-additive) 이고 비자기회귀적 (non-autoregressive) 인 복잡한 상호작용 (예: 단백질 진화의 에피스타시스) 을 포함합니다. 기존 자기회귀 (AR) 모델은 이러한 장기 의존성을 포착하는 데 한계가 있습니다.
• 이산 흐름 매칭 (Discrete Flow Matching, DFM) 의 한계: DFM 은 병렬적, 반복적 정제를 통해 복잡한 구조를 학습하는 강력한 생성 모델입니다. 그러나 DFM 은 **암시적 생성기 (implicit generator)**로, 이산 시퀀스 위의 정규화된 주변 분포 (marginal distribution, $q_\phi(x)$) 를 직접 계산할 수 없습니다.
• 활성 생성 (Active Generation) 과의 충돌:
  • **VSD (Variational Search Distributions)**와 **CbAS (Conditioning by Adaptive Sampling)**와 같은 활성 생성 프레임워크는 블랙박스 최적화를 위해 생성 모델의 로그 확률 밀도 ($\log q_\phi(x)$) 나 그 기울기를 필요로 합니다.
  • DFM 은 주변 확률 밀도를 계산하는 폐쇄형 식 (closed-form) 을 제공하지 않으며, 모든 가능한 경로에 대한 합을 요구하므로 계산이 불가능합니다.
  • 핵심 문제: DFM 의 유연성과 VSD/CbAS 와 같은 원칙적인 최적화 프레임워크를 결합하는 것이 수학적으로 불가능하여, 실험 예산이 제한된 상황에서 고적합도 (high-fitness) 영역을 탐색하는 데 어려움이 있었습니다.

2. 제안 방법론: Active Flow Matching (AFM)

저자들은 DFM 이 제공하는 **조건부 엔드포인트 분포 (conditional endpoint distributions, $q_\phi(x_1 | x_t, t)$)**를 직접 활용하여 변분 목적 함수를 재구성함으로써 위 문제를 해결했습니다.

핵심 아이디어

• 가변적 목적 함수의 재정의: 계산 불가능한 주변 분포 $q_\phi(x)$ 대신, 흐름 경로 (flow path) 상의 조건부 분포 $q_\phi(x_1 | x_t)$를 기반으로 KL 발산을 최소화하는 방식으로 목적 함수를 변경했습니다.
• 자기 정규화 중요도 샘플링 (Self-Normalised Importance Sampling, SNIS): 샘플링된 데이터에 가중치를 부여하여 목표 분포에 대한 편향을 보정합니다.

주요 변형 (Variants)

1. Forward-KL AFM (CbAS 기반):
   • 목표: 실제 조건부 엔드포인트 분포와 모델 분포 간의 Forward-KL 을 최소화.
   • 목적 함수: $L_{fwd}(\phi) \approx -E[\sum w_k \log q_\phi(x_{1,k} | x_{t,k})]$
   • 특징: 이론적으로 일관성 (consistency) 이 보장됩니다. 즉, 최적화 시 원하는 주변 분포 $p^*(x)$를 생성하는 흐름을 학습합니다.
2. Reverse-KL AFM (VSD 기반):
   • 목표: 모델 분포와 목표 분포 간의 Reverse-KL 을 최소화.
   • 목적 함수: $L_{rev}(\phi)$는 모델이 생성한 샘플에 대한 기대값을 사용합니다.
   • 특징: 모드 찾기 (mode-seeking) 성향이 강하며, 이론적 일관성 증명은 아직 완료되지 않았습니다.
3. Symmetric-KL AFM:
   • Forward 와 Reverse KL 을 모두 결합하여 탐색 (exploration) 과 활용 (exploitation) 의 균형을 시도합니다.

제안 분포 (Proposal Distribution) 설계

효율적인 중요도 샘플링을 위해 세 가지 성분이 혼합된 제안 분포 $\mu(x)$를 사용합니다:

• Prior: 균일 분포 (전체 공간 탐색).
• Flow: 이전 라운드의 흐름 모델에서 생성된 분포 (국소적 정밀화).
• Replay Buffer: 과거에 관측된 고점수 시퀀스 저장소 (고적합도 영역 집중).
• 실제 구현: 각 반복 단계에서 하나의 성분을 선택하여 샘플링하고, 해당 성분에 대한 가중치만 계산하여 계산 비용을 줄입니다.

3. 주요 기여 (Key Contributions)

1. 이론적 통합: 암시적 생성 모델 (DFM) 과 변분적 활성 생성 프레임워크 (VSD, CbAS) 를 통합하는 새로운 프레임워크인 **Active Flow Matching (AFM)**을 제안했습니다.
2. 계산 가능 목적 함수 유도: 불가능한 주변 확률 밀도 대신 조건부 분포를 사용하여 KL 발산을 재정의하고, SNIS 를 통해 이를 최적화 가능한 형태로 유도했습니다.
3. Forward-KL 의 일관성 증명: Forward-KL AFM 이 이론적으로 목표 분포를 생성하는 흐름을 학습함을 증명했습니다.
4. 실용적 제안 분포: 탐색, 활용, 정밀화를 균형 있게 수행하는 혼합 제안 분포를 설계하여 샘플링 효율성을 높였습니다.

4. 실험 결과 (Results)

다양한 단백질 및 분자 설계 태스크 (Ehrlich 합성 지형, AAV 캡시드, FoldX 기반 단백질 안정성/SASA, 분자 도킹) 에서 평가되었습니다.

• 성능:
  • Forward-KL AFM은 대부분의 태스크에서 SOTA 기법 (VSD, CbAS, LaMBO-2) 보다 경쟁력 있거나 우수한 성능을 보였습니다.
  • 특히 Ehrlich 및 AAV 태스크에서 Forward-KL AFM 은 가장 빠른 수렴 속도와 가장 낮은 단순 후회 (simple regret) 를 기록했습니다.
  • FoldX 안정성 최적화에서도 다른 방법론보다 빠르게 고안정성 변이를 발견했습니다.
  • 분자 도킹 (F2) 태스크에서는 VSD 를 크게 능가하는 도킹 점수를 달성했습니다.
• 비교 분석:
  • Reverse-KL AFM은 일부 태스크 (AAV 등) 에서 수렴이 느리거나 조기 수렴 (premature convergence) 하는 경향을 보였습니다.
  • CbAS는 초기에 발견된 시퀀스 밖으로 탐색하지 못하고 조기 수렴하는 경향이 있었습니다.
  • VSD는 긴 시퀀스 (Ehrlich-64) 에서 장기 에피스타시스 상호작용을 포착하는 데 어려움을 겪었습니다.
• 결론: 제한된 실험 예산 하에서 Forward-KL AFM 이 탐색 - 활용 (exploration-exploitation) 트레이드오프를 가장 효과적으로 수행했습니다.

5. 의의 및 결론 (Significance)

• 블랙박스 최적화의 새로운 패러다임: DFM 과 같은 최신 비자기회귀 생성 모델의 강력한 표현력을, VSD/CbAS 와 같은 확률론적 최적화 프레임워크와 결합하여, 이산적이고 고차원인 공간에서의 활성 생성 문제를 해결했습니다.
• 실용적 가치: 실험 비용이 매우 비싼 단백질 및 분자 설계 분야에서, 적은 수의 실험 (oracle calls) 으로 고품질의 설계를 발견할 수 있는 가능성을 제시했습니다.
• 미래 방향: AFM 은 암시적 생성 모델과 엄격한 의사결정 이론적 목적 함수를 통합하는 첫 걸음으로, 다목적 최적화, 제약 조건 최적화, 그리고 생물학적 시퀀스 설계 외의 다양한 분야로 확장될 수 있는 기반을 마련했습니다.

요약하자면, 이 논문은 **DFM 의 계산적 한계 (주변 확률 부재) 를 우회하여 조건부 분포를 기반으로 활성 생성을 수행하는 새로운 방법론 (AFM)**을 제시하며, 이를 통해 단백질 및 분자 설계 분야에서 기존 방법론보다 우수한 성능을 입증했습니다.