ProtQueSt: Query-Conditioned Retrieval-Augmented Generation for Protein Function Annotation

Motivation

• Known: LLM은 sequence-only 단백질 기능 예측 벤치마크에서 retrieval 기반 방법보다 지속적으로 성능이 낮으며, RAPM과 같은 retrieval-augmented 방법이나 ProteinCLIP, ProtST 같은 contrastive alignment 방법들이 protein-text 이해를 위해 활용되어 왔다.
• Gap: 기존 retrieval 기반 방법들은 embedding-space 상의 이웃 단백질로부터 주석을 그대로 transfer할 뿐, 동일한 단백질이라도 catalytic activity, domain architecture, cellular localization 등 질의(query)에 따라 다르게 기술되어야 한다는 점에 적응하지 못한다. 즉, 기존 접근은 query-independent한 nearest-neighbor transfer에 머물러 있다.
• Why: 단백질 기능은 본질적으로 query-dependent하기 때문에, 이를 반영하지 못하는 고정된 sequence-to-text 매핑은 신약 개발과 orphan protein 주석 등 실제 활용에서 부정확한 결과를 낳을 수 있어, query에 적응하는 retrieval 및 생성 방식이 중요하다.
• Approach: structure-aware retriever(ESM-2 + ProstT5 fusion)로 query-agnostic한 구조적 후보를 검색하고, FiLM 조건화와 query-pooled negative sampling으로 학습된 contrastive retriever로 query-specific한 후보를 별도로 검색한 뒤, 두 채널의 결과를 병합해 LLM에 입력하는 hybrid retrieval-augmented generation 방식을 취한다.

Achievement

Figure 5. Per-task Entity-BLEU by merging strategy (n=256,

1. Entity-BLEU 최고 성능: Prot-Inst-OOD 벤치마크에서 Entity-BLEU 48.79를 달성하여 RAPM 대비 +37% 향상되었으며, LLM-as-a-judge 점수도 최고치를 기록했다.
2. FiLM과 query-pooled negative sampling의 상호 필수성 입증: 두 요소 중 하나만 적용할 경우 structural baseline 대비 성능 향상이 없고, 둘을 함께 적용해야만 유의미한 개선이 나타남을 실험적으로 확인했다.
3. Structure-aware unified retriever 설계: ESM-2와 ProstT5 임베딩을 가중 코사인 융합(α=0.7)으로 결합해 RAPM의 dual-key(MMseqs2 + ESM-2) 파이프라인을 단일 dense score로 대체했다.
4. Hybrid retrieval 프레임워크 제안: structural retriever와 query-conditioned contrastive retriever의 상위-K 후보를 budget-matched(5+5)로 병합하여 LLM 프롬프트에 제공하는 실용적 추론 파이프라인을 구축했다.

How

Figure 1. ProtQueSt Architecture. (A) A structure-aware retriever fuses ESM-2 and ProstT5 embeddings for query-independe

• Structural retriever: 고정된 ESM-2(1280-d, evolutionary co-variation)와 ProstT5(1024-d, 3D fold 정보) 임베딩을 concat하여 h∈R2304로 융합하고, s_struct(i,j)=α·cos(p_i,p_j)+(1-α)·cos(e_i,e_j) (α=0.7)로 query-agnostic top-K 후보를 검색한다.
• Query-conditioned contrastive retriever: 동일한 fused protein embedding과 OpenAI text-embedding-3-large(3072-d)로 얻은 annotation/instruction 텍스트 임베딩을 two-tower 구조의 2-layer MLP adapter로 512-d 공유 공간에 투영한다.
• FiLM conditioning (ProtFiLM/BiFiLM): query 표현으로부터 scale γ와 shift β를 생성해 protein embedding을 차원별로 증폭/억제함으로써 질의에 따라 달라지는 protein representation z_p를 만든다.
• Query-pooled negative sampling: 서로 다른 query type 간 negative를 pooling하고 false-negative(의미적으로 유사한 annotation)를 필터링하여 query-specific한 contrastive 신호를 제공한다.
• 학습 목적함수: SigLIP loss를 사용해 protein-text compatibility score(cosine similarity + learnable temperature τ)를 최적화한다.
• 추론(Hybrid retrieval): structural retriever의 top-5와 contrastive retriever의 top-5를 결합한 10-item budget-matched context를 구성하여 사용자 질의와 함께 LLM에 입력, grounded functional description을 생성한다.

Originality

• 단백질 기능 주석을 단일 고정된 sequence-to-text 매핑이 아닌 query-conditioned retrieval 문제로 재구성한 문제 정의 자체가 새로운 프레이밍이다.
• FiLM(Feature-wise Linear Modulation)을 protein-text contrastive retrieval에 도입해, 동일 protein embedding에서 query에 따라 다른 feature를 증폭/억제하도록 한 조건화 메커니즘이 독창적이다.
• query-pooled negative sampling과 false-negative filtering을 결합하여 진정한 query-specific contrastive signal을 제공하는 학습 전략을 제시했다.
• structure-aware(ESM-2+ProstT5)와 query-conditioned contrastive 채널을 budget-matched hybrid로 결합하는 아키텍처가 RAPM의 dual-key 방식을 대체하는 새로운 설계이다.

Limitation & Further Study

• FiLM conditioning과 query-pooled negative sampling이 개별적으로는 효과가 없고 결합 시에만 효과가 있다는 결과는 두 메커니즘 간 상호작용에 대한 이론적 설명이 부족하며, 왜 그런지에 대한 심층 분석이 더 필요하다.
• 평가가 Prot-Inst-OOD라는 단일 benchmark 및 Entity-BLEU/LLM-as-a-judge라는 제한된 metric에 의존하고 있어, 다양한 실제 biological question(예: 희귀 orphan protein, novel fold)에 대한 일반화 검증이 부족하다.
• text-embedding-3-large와 같은 폐쇄형 상용 임베딩 모델에 의존하여 재현성과 비용 측면에서 제약이 있을 수 있다.
• 논문 발췌본에는 ablation의 정량적 세부 수치나 통계적 유의성 검정, downstream task(예: 실제 실험적 검증)에 대한 논의가 충분히 제시되지 않아 후속 연구에서 이를 보강할 필요가 있다.

Evaluation

총평: 단백질 기능 주석을 query-conditioned retrieval 문제로 재정의하고 FiLM 기반 contrastive retriever와 structure-aware retriever를 결합하여 state-of-the-art 성능을 달성한 실용적이고 참신한 연구로, 다만 단일 벤치마크에 대한 검증과 메커니즘의 이론적 분석 보강이 필요하다.