[논문 리뷰] Supervised Training of Conditional Monge Maps

Charlotte Bunne, Andreas Krause외 1명|arXiv (Cornell University)|2022. 06. 28.

Adversarial Robustness in Machine Learning인용 수 14

한 줄 요약

이 논문은 부분 입력 볼록 신경망(PICNNs)을 사용하여 조건부 몽헤 맵을 훈련하는 지도 학습 프레임워크인 CondOT을 소개한다. 이는 확률 측도 간의 맥락 의존적 최적 운반을 모델링하기 위해 설계되었다. 가우스 근사에 기반한 맥락 인식 초기화와 딥 세트 조합 기법을 통해, 단일 세포에 대한 새로운 약물 조합이나 유전자 변형에 대한 반응을 정확하게 예측할 수 있다.

ABSTRACT

Optimal transport (OT) theory describes general principles to define and select, among many possible choices, the most efficient way to map a probability measure onto another. That theory has been mostly used to estimate, given a pair of source and target probability measures $(μ, ν)$, a parameterized map $T_θ$ that can efficiently map $μ$ onto $ν$. In many applications, such as predicting cell responses to treatments, pairs of input/output data measures $(μ, ν)$ that define optimal transport problems do not arise in isolation but are associated with a context $c$, as for instance a treatment when comparing populations of untreated and treated cells. To account for that context in OT estimation, we introduce CondOT, a multi-task approach to estimate a family of OT maps conditioned on a context variable, using several pairs of measures $\left(μ_i, ν_i ight)$ tagged with a context label $c_i$. CondOT learns a global map $\mathcal{T}_θ$ conditioned on context that is not only expected to fit all labeled pairs in the dataset $\left\{\left(c_i,\left(μ_i, ν_i ight) ight) ight\}$, i.e., $\mathcal{T}_θ\left(c_i ight) \sharp μ_i \approx ν_i$, but should also generalize to produce meaningful maps $\mathcal{T}_θ\left(c_{ ext {new }} ight)$ when conditioned on unseen contexts $c_{ ext {new }}$. Our approach harnesses and provides a novel usage for partially input convex neural networks, for which we introduce a robust and efficient initialization strategy inspired by Gaussian approximations. We demonstrate the ability of CondOT to infer the effect of an arbitrary combination of genetic or therapeutic perturbations on single cells, using only observations of the effects of said perturbations separately.

연구 동기 및 목표

• 단일 세포에 대한 새로운 약물 또는 유전자 조합의 영향을 예측하는 데 있어, 개별 변형의 영향만 관측 가능한 상황에서의 도전 과제를 해결하기 위해.
• 원천-대상 쌍이 고정된 최적 운반을 초월하여, 치료법, 복용량 또는 유전자 변형과 같은 맥락 변수에 따라 운반 맵을 조건화함으로써 일반화를 가능하게 하기 위해.
• 맥락 벡터에 조건화된 글로벌이고 파arameterized된 운반 맵을 학습하여, 새로운 맥락에 대한 일반화를 가능하게 하기 위해.
• 입력 볼록 신경망에서의 훈련 불안정성을 해결하기 위해, 가우스 근사에 기반한 안정적이고 폐쇄형 해를 갖는 초기화 전략을 도입하기 위해.

제안 방법

• 입력 공간 $\mu$ 에서 볼록성을 보장하기 위해, 조건부 운반 맵 $\mathcal{T}_{\theta}(c)$ 를 부분 입력 볼록 신경망(PICNN)의 그래디언트로 파aram터화한다.
• PICNN의 가중치와 편향을 맥락 벡터 $c$ 가 조절함으로써, 맥락에 맞는 운반 맵을 생성하는 하이퍼네트워크를 효과적으로 구현한다.
• 입력 측도 $\mu_i$ 와 목적 측도 $\nu_i$ 의 가우스 근사를 기반으로 한 폐쇄형 해를 사용하여 PICNN의 파라미터를 초기화함으로써, 안정적이고 의미 있는 초기 맵을 확보한다.
• 훈련 과정에서, 음수 가중치를 0으로 설정하는 가중치 클리핑을 통해 PICNN의 볼록성을 강제로 유지하고, 이중 네트워크의 볼록성은 음수 가중치에 대한 미분 가능한 페널티를 통해 유연하게 완화한다.
• 다양한 변형 임베딩을 통합하여 공동 맥락 표현을 생성하기 위해 딥 세트 아키텍처를 활용함으로써, 새로운 조합의 변형에 대한 일반화를 가능하게 한다.
• 훈련 효율성과 일반화 성능을 향상시키기 위해, 고차원 단일 세포 RNA-seq 데이터를 PCA 또는 오토에인코더를 사용하여 저차원 잠재 공간으로 투영한다.

실험 결과

연구 질문

• RQ1 단일 변형의 관측만으로도, 새로운 조합의 변형에 대해 일반화할 수 있는 단일 파arameterized 운반 맵을 학습할 수 있는가?
• RQ2 가우스 근사를 통한 폐쇄형 해를 활용하여, 입력 볼록 신경망의 최적 운반 훈련을 어떻게 안정화시킬 수 있는가?
• RQ3 딥 세트 기반의 조합 모듈이 하나의 온핫 인코딩에 비해 새로운 변형 조합에 대해 얼마나 잘 일반화되는가?
• RQ4 제안된 조건부 최적 운반 프레임워크가 새로운 약물 또는 유전자 변형에 대한 단일 세포 반응 예측에서 기존 베이스라인을 능가하는가?

주요 결과

• CondOT는 딥 세트 모듈을 통해 개별 변형의 공동 임베딩을 학습함으로써, 새로운 조합의 변형에 대해 성공적으로 일반화하며, 제로샷 일반화에서 온핫 인코딩보다 뛰어난 성능을 보였다.
• 가우스 근사를 기반으로 한 제안된 초기화 전략은 PICNN 기반 운반 맵의 훈련 안정성과 수렴 속도를 크게 향상시켰다.
• MOFA+ 및 MOFA++ 데이터셋에서의 검증을 통해, CondOT는 새로운 약물 조합에 대한 단일 세포 반응 예측에서 최신 기술 수준의 성능을 달성하였다.
• 오토에인코더 기반 투영은 PCA보다 고차원 구조를 더 잘 유지하여, 후속 예측 작업에서 향상된 성능을 이끌어냈다.
• 학습 중에 관측되지 않은 변형에 대해서도 모델이 효과적으로 일반화되어, 제로샷 설정에서의 강건성을 입증하였다.
• 이중 네트워크의 음수 가중치에 대한 미분 가능한 페널티를 도입함으로써, 볼록성 제약을 효과적으로 완화하면서도 이론적 보장을 유지할 수 있었다.