PairSAE: Mechanistic Interpretability from Pair Representations in Protein Co-Folding

기술 요약: PairSAE – 단백질 공동 폴딩(Co-Folding)의 쌍 표현(Pair Representations)으로부터의 기계론적 해석 가능성

문제 정의
AlphaFold나 Boltz 모델과 같은 구조 생물학용 파운데이션 모델들은 생체 분자 구조 및 복합체 예측에서 놀라운 정확도를 달ы 성취했습니다. 그러나 이러한 출력값을 구동하는 내부 특징(internal features)을 이해하는 것은 여전히 중요한 과제로 남아 있습니다. 희소 오토인코더(Sparse Autoencoders, SAEs)는 트랜스포머 기반 언어 모델(LLM) 및 단백질 언어 모델(pLM)에서 '중첩된(superposed)' 특징들을 단일 의미적 개념(monosemantic concepts)으로 분리하는 데 성공했지만, 이를 구조 예측 모델로 깔끔하게 전이시키는 데는 어려움이 있습니다.

표준 LLM과 달리, 구조 예측 모델은 시퀀스 임베딩과 쌍별 상호작용 텐서(pairwise interaction tensors)를 모두 유지하는 페어포머(pairformer) 아키텍처를 사용합니다. SAE를 이러한 모델에 단순히 적용할 경우 두 가지 결정적인 문제가 발생합니다:

1. 이차적 폭발(Quadratic Blow-up): 쌍별 표현에 대한 별도의 딕셔너리를 학습하는 것은 시퀀스 길이에 따라 이차적으로 스케일링되어 계산적으로 매우 부담이 됩니다.
2. 가려진 개념(Obscured Concepts): 특징들이 시퀀스와 쌍 공간에 걸쳐 공동으로 중첩되어 존재할 수 있습니다. 이를 별도로 처리하면 두 표현 모두에 분포된 개념을 포착하지 못하며, 독립적으로 처리하면 학습된 특징들의 결합된 본질을 가리게 됩니다.

방법론: PairSAE
저자들은 시퀀스와 쌍 표현을 모두 공동으로 재구성하는 공유된 토큰 수준의 특징 세트를 학습하도록 설계된 프레임워크인 PairSAE를 소개합니다. 이 방법은 세 가지 주요 단계로 구성됩니다:

1. N-모드 SVD를 통한 쌍별 상호작용 요약:
   쌍별 텐서($Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$)의 이차적 스케일링 문제를 피하기 위해, 저자들은 N-모드 특이값 분해(N-mode SVD)를 수행합니다. 이들은 행과 열의 상호작용 역할을 인코딩하는 좌특이 벡터(left singular vectors)의 모드-1 및 모드-2 행렬($U^{(1)}$ 및 $U^{(2)}$)을 추출합니다. 이들은 상위 $r$개의 특이 벡터로 절단(truncated)된 후 토큰별 요약 임베딩 $m_i$로 결합됩니다. 이 단계는 이차적 복잡도 없이 상호작용 역할을 보존하면서 쌍별 정보를 토큰 수준의 표현으로 압축합니다.

2. 공유 오토인코더 아키텍처:
   PairSAE 인코더는 원래의 시퀀스 임베딩($s_i$)과 SVD로부터 유도된 상호작용 요약($m_i$)의 결합을 입력으로 받습니다. 이는 희소성을 유도하는 비선형성(BatchTopK와 ReLU의 조합)을 사용하여 희소 잠재 특징 벡터 $h_i$를 생성합니다.
   결정적으로, 디코더는 이 공 공유된 특징 세트 $h$를 사용하여 다음 두 가지를 모두 재구성합니다:

   • 시퀀스 임베딩: $\hat{s}_i = D_s h_i + b_{dec}^s$
   • 쌍 임베딩: $\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
     이는 모델이 국소적 시퀀스 문맥과 쌍별 상호작용을 모두 설명할 수 있는 통합된 특징 기저(unified feature basis)를 학습하도록 강제합니다.

3. 손실 함수 및 훈련:
   모델은 중첩된 너비($c_1 < c_2 < c_3 = D$)에 걸쳐 디코더와 특징 벡터를 슬라이싱하여 견고함을 보장하는 Matryoshka SAE 손실을 사용하여 훈련됩니다. 전체 손실은 동일한 가중치로 계산되는 시퀀스 및 쌍 표현에 대한 재구성 항과 "죽은(dead)" 특징을 방지하기 위한 보조 손한을 포함합니다. 훈련은 각 $i$에 대해 단일 $j$를 샘생하는 몬테카를로 근사법을 사용하여 쌍별 손실 항을 계산함으로써 가속화됩니다.

주요 기여

• PairSAE 아키텍처: N-mode SVD를 통해 쌍별 텐서를 요약하고 공유된 특징 세트로부터 시퀀스 및 쌍 공간으로 디코딩함으로써, 페어포머 기반 모델에 희소 딕셔너리 학습을 적용하는 새로운 방법을 제시합니다.
• 해석 가능한 특징 발견: 이 방법은 구조 생물학 모델에 기존 SAE를 적용할 때 발생하는 문제점을 피하면서, 생물물리학적 개념과 일치하는 특징들을 성공적으로 복구합니다.
• 결합 친화도(Binding Affinity) 적용: 이 프레임워크는 단백질-리간드 복합체(PLINDER 데이터셋)의 Boltz-2 활성화 값에 적용되어, 잠재 표현을 결합 친화도 예측과 연결합니다.

결과
저자들은 단백질-리간드 복합체의 Boltz-2 표현에 대해 PairSAE를 평가했습니다:

• 선형 프로빙(Linear Probing): PairSAE 특징은 UniProt 잔기 주석(residue annotations)에 대해 높은 예측력을 보여주었습니다. 레이어 33과 64(세 번째 리사이클링 단계)의 특징은 토큰 수준에서 29.1%와 24.0%, 복합체 수준에서 53.2%와 49.6%의 F1 점수를 기록하며, ESM2-650M 뉴런 기반 분류기(각각 0.8%와 19.7% 기록)를 크게 앞질렀습니다.
• 개념 정렬(Concept Alignment): 시각화를 통해 막 관통 영역(transmembrane regions), 이황화 결합(disulfide bonds), 프로테아제 활성 부위(protease active sites)와 같은 특정 구조적 모티프에 활성화되는 특정 특징들을 확인하였으며, UniProt 실제 값에 대해 높은 F1 점수를 나타냈습니다.
• 친화도 예측: Max-pooling된 PairSAE 특징에 LASSO 회귀를 사용하여, 저자들은 테스트 세트(PoseBusters)에서 $R^2 = 0.528$의 성능으로 Boltz-2 친화도 값을 예측했습니다. 고친화도 복합체와 강하게 상관관계가 있으며 리간드에 특이적으로 활성화되는 특정 특징(예: 특징 2299)이 식별되었으며, 비활성 복합체와 유의미한 그룹 차이를 보였습니다.

의의 및 주장
본 논문은 PairSAE가 구조 생물학을 위한 파운데이션 모델의 잠재 공간을 해석 가능한 구조적 개념과 성공적으로 연결한다고 주장합니다. 페어포머 모델의 특정 아키텍처 제약 조건을 해결함으로써, PairSAE는 기존 SAE가 가진 이차적 복잡도나 특징 중첩 문제를 겪지 않고 폴딩 및 결합에 관한 "모델이 무엇을 알고 있는지"를 명확히 합니다.

저자들은 이 연구를 구조 생물학의 원칙적인 평가 및 발견을 향한 단계로 규정하며, 연구자들이 예측 신뢰성과 생물물리학적 타당성을 평가할 수 있게 해준다고 말합니다. 저자들은 이 방법이 시퀀스 및 쌍 특징을 알려진 주석에 효과적으로 매핑하지만, 상세한 리간드 주석의 부재로 인해 리간드 활성화 특징을 특정 개념으로 매핑하는 데에는 여전히 한계가 있다고 언급했습니다. 향-후 연구로는 더 많은 레이어로 연구를 확장하고, 이러한 특징들을 자동화된 파이프라인에 통합하며, 해석 가능한 단백질 설계를 위한 스티어링(steering) 방법을 개발하는 것을 제안합니다.