기술 요약: 생물학적 서열 설계를 위한 유연한 흐름 (FlexFlow)

문제 정의

기능적인 생물학적 서열(DNA, RNA, 단백질, 펩타이드)을 설계하는 것은 엄격한 진화적 및 생물물리학적 제약 조건 하에서 방대하고 이산적인 조합 공간을 탐색해야 하는 작업이다. 자기회귀 모델(예: ProtGPT2, ESM-2)이 생존 가능성을 입증했으나, 이들은 누적되는 오류, 느린 순차적 샘플링, 그리고 특히 길거나 구조적으로 복잡한 서열에 대한 제한된 전역적 일관성(global coherence) 문제를 겪는다.

대안적인 생성 프레임워크인 이산 확산(Discrete Diffusion) 및 이산 플로우 매칭(Discrete Flow Matching, DFM)은 더 빠른 비자기회귀적 샘플링을 제공한다. 그러나 생물학적 응용을 위한 기존의 이산 방법론들은 세 가지 결정적인 한계에 직면해 있다:

1. 제한된 유연성: 가변 길이 서열 생성에 어려움을 겪으며, 생성 과정에 대한 미세한 제어 능력이 부족하다.
2. 정보가 없는 결합(Uninformative Couplings): 생물학적 서열 요소 간의 구조화된, 도메인 특화된 관계(예: 진화적 치환 편향)를 무시하는 균등 또는 마스킹된 결합에 의존한다.
3. 일관성 없는 가이던스: 특정 기능적 특성을 향해 생성을 유도할 때, 복잡한 조절 또는 결합 모티프에 필요한 전역적 의존성을 포착하지 못하는 독립적인 토큰 단위 조정을 사용한다.

방법론: FlexFlow

저자들은 세 가지 핵심적인 아키텍처 혁신을 통해 이러한 한계점을 해결하도록 설계된 이산 플로우 매칭 프레임워크인 FlexFlow를 제안한다.

1. 구조화된 순방향 프로세스 (Structured Coupling)

비정보적인 베이스라인을 사용하는 대신, FlexFlow는 소스 분포를 "기울임(tilting)"으로써 구조적 귀납 편향(structural inductive bias)을 부과한다.

• 메커니즘: 이 방법은 표준 균등 결합을 도메인 특화된 치환 선호도(예: 단백질의 경우 BLOSUM62, DNA/RNA의 경우 JC69/HKY85)를 인코딩하는 전이 커널 $K_\gamma$로 정의된 구조적 결합 $\pi(\tilde{x}_0, x_1, K_\gamma)$로 대체한다.
• 효과: 소스 분포가 전체 심플렉스(simplex)가 아닌 생물학적으로 타당한 치환 근방(neighborhood)을 향하도록 편향된다. 이는 플로우 목적 함수나 훈련 절차를 수정하지 않고도 생성 경로를 선호되는 기능적 영역으로 유도한다. 이 결합은 커널 $K_\gamma$를 교체하는 것만으로 서열 분류 체계 전반에 걸쳐 일반화된다.

2. 잠재 편집 기반 속도 파라미터화 (Latent Edit-based Rate Parameterization)

가변 길이 생성을 지원하기 위해, 모델은 공유된 전역 잠재 변수에 조건화된 편집 연산(삽입, 삭제, 치환)을 통해 역방향 플로우 속도를 파라미터화한다.

• 계층적 파라미터화: 노이즈가 섞인 서열 $x_t$가 주어지면, 인코더 $\Psi(x_t)$를 통해 전역 잠재 컨텍스트 $r$이 추론된다. 이 잠재 변수 $r$은 각 위치에서의 편집 연산의 범주형 분포 $p_\Phi(op_i | x_i^t, r)$를 조건화한다.
• 트랙터빌리티(Tractability): 이 계층적 접근 방식은 공유된 잠재 변수 $r$을 통해 토큰 수준의 연산 간 결합을 유도하면서도, 위치 간의 조건부 독립성을 유지하여 모델의 계산 가능성을 보장한다.
• 테스트 타임 제어 (디리클레 사전 확률): 연산 확률은 디리클레 사전 확률을 통해 모델링된 위치 의존적 온도 파라미터 $\beta_i$를 사용하여 재조정된다. 이를 통해 추론 시(inference) 재학습 없이도 삽입, 치환, 삭제의 상대적 비율을 조절하는 "예산 컨트롤러(budget controller)" 역할을 수행할 수 있다.

3. 잠재 분류기-자유 가이던스 (Latent Classifier-Free Guidance)

저자들은 표준 속도 공간 가이던스와 대조되는, 완전히 연속적인 잠재 공간 내에서 작동하는 새로운 가이던스 메커즘을 도입한다.

• 메커니즘: 훈련 중에 인코더 $\Psi$는 조건 $c$가 주어졌을 때의 조건부 잠재 $r_c$와 무조건부 잠재 $r_\emptyset$를 모두 예측하도록 학습된다. 테스트 시에는 가이드된 잠재 변수가 $r_{LG} = r_\emptyset + w(r_c - r_\emptyset)$로 계산된다.
• 장점: 이 접근 방식은 가이던스 신호가 공유된 잠재 변수 $r$을 통해 모든 토큰별 연산에 동시에 전달되므로 전역적 구조적 일관성을 보장한다. 이는 토큰 단위 가이던스의 불일치 문제를 피하며, 구조적으로 속도의 비음성(non-negativity)을 보존한다.
• 이론적 근거: 저자들은 이 잠재 보간(latent interpolation)이 잠재 공간 내의 암시적 분류기의 로그 가능도(log-likelihood)에 대한 경사 상승 단계에 해당함을 증명한다.

주요 기여

1. 구조적 결합: 균등 결합을 대체하여 생성 경로를 기능적 근방으로 유도하는 생물학적 정보 기반의 사전 확률을 도입함.
2. 잠재 편집 기반 속도: 공유된 전역 잠재 변수와 디리클레 사전 확률 온도 스케일링을 통해 가변 길이 생성을 가능하게 하는 파라미터화 방식 제공.
3. 잠재 가이던스: 전역적 일관성을 보장하기 위해 연속 잠재 공간에서 작동하는 가이던스 체계를 제시하며, 이를 속도 공간의 분류기-자유 가이던스의 대안으로 제시함.
4. 새로운 벤치마크: MHC 조건부 펩타이드 서열을 생성하고 평가하기 위한 peptide-MHC II 벤치마크 도입.
5. 실증적 검증: 밀도 추정, 무조건부/조건부 DNA 생성, 펩타이드 생성 등 다양한 작업에서 최첨단(SOTA) 성능 입증.

실험 결과

저자들은 다양한 작업에 대해 FlexFlow를 확산 및 플로우 베이스라인(예: Bit Diffusion, Dirichlet FM, Fisher FM, Gumbel-Softmax FM)과 비교 평가하였다.

• 토이 데이터셋: 1D 및 2D 이산 패턴(Sine-wave, Spiral, Rings, Checkerboard)에서 FlexFlow는 가장 낮은 Wasserstein 거리($W_2$)를 달성하며 위상적 특징을 정확하게 재구성했다.
• 무조건부 DNA 설계 (인핸서): 초파리 뇌 및 인간 멜라노마 인핸서 데이터셋을 대상으로 평가했다. FlexFlow는 가장 낮은 Fréchet Biological Distance (FBD) (멜라노마 3.4, 초파리 뇌 4.2)를 기록하며 Dirichlet FM 및 Fisher FM을 능가했다.
• 조건부 DNA 설계 (프로모터): 전사 프로파일이 주어졌을 때 프로모터 서열을 생성했다. 잠재 가이던스를 적용한 FlexFlow는 0.022의 평균 제곱 오차(MSE)를 달와서 Dirichlet FM(0.026) 및 DDSM(0.033)보다 우수한 성능을 보였다.
• Peptide-MHC II 생성: 새로운 peptide-MHC II 벤치마크에서, 잠재 가이던스를 사용한 FlexFlow는 0.66의 DeepMHCII 판별기 점수(Dirichlet FM의 0.21 대비)와 0.84의 Fréchet Protein Distance (FPD)를 달성했다. 이 결과는 잠재 가이던스가 결합 타당성(binding plausibility)을 우선시하는 반면, 속도 가이던스는 분포 커버리지를 선호하는 다른 트레이드오프를 제공함을 나타낸다.
• 효율성: 추론 속도는 베이스라인과 유사한 수준이었다 (500 길이의 DNA 서열에 대해 100 스텝 기준 약 4.0초).

의의 및 주장

본 논문은 생물학적 도메인 지식을 플로우 매칭 프레임워크에 직접 통합함으로써 기존의 이산 생성 방법론의 핵심적인 한계를 해결한다고 주장한다. 구조적 결합을 활용함으로써, 이 방법은 훈련 목적 함수를 변경하지 않고도 생성 과정을 타당한 생물학적 영역으로 편향시킨다. 잠재 편집 기반 파라미터화와 잠재 가이던스 체계는 이전의 방법들이 갖지 못했던 유연한 길이의 생성과 일관되고 미세한 제어 메커니즘을 제공한다.

저자들은 FlexFlow를 무조건부 및 조건부 설정 모두에서 강력한 확산 및 플우 매칭 베이스라인을 능가하는 견고한 생물학적 서열 설계 프레임워크로 포지셔닝한다. 또한, 더 적은 샘플링 단계에서 고품질의 서열 생성을 가능하게 하기 위해 증류(distillation) 기반 훈련 목적 함수를 적용하는 것을 향후 연구 과제로 언급했다. 본 연구는 프리프린트로 발표되었으며, 저자들은 핀란드 연구 위원회 및 관련 바이오 디자인 AI 프로젝트의 지원을 받았음을 밝히고 있다.