기술 요약: 생성 모델의 출력에 대한 샘플 효율적 최적화 (O3)

문제 정의

현대적인 생성형 AI 모델, 특히 확산 (diffusion) 및 흐름 매칭 (flow-matching) 모델은 풍부한 데이터 분포로부터 샘플링하는 데 탁월한 성능을 보입니다. 그러나 많은 과학 및 공학 응용 분야는 단순한 무작위 샘플링 이상의 것을 요구합니다. 즉, 모델의 분포 내에서 특정 목적 함수 (예: 단백질 안정성, 프롬프트 정렬) 를 최적화하는 샘플을 찾기 위해 해당 분포 내에서 탐색이 필요합니다. 이러한 목적 함수는 종종 비용이 많이 들고 블랙박스 형태를 띱니다.

이러한 "잠재 공간 최적화 (Latent Space Optimisation, LSO)"에 대한 전통적인 접근 방식은 현대적인 확률적 보간 모델에 적용될 때 다음과 같은 중대한 장애물에 직면합니다:

1. 차원성: 변분 오토인코더 (VAE) 와 달리, 확산 및 흐름 매칭 모델은 종종 데이터 자체와 비교 가능한 차원을 가진 잠재 공간에서 작동하므로, 직접적인 최적화가 계산적으로 불가능합니다.
2. 통계적 불일치: 잠재 벡터의 단순한 선형 결합은 종종 생성 모델이 요구하는 통계적 속성을 위반하여 열악하거나 비현실적인 출력을 초래합니다.
3. 재학습 비용: 이러한 모델을 최적화하는 기존 방법들은 종종 생성 모델 자체를 파인튜닝하거나 보조 모델을 학습시켜야 하므로, 계산 비용이 많이 들고 일반화 능력을 제한합니다.

방법론: O3 (생성 모델의 출력에 대한 최적화)

본 논문은 재학습이나 파인튜닝 없이 사전 훈련된 연속 변수 생성 모델의 출력에서 직접 작동하는 샘플 효율적 블랙박스 최적화를 위한 프레임워크인 O3를 제안합니다. 핵심 혁신은 **대리 잠재 공간 (surrogate latent spaces)**의 구축에 있습니다.

1. 대리 잠재 공간

O3 는 유효하고 고품질의 생성물을 매핑하는 저차원 유계 유클리드 탐색 공간 $U$를 구축합니다. 이 과정은 세 가지 주요 단계를 포함합니다:

• 시드 잠재 (Seed Latents): 이 방법은 모델의 잠재 분포에서 $K$개의 "시드" 잠재 벡터 ($z_1, \dots, z_K$) 를 선택하는 것으로 시작합니다. 이러한 시드는 무작위로 선택되거나, 높은 점수를 기준으로 필터링되거나, 알려진 데이터 포인트의 역전 (inversion) 을 통해 얻을 수 있습니다.
• 통계적 보존을 위한 선형 결합: 생성된 샘플이 모델의 신뢰할 수 있는 작동 영역 내에 있도록 보장하기 위해, O3 는 **잠재 최적 선형 결합 (Latent Optimal Linear combinations, LOL)**을 사용합니다. 임의의 선형 결합 대신, O3 는 단위 노름 선형 결합으로 새로운 잠재 $z^*$를 구성합니다:
  $$z^* = \sum_{i=1}^K w_i z_i \quad \text{where} \quad \sum w_i^2 = 1$$
  이 제약 조건은 결과적인 잠재 $z^*$가 생성 모델이 기대하는 통계적 분포 (예: $N(0, I)$) 를 유지하도록 보장하여, REMBO 와 같은 방법들이 겪는 "불일치"를 방지합니다.
• 대리 차트 ($\phi_w$): 가중치의 서로 다른 스칼라 배수로 동일한 잠재를 표현하는 중복성을 피하고 표준 최적화를 용이하게 하기 위해, 단위 노름 가중치는 유계 유클리드 공간 $U = [0, 1]^{K-1}$에서 단위 초구 (unit hypersphere) $S^{K-1}_+$의 양의 사분면으로 매핑됩니다.
  • 본 논문은 Knothe–Rosenblatt (KR) 차트의 변형을 활용합니다. 이 쌍사상은 초구를 평탄화하여 베이지안 최적화나 CMA-ES 와 같은 표준 블랙박스 최적화기가 잠재 공간의 기하학적 구조 (코사인 유사도) 를 보존하면서 단순한 유클리드 공간에서 작동할 수 있게 합니다.

2. 워크플로우

1. 초기화: $K$개의 시드 잠재 벡터를 선택합니다.
2. 최적화: 대리 공간 $U$에서 표준 최적화기 (예: BO, CMA-ES) 를 실행합니다.
3. 생성: 각 후보 $u \in U$에 대해, 가중치 $w = \phi_w(u)$로 매핑하고, 잠재 $z^* = \sum w_i z_i$를 계산한 후 $x = g(z^*)$를 디코딩합니다.
4. 평가: 블랙박스 목적 함수를 사용하여 $x$에 점수를 매깁니다.

주요 기여

1. 모델 중립적 프레임워크: O3 는 사전 훈련된 모델에 수정이 필요 없이, 보간에 적합한 잠재 분포를 가진 임의의 연속 변수 생성 모델 (예: 확산, 흐름 매칭) 에 적용 가능합니다.
2. 통계적 유효성: 시드 결합에 단위 노름 제약을 부과함으로써, O3 는 검색 제안이 신뢰할 수 있는 생성에 필수적인 것으로 Bodin 등 [5]이 지적한 생성 모델의 잠재 통계에 부합하도록 보장합니다.
3. 대리 기하학: KR 차트의 도입은 잠재 매니폴드의 유사도 구조를 보존하는 유계 비제약 유클리드 탐색 공간을 제공하여 표준 최적화기의 사용을 가능하게 합니다.
4. 제로 재학습 비용: 파인튜닝 (예: DPO, Adjoint Matching) 이나 보조 모델을 요구하는 방법들과 달리, O3 는 표준 생성 이상의 미미한 계산 오버헤드만 발생합니다.

실험 결과

저자들은 이미지, 오디오, 비디오, 단백질 설계 작업 전반에 걸쳐 O3 를 검증했습니다:

• 이미지 생성 (Flux 및 Stable Diffusion):
  • O3 는 표준 "Best-of-N" 샘플링과 잠재 통계를 보존하지 않는 방법들 (예: REMBO) 보다 훨씬 뛰어난 성능을 보입니다.
  • 고차원 대리 공간 (예: 89 차원) 에서 O3 와 베이지안 최적화는 저차원 대안이나 무작위 탐색보다 훨씬 높은 점수의 이미지를 찾습니다.
  • O3 는 관련 학습 비용 없이 파인튜닝된 모델 (예: Importance Fine-tuning, DPOK) 과 비교 가능한 성능을 달성합니다.
• 단백질 설계 (BOLTZ-2 및 RFDIFFUSION):
  • 입체 구조 편향 (Conformation Biasing): BOLTZ-2 를 특정 칼모듈린 입체 구조 (1CLL) 로 편향시킬 때, O3 는 낮은 오라클 예산 ($N \le 100$) 에서 FK-스티어링 및 DPO 보다 우수한 성능을 보이며, 높은 예산에서도 경쟁력 있는 성능을 유지합니다. 이는 파인튜닝 없이 이루어집니다.
  • De Novo 설계: RFDIFFUSION 을 사용하여 600 잔기 단백질을 생성할 때, 표준 샘플링은 거의 회복 기준 (RMSE < 2.0 Å) 을 충족하는 설계를 산출하지 못합니다. 반면, O3 는 대리 공간에서 CMA-ES 를 사용하여 무작위 시드 잠재로도 성공적인 설계를 안정적으로 찾습니다 (중앙값 RMSE 약 1.08 Å).
• 일반화: 이 방법은 오디오 (StableAudio2.0) 와 비디오 (HunyuanVideo) 에 대해 매끄러운 저차원 매니폴드를 성공적으로 구축하여, 모달리티 간 적용 가능성을 입증했습니다.

의의 및 주장

본 논문은 O3 를 사전 훈련된 생성 모델을 최적화하기 위한 범용적인 "플러그 앤 플레이" 레이어로 위치시킵니다. 그 주요 의의는 검색 과정을 모델 학습과 분리하는 데 있습니다.

• 효율성: 대규모 모델을 재학습하거나 파인튜닝하는 막대한 계산 비용 없이, 과학 및 공학에서 흔히 발생하는 비미분 가능한 고비용 목적 함수에 대해 매우 샘플 효율적인 최적화를 가능하게 합니다.
• 접근성: 작업별 모델 적응의 필요성을 제거함으로써, O3 는 사용자가 생성 모델을 특정 과학적 또는 공학적 목표 방향으로 유도하는 장벽을 낮춥니다.
• 견고성: 이 방법은 새로운 매개변수를 학습하는 대신 잠재 공간의 기하학적 속성에 의존하므로, 최적화가 사전 훈련된 모델의 "안전" 영역 내에 머무르도록 보장합니다.

저자들은 대리 잠재 공간이 다양한 모달리티에 걸쳐 최적화를 위한 처리 가능한 구조화된 매니폴드를 제공하며, 블랙박스 최적화 작업에 대한 파인튜닝의 실용적인 대안을 제공한다고 결론지었습니다. 향후 작업으로는 적응형 시드 선택 및 다목적 최적화가 제안됩니다.