생성과 예측을 동시에 구현한 효율적인 인과 단백질 언어 모델 Proust

초록

Proust는 309 M 파라미터 규모의 인과( causal) 단백질 언어 모델로, 그룹화된 쿼리 어텐션(GQA‑S2), 크로스‑레이어 값 레지듀얼, 키 오프셋 및 Canon 깊이별 인과 컨볼루션 등 최신 LLM 설계를 도입해 33 B 토큰을 40 GPU‑hour( B200)만에 학습하였다. ProteinGym 변이 예측에서 MLM 수준의 Spearman ρ = 0.390을 달성하고, indel 과제에서는 ρ = 0.521로 현존 최고 기록을 경신한다. 또한 로그잇 렌즈를 활용한 해석을 통해 위치별 엔트로피 분산이 검색 기반 보강의 효용을 예측한다는 흥미로운 통찰을 제공한다.

상세 분석

Proust는 전통적인 마스크드 언어 모델(MLM)과 인과 언어 모델(CLM)의 기능 격차를 해소하기 위해 설계된 309 M 파라미터 규모의 디코더‑전용 트랜스포머이다. 핵심 설계는 크게 네 가지 혁신으로 요약된다. 첫째, 그룹화된 쿼리 어텐션(GQA‑S2)을 채택해 키와 밸류를 동일한 선형 투사로 공유함으로써 KV 캐시 메모리를 크게 절감하고, 남은 파라미터를 헤드 차원 확대에 재배치한다. 헤드 차원은 96개의 NoPE(위치 정보 비포함)와 32개의 RoPE(회전 위치 임베딩)로 분할되어, 내용 매칭과 위치 매칭을 명확히 구분한다. RoPE는 K와 V 모두에 적용되고, V‑output에 역 RoPE(V O‑RoPE)를 부여해 절대 위치 정보를 복원한다. 둘째, 각 레이어의 값 출력에 첫 번째 레이어의 값을 가중합하는 크로스‑레이어 값 레지듀얼을 도입해 그래디언트 흐름을 개선하고 초기 레이어의 표현을 깊이 전파한다. 셋째, 키 오프셋 기법을 사용해 NoPE 키를 한 위치 앞으로 이동시켜, 단일 레이어에서 “A 다음에 B가 온다”와 같은 빅그램 패턴을 직접 인식하도록 한다. 넷째, Canon 레이어라는 깊이별 인과 컨볼루션을 세 곳(Attention 전, FFN 전, FFN 내부)에 삽입해 지역적인 패턴(예: 반복 모티프)을 효율적으로 학습한다. 이 컨볼루션은 채널별(depthwise)이며, 4‑step 커널을 사용해 연산량을 최소화한다.
학습 측면에서는 UniRef50과 메타게놈, 바이러스, 식물 등 다양한 서브셋을 결합한 33 B 토큰 데이터셋을 구축했으며, Muon 옵티마이저와 Polar Express 정규화를 결합해 스펙트럴 노름을 1로 고정, 높은 학습률(0.015)에서도 안정적인 수렴을 달성했다. 배치 크기는 131 K 토큰(패킹)이며, FlashAttention‑4와 torch.compile 기반의 CUDA 그래프 최적화로 B200 GPU에서 19 % MFU(실제/이론) 효율을 기록했다. 전체 FLOPs는 6.3 × 10¹⁹로, 기존 ESM‑2‑650M(3.9 × 10²¹) 대비 62배, E1‑600M(1.44 × 10²²) 대비 229배 적게 소모한다.
성능 평가에서는 ProteinGym의 217개 DMS(대체 변이) 어세이에서 Spearman ρ = 0.390을 달성해 ProGen2‑6.4B(ρ = 0.391)와 동등한 수준을 보였으며, 훈련 비용은 41–213배 적었다. 인델(삽입·삭제) 어세이에서는 ρ = 0.521로, 6 B 규모 모델보다도 우수한 결과를 얻었다. EVEREST 바이럴 피트니스 벤치마크에서도 구조 정보를 활용한 최신 방법에 근접한 성능을 보이며, 순수 시퀀스 기반 모델임에도 불구하고 강력한 일반화 능력을 입증했다.
해석적으로는 로그잇 렌즈를 적용해 레이어별 예측 분포를 시각화했으며, 특히 말단 레이어에서 최종 확률이 형성되는 과정을 확인했다. 또한 포지션별 엔트로피 표준편차가 검색 기반 보강(예: PoET, E1)의 효과와 음의 상관관계(ρ = ‑0.40)를 보인다는 사실을 발견했다. 즉, 엔트로피 변동이 작을수록 외부 동형체 검색이 유용하고, 변동이 클수록 모델 자체가 중요한 부위를 이미 파악하고 있어 검색이 오히려 해를 끼칠 수 있음을 시사한다. 이러한 통계적 힌트는 테스트‑타임에 검색 비용을 동적으로 조절하는 저비용 휴리스틱으로 활용 가능하다.