E1: Retrieval-Augmented Protein Encoder Models

Motivation

• Known: ESM-2/ESM-C 같은 single-sequence protein language model은 대규모 protein database로부터 진화적 패턴을 학습하여 구조·기능 정보를 담은 representation을 생성하며, 이러한 likelihood는 기능과 상관관계를 가진다. MSA Transformer나 PoET 같은 retrieval-augmented 접근은 homologous sequence를 training/inference에 명시적으로 제공함으로써 진화적 맥락을 직접 활용해 성능을 향상시킬 수 있음이 알려져 있다.
• Gap: 기존 single-sequence PLM은 진화적 맥락을 모델 파라미터에 압축해 학습해야 하므로 parameter inefficiency, underrepresented family에서의 성능 저하, dataset bias 반영(phylogenetic bias) 등의 문제를 겪으며, fine-tuning으로 이를 보완하는 것은 비용이 크고 data-scarce 상황에서 비현실적이다. 또한 MSA Transformer나 MSA Pairformer 같은 기존 retrieval 기반 모델은 homologous sequence 간 alignment를 요구한다는 제약이 있다.
• Why: Protein engineering의 핵심은 sequence-function mapping인데, retrieval augmentation을 통해 진화적 맥락을 명시적으로 활용하면서도 alignment-free하고 fine-tuning 없이 다양한 규모의 family에 유연하게 적용 가능한 모델은 fitness prediction, variant ranking, 구조 예측 등 실질적인 protein engineering 응용에 즉각적인 가치를 제공한다.
• Approach: alignment 없이 homologous sequence들을 query sequence 앞에 concatenate하여 multi-sequence input을 구성하고, intra-sequence attention과 block-causal multi-sequence attention을 교대로 사용하는 Transformer 기반 masked language model을 대규모 protein atlas 데이터로 학습한다.

Achievement

Figure 2. Examples from CAMEO dataset where retrieval augmentation helps E1 identify contact it may have mispredicted wh

1. State-of-the-art zero-shot fitness 예측: ProteinGym 217개 Deep Mutational Scan assay에서 E1이 single-sequence 모드로도 ESM 계열을 능가하고, retrieval-augmented 모드에서는 PoET, MSA Pairformer 등 기존 retrieval 기반 모델도 능가함.
2. Unsupervised contact-map 예측 우수성: single-sequence 모드에서 ESM 계열보다 우수하며, retrieval 추가 시 추가적인 큰 성능 향상을 보임.
3. Parameter scaling 검증: 150M부터 600M 파라미터까지 모델 크기에 따라 성능이 일관되게 향상됨을 확인.
4. 유연한 추론 모드: 동일 모델이 single-sequence 또는 retrieval-augmented 모드로 fine-tuning 없이 전환 가능하여 fitness prediction, variant ranking, 구조 관련 embedding 생성에 모두 활용 가능.
5. 오픈 소스 공개: 150M/300M/600M 세 가지 E1 모델을 연구 및 상업적 사용을 위해 무료로 공개.

How

Figure 1. E1 Architecture. The E1 model can take in homologous sequences in addition to an input query sequence. The hom

• Homology search(Diamond)를 통해 PPA-1 및 UniRef 데이터셋에서 50% identity cluster 대표서열 기준으로 homolog set을 구성
• 각 training instance는 homolog set을 크기에 반비례하는 확률로 샘플링하고, cluster 대표서열을 90% ID subcluster 크기에 반비례하여 실제 서열로 치환하여 구성
• 최대 512개 서열을 <bos>/<eos> 특수 토큰과 sequence ID embedding으로 구분하여 하나의 multi-sequence instance로 tokenize
• 3개 layer마다 global block-causal attention을, 나머지는 intra-sequence attention을 적용하는 alternating attention 구조 사용 (Warner et al., 2025 방식)
• RoPE를 사용하되 intra-sequence attention layer에서는 각 서열마다 position ID를 1부터 재시작, global attention layer에서는 concatenated 전체 서열 내 절대 위치 사용
• BERT-style masking(80/10/10)을 적용한 masked language modeling으로 4조 token, WSD learning rate schedule과 Stable AdamW optimizer를 사용해 학습
• Curriculum learning으로 multi-sequence instance의 길이(8192→32768)와 서열 수(2→512)를 점진적으로 증가; 처음 1.5조 token은 PPA-1만, 이후 PPA-1:UniRef를 60:40으로 혼합
• ProteinGym v1.3의 217개 DMS assay에서 masked marginal method로 zero-shot fitness score 계산, Spearman correlation과 NDCG로 평가

Originality

• Alignment-free homologous sequence concatenation과 block-causal multi-sequence attention을 결합하여 MSA 기반 모델(MSA Transformer, MSA Pairformer)의 alignment 의존성을 제거함
• Intra-sequence attention과 block-causal global attention을 교대로 배치하는 alternating attention 구조를 protein encoder 도메인에 도입해 효율성과 성능을 동시에 확보
• 서열별 position ID 재시작 방식과 절대 위치 방식을 attention 유형에 따라 다르게 적용하는 RoPE 활용 방식
• Homolog set 샘플링 시 cluster 크기에 반비례한 가중치를 부여하는 데이터 구성 전략(PoET 방식 확장 및 PPA-1/UniRef 혼합 커리큘럼)
• 단일 모델이 fine-tuning 없이 single-sequence와 retrieval-augmented 모드 양쪽에서 SOTA를 달성하는 유연성

Limitation & Further Study

• 발췌된 본문에서는 최대 600M 파라미터까지만 실험되어 있어 더 큰 규모(수십억 파라미터급)에서도 scaling이 지속되는지는 불확실함
• Encoder 기반 masked language model 구조로, generative design task(서열 생성)에 대한 직접적 적용 가능성은 제한적일 수 있음
• Retrieval 성능은 homology search(Diamond) 품질과 database 커버리지에 의존하므로, 매우 novel하거나 orphan protein family에 대해서는 이득이 제한될 가능성
• 512개 서열, 특정 길이 budget 내로 제한된 multi-sequence instance 구성이 매우 큰 protein family나 매우 긴 서열에 대한 일반화에 제약을 줄 수 있음
• 후속 연구로 더 큰 모델 규모, 구조 정보 통합, generative 태스크로의 확장 등이 고려될 수 있음

Evaluation

총평: Alignment-free retrieval augmentation과 효율적인 alternating block-causal attention을 결합해 single-sequence 및 retrieval 모드 모두에서 SOTA를 달성하고 모델을 오픈소스로 공개했다는 점에서 실용적 가치와 학술적 기여도가 높은 연구다.