Contextualizing Biological Language Models across Modalities via Logit-Space Contrastive Alignment

Motivation

• Known: Pretrained BioFM(ESM-2, TCRLang, SELFormer 등)은 masked-token prediction을 통해 per-token 확률 분포를 제공하며, 이는 zero-shot mutation scoring, sequence design, 해석 등에 활용된다. 그러나 이 분포는 광범위한 unlabeled corpus로 학습되어 interaction partner, cellular environment, therapeutic intervention 같은 task-specific biological context에 자연스럽게 조건화되어 있지 않다.
• Gap: 기존 contextual matching 방법은 pooled embedding, contrastive latent space, task-specific prediction head를 사용해 이 likelihood interface를 훼손하며, 이는 score가 더 이상 token likelihood가 아니게 되어 mutated residue에 localize하거나 likelihood 기반 generation을 지원하지 못하게 만든다. 반면 conditional MLM 방식은 token interface는 보존하지만 reconstruction loss가 matched/mismatched context를 직접 구분하지 못하며, sparse하고 noisy한 paired biological data에서 국소적(local) functional signal을 포착하는 데 한계가 있다.
• Why: context-conditioned likelihood는 protein–ligand binding, TCR–peptide activity, drug resistance 등 치료적으로 중요한 문제에서 변이 및 결합 예측의 핵심 인터페이스이며, 이를 embedding 공간이 아닌 native token 확률 공간에서 정의함으로써 해석 가능성과 generation 능력을 유지한 채 sparse paired data 학습이 가능해진다는 점에서 중요하다.
• Approach: LogiCA는 gated cross-modal adapter를 이용해 서로 다른 vocabulary를 가진 두 BioFM의 native token head를 그대로 유지하면서, context-conditioned token log-likelihood를 매칭 점수로 사용하는 InfoNCE 유사의 logit-space contrastive objective를 정의한다.

Achievement

Figure 2. Two scaling regimes for protein–ligand LOGICA. (A) Held-out likelihood-margin trajectories during pretraining

1. Logit-space contrastive alignment 프레임워크 제안: pooled embedding이나 별도 예측 head 없이, native token head가 산출하는 context-conditioned log-likelihood를 직접 compatibility score로 사용하는 LogiCA를 제안하였다.
2. Mutation-local variant ranking에 대한 이론적 정당화: 동일 mutation site를 공유하는 변이 쌍에 대해 wild-type 공통 항이 상쇄되어 Bradley–Terry 형태의 랭킹 확률로 환원됨을 명제(Proposition 2.1)로 증명하였다.
3. 다중 도메인 실증 성능 향상: protein–ligand binding, TCR–peptide activity, drug-conditioned resistance prediction 세 과제에서 matched latent-contrastive 및 conditional-MLM baseline 대비 상당한 성능 향상을 보였으며, held-out-gene 단일 변이 약물 저항성 예측에서 AUC를 latent-space baseline의 ~0.55에서 ~0.65로 개선하였다.
4. 서로 다른 vocabulary/tokenizer 간 정렬: 공유 tokenizer, decoder, embedding space 없이도 서로 다른 vocabulary를 가진 모델(ESM-2, TCRLang, SELFormer) 간 sparse paired data 정렬을 가능하게 하였다.

How

Figure 1. Overview of LOGICA: pretrained biological language models are coupled by cross-modal adapters that preserve na

• 각 위치 i에서 context y에 조건화된 token 분포 πθ(·|x\i, y)를 정의하고, anchor-candidate 쌍에 대한 compatibility score s(a,c)를 InfoNCE 형태의 ranking distribution (Eq.1)에 대입한다.
• LogiCA의 score는 pooled latent inner product 대신 site-averaged conditional log-likelihood ℓA(x|y) (Eq.2)로 정의되며, A는 전체 시퀀스 또는 mutation site 등 특정 위치 집합으로 설정 가능하다.
• Mutation-local variant ranking을 위해 wild-type 대비 mutant token의 log-likelihood ratio를 mutation set M에 대해 평균한 context-conditioned mutation score sM (Eq.3)을 정의한다.
• gated cross-modal adapter가 서로 다른 두 BioFM(BioFM1, BioFM2)의 표현을 native token head와 호환되도록 변환하여, 각 모델의 pretrained likelihood interface를 보존한 채 상호 contrastive 학습을 수행한다.
• ESM-2, TCRLang, SELFormer 세 모델에 적용하여 protein–ligand binding, TCR–peptide ranking, drug-conditioned resistance scoring 실험을 수행한다.

Originality

• pooled embedding이나 별도 prediction head를 우회하여, contrastive alignment를 output logit(즉 native token 확률) 공간에서 직접 수행한다는 개념적 전환이 새롭다.
• mutation-local ranking 상황에서 wild-type 공통 항이 상쇄되어 InfoNCE/Bradley–Terry 형태로 환원됨을 수학적으로 증명한 점(Proposition 2.1)이 독창적이다.
• 서로 다른 vocabulary/tokenizer를 가진 이종 biological language model들을 공유 embedding space 없이 정렬한다는 아이디어는 기존 CLIP류 latent alignment와 차별화된다.

Limitation & Further Study

• 발췌된 본문에서는 gated cross-modal adapter의 구체적 구조, 학습 안정성, 계산 비용에 대한 세부 설명이 충분히 제시되지 않아 확장성 평가가 어렵다.
• position-aligned substitution만을 고려하고 insertion/deletion은 배제하고 있어, indel이 포함된 변이나 구조적 변화가 큰 경우로의 일반화가 제한적이다.
• 세 가지 특정 도메인(protein-ligand, TCR-peptide, drug resistance)에 대해서만 검증되었으며, 더 다양한 modality 조합이나 더 큰 규모의 BioFM으로의 일반화는 후속 연구로 남아있다.
• baseline 대비 성능 향상의 통계적 유의성이나 hyperparameter(τ 등) 민감도에 대한 분석이 발췌 부분에서는 드러나지 않는다.

Evaluation

총평: 사전학습된 biological language model의 핵심 자산인 token-level likelihood interface를 훼손하지 않으면서 context-conditioned contrastive learning을 가능케 한 개념적으로 신선하고 실용적인 기여이며, 다양한 biological modality에 걸친 실증 결과가 이를 뒷받침한다.