[논문 리뷰] Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval

Pascal Notin, Mafalda Dias외 5명|arXiv (Cornell University)|2022. 05. 27.

Genomics and Phylogenetic Studies인용 수 123

한 줄 요약

Tranception은 homologous 시퀀스의 추론 시 조회를 통해 단백질 적합도 예측을 향상시키는 autoregressive 트랜스포머를 도입하며, 특히 얕은 정렬 및 인델에서 성능을 개선합니다.

ABSTRACT

The ability to accurately model the fitness landscape of protein sequences is critical to a wide range of applications, from quantifying the effects of human variants on disease likelihood, to predicting immune-escape mutations in viruses and designing novel biotherapeutic proteins. Deep generative models of protein sequences trained on multiple sequence alignments have been the most successful approaches so far to address these tasks. The performance of these methods is however contingent on the availability of sufficiently deep and diverse alignments for reliable training. Their potential scope is thus limited by the fact many protein families are hard, if not impossible, to align. Large language models trained on massive quantities of non-aligned protein sequences from diverse families address these problems and show potential to eventually bridge the performance gap. We introduce Tranception, a novel transformer architecture leveraging autoregressive predictions and retrieval of homologous sequences at inference to achieve state-of-the-art fitness prediction performance. Given its markedly higher performance on multiple mutants, robustness to shallow alignments and ability to score indels, our approach offers significant gain of scope over existing approaches. To enable more rigorous model testing across a broader range of protein families, we develop ProteinGym -- an extensive set of multiplexed assays of variant effects, substantially increasing both the number and diversity of assays compared to existing benchmarks.

연구 동기 및 목표

• 다양한 단백질 가족 전반에 걸쳐 단백질 적합도 지형을 정확하게 모델링하는 것을 동기 부여하며, 정렬이 어렵거나 무질서한 영역을 포함하는 단백질도 대상으로 한다.
• 추론 시 동족 정보를 도입하기 위해 검색을 활용할 수 있는 비-MSA-훈련 단백질 언어 모델을 개발한다.
• 종 간에 강건한 성능으로 치환, 삽입 및 제거에 대한 예측을 향상시킨다.
• 다양한 분석법 across assays에서 돌연변이 효과 예측기를 엄격하게 평가하기 위한 대규모의 다양하고 포괄적인 벤치마크(ProteinGym)를 제공한다.

제안 방법

• Tranception을 제안한다. 다중 k-mer 패턴을 포착하기 위해 그룹화 커널 컨볼루션을 사용하는 특수한 주의 메커니즘(Tranception 주의)을 갖춘 autoregressive 트랜스포머이다.
• 표준 위치 인코딩을 Grouped ALiBi로 교체하여 헤드별 거리 인지 주의(attention)와 더 긴 컨텍스트 모델링을 가능하게 한다.
• 정렬되지 않은 UniRef 시퀀스에서 학습한다(700M 매개변수 모델; 컨텍스트 크기 1024) 및 양방향 점수화를 개선하기 위해 시퀀스 미러링을 적용한다.
• 돌연변이 시퀀스와 와일드타입 시퀀스 사이의 대로그가능도 비율을 계산하여 적합도를 점수화한다(식 2).
• 추론 시, 추론 시점에 검색된 MSA로부터의 검색 기반 예측(P_R)과 autoregressive 예측(P_A)을 결합한다(식 3, 4).
• 검색을 활용하여 검색된 MSA에서 각 위치의 아미노산 분포를 의사계(counts)와 라플라스 평활화를 사용해 얻고, 샘플링 바이어스를 보정하도록 재가중한다(Hopf et al. 2017).
• 삽입/삭제를 처리하기 위해 검색된 MSA 열을 맞춤화하고 새로운 위치에서는 autoregressive 모드를 의존하며, 안정성을 위해 좌→우 점수와 우→좌 점수를 평균한다.

실험 결과

연구 질문

• RQ1 정렬되지 않은 시퀀스에서 학습된 autoregressive 트랜스포머가 학습 중에 MSA에 의존하지 않고도 최신의 단백질 적합도 예측에서 최첨단 성능을 달성할 수 있는가?
• RQ2 추론 시점의 동족 시퀀스 검색이 예측을 개선하는가, 특히 얕은 또는 MSAs가 거의 없는 단백질 및 인델에 대해?
• RQ3 Tranception은 치환, 다중 돌연변이, 인델 및 다양한 분류군에서 정렬 기반 및 다른 단백질 언어 모델과 비교하여 어떤 성능을 보이는가?
• RQ4 모델은 MSA 깊이에 강건하며 정렬하기 어렵거나 무질서한 영역을 점수화할 수 있는가?
• RQ5 어떤 벤치마크가 다양한 분석법과 분류군 전반에 걸친 돌연변이 효과를 포괄적으로 평가할 수 있는가(ProteinGym)?

주요 결과

• Tranception (with retrieval) outperforms all baselines on ProteinGym substitution and indel benchmarks.
• Retrieval substantially improves performance, with the largest gains for proteins with shallow MSAs and for multiple mutations.
• Without retrieval, Tranception already exceeds non-retrieval baselines and competitive alignment-free models; with retrieval, it surpasses alignment-based methods.
• The model shows robustness to MSA depth and can score regions that are poorly aligned or disordered, providing broad coverage across proteins (e.g., BRCA1, viral proteins).
• Tranception excels in extrapolating far in sequence space, yielding larger gains for multiple mutants than for single mutants.
• ProteinGym provides a diverse, expansive benchmark (including indels) that reveals clear advantages of Tranception over prior methods.