PRIMA: a bidirectional state-space architecture and training approach for sequence modelling of protein-protein interactions

Motivation

• Known: Transformer 기반 PLM(ESM-1, ESM-2 등)은 masked language modeling(MLM)을 통해 진화적·구조적·기능적 제약을 학습하는 데 성공적이었으나, self-attention의 O(L^2) 복잡도로 인해 수천 개 아미노산에 달하는 다중 체인 복합체 처리에 한계가 있다.
• Gap: 기존 연구(MINT, LC-PLM, ProtMamba)들은 각각 transformer 기반 cross-attention, 그래프 random walk 기반 다중 단백질 학습, MSA 기반 causal Mamba 등을 시도했지만, bidirectional SSM을 실제 concatenated PPI 시퀀스에 직접 학습시키고 PPI 표현 품질과 binding affinity 예측에서 평가한 연구는 없었다.
• Why: PPI는 면역 반응, 신호 전달, 대사 조절, 유전자 발현 등 핵심 생물학적 과정을 지배하며, 이를 효율적으로 모델링할 수 있는 확장 가능한(scalable) foundation model은 신약 개발과 질병 기전 이해에 중요한 도구가 될 수 있다.
• Approach: BiMamba 블록을 쌓은 아키텍처를 설계하고, 단일 단백질 서열(UniRef50)로 일반 표현을 학습한 뒤 concatenated PPI 쌍(STRING)으로 상호작용 특화 문맥을 학습하는 2단계 MLM 학습 전략을 적용했다.

Achievement

Figure 1. Performance of PRIMA and PET on binding prediction

1. 성능 우위: 약 8M 파라미터 규모에서 PRIMA가 동일 파라미터 수의 transformer 기반 baseline(PET, ESM-2 8M 아키텍처)보다 표준 길이 입력, 긴 시퀀스, length extrapolation 과제 모두에서 더 낮은 perplexity와 우수한 binding affinity 예측 성능을 보였다.
2. 계산 효율성: 선형 복잡도(O(L))의 SSM 구조 덕분에 동일 성능을 훨씬 낮은 계산 비용으로 달성하여, 단일 GPU 환경에서 긴 concatenated PPI 시퀀스 학습이 가능함을 입증했다.
3. 위치 인코딩 불필요: positional encoding 없이도 recurrent state를 통해 자연스러운 length extrapolation을 달성하여, RoPE 기반 모델이 학습 분포를 벗어난 길이에서 성능이 저하되는 문제를 해결했다.

How

Figure 2. Comprehensive Architecture of PRIMA. (1) Macroscopic pipeline showing the embedding layer, lack of positional

• 아키텍처: N개의 bidirectional Mamba(BiMamba-S) 블록, 토큰 임베딩 레이어, tied LM head로 구성. 각 블록은 forward/reverse 두 방향의 독립적 Mamba pass를 수행하며 입출력 projection은 공유하되 SSM 파라미터는 방향별로 분리, 출력을 합산 후 residual connection 적용.
• 토큰화: 아미노산 단위 tokenization(29개 vocabulary), PPI 쌍은 <cls>+[A]+<inter>+<cls>+[B] 형태로 구조화.
• 위치 인코딩 없음: SSM의 recurrent state에 위치 정보가 암묵적으로 인코딩됨.
• 모델 규모: PRIMA는 7.3M 파라미터(N=16, D=256), PET는 ESM-2 8M 아키텍처(7.8M 파라미터, N=6, D=320)로 파라미터 수를 맞춤.
• 데이터: UniRef50에서 길이 50–1000 AA로 필터링한 500만 서열(UniRef50 5M)로 사전학습, STRING v12에서 길이 200–1000 AA 조건을 만족하는 쌍을 stratified degree-weighted sampling으로 추출한 1200만 PPI 쌍(STRING 12M)으로 후속학습.
• 학습 전략: BERT 프로토콜을 따르는 15% masking MLM, phase 1은 UniRef50 5M에서 50 epoch, phase 2는 STRING 12M에서 8 epoch, 단일 80GB H100 GPU에서 학습.
• 평가: STRING TEST SS(8–32,000 AA 길이의 236k 개별 서열)와 STRING TEST PPI(138k PPI 쌍, 일부는 학습 분포 밖의 길이 포함)에서 zero-shot perplexity 측정, binding affinity 예측 과제로 다운스트림 성능 평가.

Originality

• Bidirectional SSM(BiMamba)을 실제 concatenated PPI 서열에 직접 적용하고 PPI 표현 품질 및 binding affinity 예측까지 평가한 최초의 시도로, 기존 MINT(transformer 기반 cross-attention), LC-PLM-G(그래프 random walk 기반, 실제 concatenated 복합체 미학습), ProtMamba(causal, MSA 기반, cross-protein 상호작용 미지향)와 명확히 차별화됨.
• 단일 서열 사전학습 후 PPI 쌍 후속학습이라는 2단계 학습 전략을 SSM 아키텍처에 적용.
• 동일 파라미터 수, 동일 데이터, 동일 프로토콜로 transformer baseline(PET)과 직접 비교하는 통제된 실험 설계.

Limitation & Further Study

• 모델 규모가 8M 파라미터에 불과한 proof-of-concept 수준으로, 실제 SOTA PLM(수억~수십억 파라미터급)과의 비교나 확장성 검증이 부족함.
• 학습 데이터가 UniRef50과 STRING의 서브샘플(각각 500만, 1200만)로 제한되어 있어, 전체 데이터셋 규모에서의 성능이나 대규모 학습 시의 거동은 알 수 없음.
• MINT와 같은 대규모 transformer 기반 PPI 모델(ESM-2 650M 기반)과의 직접적 성능 비교가 제시되지 않아, 실제 활용 가능한 SOTA 대비 우위인지는 불명확함.
• 후속 연구로는 모델 및 데이터 규모 확장, 다양한 다운스트림 과제(예: mutational effect prediction, 구조 예측과의 연계) 평가, 그리고 3개 이상 체인의 복합체로의 일반화 검증이 필요함.

Evaluation

총평: PPI 모델링을 위한 선형 복잡도 SSM 아키텍처를 최초로 체계적으로 검증한 proof-of-concept 연구로, 작은 규모임에도 명확한 방법론적 기여와 통제된 비교 실험을 통해 BiMamba의 잠재력을 설득력 있게 보여준다. 다만 대규모 확장 및 SOTA 모델과의 직접 비교가 후속 과제로 남아있다.