기술 요약: 리워드 트랜스포트(Reward Transport) – 플로우 매칭에서의 노이즈 공간 정렬을 통한 속성 제어

1. 문제 정의

플로우 매칭(Flow Matching, FM) 및 관련 확산 기반 생성 모델에서, 커플링(coupling)—훈련 중에 어떤 노이즈 벡터가 어떤 데이터 포인트와 쌍을 이룰지를 결정하는 규칙—은 일반적으로 계산적 휴리스틱으로 취급됩니다. 표준 방식들은 단순성을 위해 독립적인 페어링을 사용하거나, 궤적을 직선화하고 분산을 줄이기 위해 미니배치 최적 운송(Optimal Transport, OT)을 사용합니다. 이러한 패러다임에서 커플링은 오로지 훈련을 용이하게 하기 위해 최적화될 뿐, 학습된 플로 필드(flow field)에 특정 구조적 속성을 인코딩하기 위해 최적화되지 않습니다.

결과적으로, 제어 가능한 생성(예: 특정 logP 또는 QED 점수를 가진 분자 생성)을 달升하기 위해서는 다음과 같은 명시적인 추론 시 메커니즘이 필요합니다:

• 클래시파이어 프리 가이던스(Classifier-free guidance): 타겟 속성 값을 입력으로 요구함.
• 그래디언트 가이던스(Gradient guidance): 샘플링 매 단계마다 속성 예측기의 그래디언트를 계산함.
• 강화 학습(Reinforcement Learning, RL): 오라클(oracle)에 대항하여 생성기를 미세 조정하며, 이는 종종 샘플링 중에 오라클 쿼리를 요구함.

이러한 방법들은 추론 시의 오버헤드, 외부 보상 모델 또는 오라클에 대한 의존성, 혹은 복잡한 훈련 파이프라인을 초래합니다. 본 논문은 커플링 자체가 정렬 인터페이스(alignment interface) 역할을 수행하여, 훈련 중에 제어 가능한 구조를 플로 필드에 직접 임베딩함으로써 추가적인 계산이나 오라클 호출 없이도 제어가 가능하게 할 수 있는지 질문합니다.

2. 방법론: 리워드 트랜스포트 (Reward Transport)

저자들은 노이즈-데이터 커플링을 태스크 구조를 플로 필드에 주입하는 메커니즘으로 재정의하는 방법론인 리워드 트랜스포트를 제안합니다.

2.1 속성 정렬 최적 운송 커플링

노이즈와 데이터를 독립적으로 페어링하는 대신, 리워드 트랜스포트는 타겟 분자 속성 $y(x)$ (예: logP, QED)를 기반으로 한 **단조 커플링(monotone coupling)**을 구축합니다.

1. 스칼라 키(Scalar Keys): 각 노이즈 벡터 $z$(토큰 임베딩의 시퀀스)에 대해 스칼라 정렬 키 $s(z)$를 계산합니다. 주요 선택지는 평균 풀링된 노이즈의 유클리드 노름입니다: $s(z) = \|\bar{z}\|_2$.
2. 순위 기반 페어링(Rank-Based Pairing): 노이즈 벡터 집합 $\{z_n\}$과 속성 $\{y_n\}$을 가진 분자 집합 $\{x_n\}$이 주어졌을 때, 두 집합 모두 각각의 키($s$와 $y$)에 따라 정렬됩니다. 그 후, $k$번째로 작은 $s$를 가진 노이즈 벡터는 $k$번째로 작은 $y$를 가진 분자와 페어링됩니다.
3. 훈련 목적 함수: 플로 매칭 모델 $v_\theta$는 이러한 정렬된 쌍을 운송하도록 훈련됩니다. 손실 함수는 엔드포인트 예측(연속 임베딩 공간 내)에 대한 표준 MSE와, 디코딩된 토큰이 유효한 분자를 형성하도록 보장하는 교차 엔트로피(Cross-Entropy, CE) 손실을 결합합니다.

2.2 방향 조건화 (Direction Conditioning)

모델이 추론 중에 노이즈 공간 좌표 $s$에 접근할 수 있도록 하기 위해, **방향 MLP(Direction MLP)**가 도입됩니다. 이 2층 MLP는 정규화된 스칼라 $s$를 $d$차원 임베딩으로 매핑하며, 이는 시간 및 위치 임베딩과 함께 트랜스포머의 입력에 더해집니다. 이를 통해 모델은 단조 페어링에 사용된 특정 노이즈 좌표에 따라 자신의 속도 필드(velocity field)를 조건화할 수 있습니다.

2.3 추론 메커니즘

추론 시, 사용자는 스칼라 방향 신호 $s^*$ (타겟 속성 값이 아닌 노이즈 공간의 좌표)를 지정합니다. 모델은 $s(z) \geq s^*$인 노이즈 벡터들로부터 플로 필드를 적분하여 분자를 생성합니다.

• 이론적 보장 (Proposition 1): 학습된 플로가 $s$와 $y$ 사이의 순위 대응 관계를 완벽하게 보존하는 극한(coupling-preserving limit)에서, $s \geq s^*$로 조건화하는 것은 **교차 엔트로피 방법(Cross-Entropy Method, CEM)**의 1단계 선택과 동일한 분포를 유도합니다. 구체적으로, 이는 생성된 분포를 속성 분포의 상위 $\epsilon$ 분율로 절단(truncate)하며, 여기서 $\epsilon$은 임계값 $s^*$에 의해 결정됩니다.
• 제로 오버헤드: 이 제어는 추론 시 오라클, 보상 모델, 그래디언트 계산, 또는 거부 샘플링(rejection sampling)을 필요로 하지 않습니다.

3. 주요 기여

1. 정렬 인터페이스로서의 커플링: 본 논문은 플로 매칭에서의 노이즈-데이터 커플링을 단순한 훈련 휴리스틱에서 분포 수준의 정렬 인터페이스로 재구성합니다. 커플링을 속성에 기반하여 선택하는 것이 학습된 플로 필드를 조직하여 노이즈 공간 좌표가 속성 값에 대응하도록 함을 입증합니다.
2. 단일 노브(Single-Knob) 분자 스티어링: 이 방법은 단일 스칼라 노이즈 좌표를 사용하여 분자 속성(logP 및 QED)에 대한 단조 제어를 달성합니다. 결정적으로, 동일한 노브가 서로 다른 타겟에 대해 반대되는 구조적 반응을 생성합니다: $s$를 증가시키면 logP(친유성)를 위해 더 큰 분자가 생성되지만, QED(약물 유사성)를 위해서는 더 작은 분자가 생성됩니다. 이는 일반적인 크기 편향(size bias)이 아님을 입증하며, 커플링이 타겟 특이적 구조를 플로에 기록함을 확인시켜 줍니다.
3. 이산 시퀀스를 위한 리워드 트랜스포트: 저자들은 가변 길이 SELFIES 토큰 시퀀스에 대해 이를 구현하며, 이산 데이터를 생성하는 플로 모델에서 발생하는 두 가지 특정 실패 모드를 식별하고 해결했습니다:
   • 암시적 신호 소멸(Implicit Signal Erasure): Pre-LayerNorm 아키텍처에서는 크기 정보($\|\bar{z}\|$에 인코딩된)가 소멸됩니다. 방향 MLP는 학습된 방향 패턴을 주입함으로써 이를 해결합니다.
   • 길이 혼동(Length Confound): 패딩 토큰에 대한 unmasked MSE 손 loss가 없으면, 모델은 노름 제약을 만족하기 위해 단순히 시퀀스 길이를 늘리는 지름길(shortcut)을 학습합니다. Unmasked MSE는 모델이 패딩 위치의 원시 노이즈를 회귀하도록 강제하여, 크기 신호를 고정하고 이러한 지름길을 방지합니다.

4. 실험 결과

방법론은 ZINC-250K 및 GuacaMol 데이터셋에서 평가되었습니다.

• 제어 성능:
  • logP: ZINC-250K에서 $s$를 -3에서 +7까지 스윕(sweep)하면 평균 logP가 1.50에서 5.44(+137%)로 이동했으며, $s$와 평균 속성 간의 Spearman 상관계수 $\rho=1.000$을 보였습니다. 효과 크기(Cohen's $d$)는 1.63에 달했습니다.
  • QED: 동일한 노브가 QED를 단조적으로 변화시켰습니다 (예: ZINC에서 $s=5$일 때 +0.059). $\rho=1.000$을 기록했습니다.
  • 구조적 특이성: logP의 경우, $s$를 증가시키면 헤비 원자(heavy atoms) 수가 증가했습니다 (12 $\to$ 23). QED의 경우, $s$를 증가시키면 헤비 원자 수가 감소했습니다 (23 $\to$ 16).
• 일반화: Guacaocal(1.6M 분자)에서도 결과가 재현되었으며, logP에 대해 더 강력한 효과 크기($d=2.31$)와 유효 범위 내에서의 일관된 QED 단조성을 보여주었습니다.
• 베이스라인과의 비교:
  • 조건부 플로 매칭(Conditional Flow Matching) (속성 값이 명시적 입력으로 들어가는 방식)과 비교했을 때, 리워드 트랜스포트는 유사한 단조성($\rho=1.000$)을 달성하지만, 명시적 타겟 적중보다는 분포 절단을 통해 작동합니다.
  • **$\epsilon$-예측 확산($\epsilon$-prediction diffusion)**과 비교했을 때, 본 방법은 커플링 수준의 정렬이 $\epsilon$-예측에는 구조적으로 부재함을 보여줍니다. 왜냐나 베이즈 최적 타겟(Bayes-optimal target)이 커플링과 독립적이어서 그래디언트 신호를 약화시키기 때문입니다.
• 훈련 역학: 순위 보존(rank preservation, $s$와 생성된 $y$ 사이의 상관관계)은 미세 조정 중 빠르게 나타나며(1 에포크 이내), 1.0 미만(예: logP의 경우 0.57)에서 정체됩니다. 이는 Proposition 1에서 설명된 "분포적 근사(distributional approximation)" 격리와 일치합니다.

5. 의의 및 주장

본 논문은 리워드 트랜스포트가 플로 매칭에서의 커플링을 정렬 구조를 전달할 수 있는 독자적인 분포 수준의 제어 채널로 확립한다고 주장합니다.

• 효율성: 추론 시 오라클 호출, 보상 모델, 또는 반복적 최적화를 제거하여, RL이나 가이던스 기반 방법들에 비해 계산 비용을 크게 줄입니다.
• 메커니즘: 제어 가능한 생성의 "스티어링(steerability)"이 반드시 명시적인 조건부 입력에 의존할 필요는 없음을 입증합니다. 대신, 훈련 커플링을 통해 모델의 기하학적 구조에 직접 심어질 수 있습니다.
• 한계 및 범위: 저자들은 연구 범위를 겸허하게 설정했습니다. 이 방법은 1D 분포적 노브(분포를 절단함)를 제공하는 것이지, 개별 샘플에 대한 정확한 값 할당을 제공하는 것이 아닙니다. 이 방법은 크기나 극성 같은 전역적 분자 특성과 상관관계가 있는 스칼라 속성에는 잘 작동하지만, 복잡한 구조를 구분할 수 없는 합성 가능성(SA)과 같은 위상적 속성에는 작동하지 않습니다. 또한, 이 방법이 $\epsilon$-예측 타겟과는 호환되지 않음을 보여줌으로써, 인터페이스가 네트워크의 특정 예측 타겟에 의존한다는 점을 명시했습니다.

요약하자면, 이 연구는 제어 가능한 생성의 패러다임을 "추론 시 궤적을 가이드하는 것"에서 "훈련 분포를 정렬하는 것"으로 전환하여, 플로 매칭에서 속성 제어를 위한 원칙적이고 오버헤드가 없는 메커니즘을 제시합니다.