In-Context Learning for Latent Space Bayesian Optimization

기술 요약: 잠재 공간 베이지안 최적화를 위한 인컨텍스트 학습 (In-Context Learning for Latent Space Bayesian Optimization)

문제 정의
베이지안 최적화(Bayesian Optimization, BO)는 분자 설계 및 단백질 공학과 같은 구조화된 공간에서 비용이 많이 드는 블랙박스 목적 함수를 최적화하기 위한 표준적인 접근 방식이다. 잠재 공간 베이지안 최적화(Latent Space Bayesian Optimization, LSBO)는 인코더-디코더(일반적으로 VAE)를 통해 연속적인 잠재 표현을 학습함으로써 이러한 조합론적 공간의 문제를 해결한다.

표 형식의 파운데이션 모델(Tabular Foundation Models, TFMs)과 사전 데이터 적합 네트워크(Prior-Data Fitted Networks, PFNs)는 인컨텍스트 학습을 통해 근사적인 베이지안 예측을 학습함으로써 회귀 작업에서 최첨단 성능을 달아가 왔으나, 이들을 LSBO에 직접 적용하는 데에는 결정적인 불일치 문제가 발생한다. 현재의 TFM들은 특정 잠재 공간의 기하학적 구조나 분자 목적 함수(예: logP, QED)의 특성을 반영하지 않는 대규모 합성 회귀 작업들에 대해 사전 학습되어 있다. 결과적으로, 이러한 모델들이 사전 학습 중에 학습한 "암묵적 사전 확률(implicit prior)"은 디코더와 도메인 특화 속성에 의해 형성되는 잠재 코드와 목적 함수 사이의 매핑을 따르는 LSBO의 배포 환경과 일치하지 않는다. 따라서 TFM을 일반적인 플러그인 회귀기로 취급하는 것은 이 특수한 환경에서 모델의 잠재력을 온전히 활용하지 못하게 만든다.

방법론
저자들은 사전 학습된 표 형식 파운데이션 모델(특히 TabPFN-3)을 지속적 사전 학습(continued pretraining)을 통해 LSBO 도메인에 적응시키는 프레임워크인 LILBO(Latent In-context Learning for Bayesian Optimization)를 제안한다. 이 방법론은 두 가지 주요 구성 요소로 이루어진다:

1. BO 인식 합성 작업 생성 (BO-Aware Synthetic Task Generation):
   일반적인 회귀 사전 확률을 사용하는 대신, 저자들은 분자 목적 함수에 기반을 둔 합성 사전 학습 분포를 구축한다.

• 기본 작업(Base Tasks): 합성 목적 함수는 GuacaMol 데이터베이스에서 가져온 분자 기본 작업 라이브러리(예: logP, QED, 유사도, 재발견)로부터 유도된다.
• 목적 함수 구축: 각 학습 에피소드마다 기본 작업의 하위 집합이 샘플링되며, 선형 결합(디리클레 분포에 의해 가중치가 부여됨) 또는 비선형 결합(무작위 MLP 사용)을 통해 합성 스칼라 목적 함수가 생성된다.
• 컨텍스트 샘플링(Context Sampling): 고가치 영역에 집중하는 BO의 동작을 모방하기 위해, 컨텍스트 데이터는 균등하게 샘플링되지 않는다. 대신, VAE 잠재 공간 내의 유망한 영역을 강조하면서도 다양성을 유지하도록 분자 풀(pool) 상에 가치 편향 분포(Boltzmann distribution)가 유도된다.

2. 정규화를 동반한 지속적 사전 학습 (Continued Pretraining with Regularization):
   저자들은 모델을 처음부터 학습시키는 대신 사전 학습된 TabPFN-3 체크포인트로부터 초기화한다. 이는 대규모 사전 학습을 통해 학습된 광범위한 회귀 사전 확률을 보존하기 위함이다.

• 목표: 모델은 합성된 분자 LSBO 에피소드에 대해 지속적 사전 학습을 거친다.
• 정규화: 파괴적 망각(catastrophic forgetting)과 과도한 전문화를 방지하기 위해 L2-SP(L2-Soft Parameter) 정규화가 적용된다. 이는 적응된 모델 가중치($\theta$)를 원래의 체크포인트 가중치($\theta_0$)에 고정하여, 모델이 분자 잠재 공간의 특정 기하학에 적응하면서도 일반적인 목적의 회귀 능력을 유지하도록 보장한다.
• 학습 루프: 모델은 약 500,000개의 분자 LSBO 에피소드에 대해 8,000 스텝 동안 학습된다.

주요 기여

• LSBO 특화 적응: 본 논문은 일반적인 TFM 사전 학습과 LSBO 배포 시 발생하는 분포 불일치를 해결하기 위해, 인컨텍스트 학습 서로게이트를 특정 잠재 공간 최적화에 적응시키는 방법을 소개한다.
• 합성 분포 설계: 분자 기본 작업과 가치 편향 컨텍스트 샘플링을 포함하는 합성 작업 생성 파이프라인을 설계하여, BO의 요구 사항(고가치 영역에 집중)에 더 잘 부합하도록 하였다.
• 정규화된 지속 학습: 정규화된 지속적 사전 학습 목적 함수를 사용하여, 파운데이션 모델의 광범위한 사전 확률을 버리지 않으면서도 분자 LSBO에 특화될 수 있도록 한다.
• 진단(Diagnostics): 합성 사전 학습 분포가 관련 잠재 매니폴드 및 목적 함수 공간의 영역을 커버하는지 확인하기 위해 UMAP을 통한 잠재 공간 및 목적 함수 공간 진단을 도입하였다.

결과
제안된 모델인 LILBO는 사전 학습 중에 보지 못한 GuacaMol 스위트의 홀드아웃(held-out) 분자 최적화 벤치마크(예: Osimertinib, Fexofenadine, Ranolazine)에서 평가되었다.

• 성능: LILBO는 작업 전반에 걸친 평균 순위 측면에서 오프더쉘프(off-the-shelf) TabPFN-3 플러그인 베이스라인과 여러 특화된 LSBO 베이스라인(GP-LSBO, LOLBO, CoBO, InvBO 포함)을 능가하는 강력한 성능을 달성했다.
• 효율성: 표준 TabPFN-3 베이스라인 대비 개선 정도는 완만하지만 주목할 만하며, 이는 원래 파운데이션 모델의 거대한 규모(조 단위의 토큰)와 비교했을 때 상대적으로 가벼운 적응(50만 에피소드)을 통해 달성되었다.
• 실행 가능성: 본 결과는 인컨텍스트 확률적 서로게이트가 사전 학습 분포가 도메인에 적절히 맞춰질 경우, 현대적인 LSBO 파이프라인에서 가우시안 프로세스(GP)를 대체할 수 있는 실행 가능한 대안임을 입증한다.

의의 및 주장
본 논문은 표 형식의 파운데이션 모델이 LSBO에서 효과적으로 작동하려면, 적응 분포가 일반적인 사전 학습 작업과 배포 시 마주치는 잠재 공간 목적 함수 및 BO 이력 간의 불일치를 줄여야 한다고 주장한다. 저자들은 과학적 설계를 위한 인컨텍스트 서로게이트의 잠재력을 끌어올리기 위해 LSBO 특화 적응이 필수적이라고 강조한다.

본 연구는 현재의 접근 방식이 사전 학습 중에는 고정된 VAE를 가정하고 있으나, 실제 배포 시에는 주기적인 VAE 재학습(잠재 드리프트)이 발생한다는 한계점을 인정한다. 향-후 방향으로는 학습 컨텍스트를 위한 짧은 최적화 궤적 시뮬레이션, 사전 확률에서의 잠재 드리프트 고려, 다중 목적 최적화로의 확장, 그리고 단백질 설계와 같은 다른 구조화된 도메인으로의 적용 등이 제시되었다.