VFUSE: Virulent Feature Understanding with Sparse autoEncoders

기술 요약: VFUSE – 희소 오토인코더를 이용한 악성 특징 이해(Virulent Feature Understanding with Sparse autoEncoders)

문제 정의

단백질 설계를 위한 생성 모델, 특히 RoseTTAFold3 (RF3) 및 RFDiffusion3 (RFD3)와 같은 전원자(all-atom) 확산 모델은 위험한 템플릿 주변에 새로운 구조를 스캐폴딩함으로써 유해한 단백질(예: 신경독소, 면역 회피 단백질)을 설계할 수 있는 이중 용도(dual-use) 능력을 보유하고 있습니다. 현재의 생물 보안 스크리닝 방법은 서열 상동성이나 서열 수준의 분류기에 의존합니다. 이러한 접근 방식은 두 가지 결정적인 한계점을 가집니다. 첫째, 사후적(post-hoc)으로 작동하여(생성 후 서열을 플래그함) 위험 신호를 감지하며, 둘째, 기계적 해석 가능성이 부족하여 어떤 특정 구조적 특징이 위험 플래그를 트리거했는지 설명하거나 샘플링 과정 중에 개입할 수 없습니다. 또한, 희소 오토인코더(Sparse Autoencoders, SAE)가 단백질 언어 모델(PLM) 및 시각 모델에는 성공적으로 적용되었으나, 전원자 확산 모델의 내부 표현을 조사하는 데에는 아직 활용되지 않았습니다.

방법론

저자들은 확산 트랜스포머 모델의 활성화 함수에 대해 희소 오토인코더를 학습시켜 위험 관련 특징을 감사하는 기계적 해석 가능성 프레임워크인 VFUSE를 소개합니다.

1. 모델 및 활성화 수집:

   • 본 연구는 RFDiffusion3(18개 블록의 전원자 확산 모델)와 RoseTTAFold3(그에 대응하는 폴딩 모델)를 대상으로 합니다.
   • 활성화($d=768$)는 특정 잔차 스트림(residual stream) 블록인 RFD3의 경우 블록 6, 8, 12에서, RF3의 경우 블록 12, 16에서 PyTorch 훅(hook)을 통해 캡처됩니다.
   • RFD3의 경우, 위험한 템플릿의 PDB 구조에 부분 확산(partial diffusion, $\partial t = 5$, 약 5Å 노이즈)을 적용하여 모델의 설계 과정을 시뮬레이션함으로써 입력을 생성합니다. RF3 입력은 직접 폴딩됩니다.

2. 데이터셋 구축:

   • 악성 서열: SafeProtein(275쌍)과 ToxinPred3(RF3 전용 1,200개 서열)에서 추출되었습니다.
   • 양성 서열: UniProt에서 유래하였으며, 독소/독성 키워드를 제외하도록 필터링된 길이 매칭(100–300 aa) 서열입니다.
   • 상동성 제어: 프로브가 단순한 폴드 패밀리의 정체성을 기억하는 것이 아니라 독성 신호를 학습하도록 하기 위해, mmseqs2를 사용하여 30% identity 기준으로 서열을 클러스터링했습니다. 훈련/테스트 분할은 근친 상동체가 양측에 모두 나타나지 않도록 보장합니다.

3. SAE 아키텍처 및 학습:

   • 아키텍처: Matryoshka BatchTopK SAE가 채택되었습니다. 이는 다음을 결합합니다:
     • Matryoshka: 다양한 정밀도로 재구성을 허용하는 중첩된 접두사 크기(nested prefix sizes).
     • BatchTopK: L1 페널티 대신 배치 전체에서 상위 $K$개의 사전 활성화(pre-activations)만을 유지하여 희소성을 강제합니다.
   • 하이퍼파라미터: 사전 크기 $m=12,288$ (16배 확장), 활성 잠재 변수 $K=80$, 20,000 스텝 동안 학습되었습니다.
   • 성능: 학습된 SAE는 홀드아웃 활성화 분산의 96.9%를 설명하며, 평균 $L_0$ 희소성은 79.9입니다.

4. 평가 프로토콜:

   • 프로빙(Probing): 평균 풀링된 SAE 잠재 벡터($z$)와 원시 활성화($x$)를 사용하여 L2 정규화 로지스틱 회귀 프로브를 학습시킵니다.
   • 지표: 무작위 계층적 분할(random stratified splits)과 상동성 클러스터링 분할(homology-clustered splits) 두 가지 체제에서 AUROC를 평가합니다.
   • 특징 점수화: 각 클래스 레이블에 대해 단변량 AUROC를 계산합니다. 유의성은 Benjamini–Hochberg FDR 보정($q < 0.05$)을 거친 Mann–Whitney U 검정을 통해 결정됩니다.

주요 결과

1. SAE 잠재 변수가 원시 활성화보다 우수함:

   • RFD3 블록 12에서, SAE 인코딩된 특징은 상동성 클러스터링 설정에서 원시 특징(AUROC 0.763) 대비 유의미하게 높은 0.817의 AUROC를 달성하며 원시 활성화를 능가했습니다 (+0.054 개선).
   • 이러한 개선은 블록 12에 국한됩니다. 초기 블록(6, 8)은 SAE로부터 중립적이거나 부정적인 이득을 보였습니다.
   • RF3(ToxinPred3 데이터셋)의 경우, 블록 12 SAE 특징 또한 클러스터링된 분할에서 원시 특징 대비 +0.041의 개선을 보였습니다.

2. 상동성 기억 현상의 감소:

   • 무작위 분할과 상동성 클러스터링 분할 간의 AUROC 차이(패밀리 기억의 대리 지표)는 RFD3 블록 6(0.22 차이)에 비해 RFD3 블록 12 SAE 특징(0.13 차이)에서 가장 작았습니다. 이는 블록 12 SAE가 폴드 패밀리를 단순히 암기하는 것이 아니라, 클래스 판별적이고 패밀리 일반적인 구조적 신호를 포착함을 시사합니다.

3. 단일 의미적(Monosemantic) 위험 특징 식별:

   • 연구는 악성 설계에 선택적으로 반응하는 개별 SAE 특징을 식별했습니다.
   • 최고 성능 특징: RFD3 블록 12의 특징 639는 viper 신경독소를 감지하는 데 있어 0.819의 AUROC($q \approx 10^{-13}$)를 달성했습니다.
   • 국소화(Localization): 이러한 특징들은 잔기(residue) 수준의 국소성을 보입니다. 특징 639의 경우, 상위 3개 활성화 잔기(99, 104, 109)는 Ammodytoxin A 구조 내의 단일 알파 헬릭스 내에 클러스터링되어 있습니다. 다른 특징들도 원숭이두창(Monkeypox) 면역 회피 단백질, 리보솜 불활성화 단백질, 모기 타액 독소의 특정 영역으로 유사하게 국소화됩니다.

4. 경쟁력 있는 성능:

   • 최고 성능의 VFUSE 프로브(RFD3 블록 12 SAE, 무작위 분할, AUROC 0.877)는 단순한 로지스틱 프로브와 더 작은 데이터셋을 사용했음에도 불구하고 최신 서열 전용 분류기(예: DTVF, 0.92)와 경쟁할 만한 수준입니다. 결정적으로, VFUSE는 서열 전용 모델이 갖지 못한 잔기별 근거(rationale)를 제공합니다.

의의 및 주장

본 논문은 전원자 확산 모델에 대해 학습된 최초의 SAE이자, 단백질 설계 모델에 대한 최초의 특징 수준 독성 감사라고 주장합니다.

• 기계적 통찰: 이 작업은 특정 깊이(블록 12)에서의 딕셔너리 학습이 다의적인(polysemantic) 구조적 개념을 저용량 위험 탐지에 더 적합한 단일 의미적(monosemantic) 방향으로 풀어낼 수 있음을 입증합니다.
• 안전성 및 해석 가능성: VFUSE는 생성형 단백질 모델의 런타임 모니터링 및 해석 가능한 감사를 위한 경로를 구축합니다. 독성과 관련된 특정 잔기와 구조적 모티프를 식별함으로써, 이 방법은 "블랙박스" 스크리닝을 넘어 설명 가능한 안전 점검으로 나아갑니다.
• 향후 방향: 저자들은 한계점(작은 데이터셋, 예비 테스트에서의 스티어링 성공 부족)을 언급했으나, 본 연구는 서열 수준을 넘어 특징 수준에서 유해한 특징을 탐지할 수 있게 함으로써 안전하고 해석 가능한 단백질 설계를 위한 길을 열었습니다.