기술 요약: Manifold Restore Mixing을 통한 단백질 표현 학습 강화

1. 문제 정의

단백질 표현 학습(Protein Representation Learning, PRL)은 실험적인 3D 구조 데이터의 부족을 극복하기 위해 데이터 증강(Data Augmentation, DA)에 크게 의존합니다. 그러나 잔기 치환(residue substitution), 좌표 노이즈 주입(coordinate noise injection), 부분 구조 샘플링(substructure sampling)과 같은 주류 DA 방법들은 단백질의 1D 아미노산 서열과 3D 입체 구조 사이의 결정적인 대응 관계를 깨뜨리는 경우가 많습니다.

저자들은 근본적인 딜레마를 식별했습니다. 즉, 섭동 기반(perturbation-based) 방법들은 다양한 변형을 생성하지만, 핵심 기능 도메인의 구조적 무결성을 파괴하여 증강된 데이터를 생물학적으로 불가능하거나 기능적으로 일관되지 않게 만들 위험이 있습니다. 반대로, 상동성 모델링(homology modeling)이나 확산 기반(diffusion-based) 생성 도구들은 충분한 구조적 다양성이 부족한 경우가 많습니다. 본 논문의 실증적 분석에 따르면, 이러한 표준 증강 기술을 다운스트림 태스크에 적용할 경우 모델이 구조적으로 손상된 샘플로부터 정확한 표현을 학습하는 데 어려움을 겪게 되어, 성능 향상이 아닌 상당한 성능 저하를 초-래한다는 것을 보여줍니다.

2. 방법론: Manifold Restore Mixing (MRM)

기존 DA의 구조적 결함을 해결하기 위해, 저자들은 원시 입력 공간(raw input space)이 아닌 잠재 표현 공간(latent representation space)에서 작동하는 방법인 **Manifold Restore Mixing (MRM)**을 제안합니다. 이 방법론은 세 가지 핵심 구성 요소로 이루어집니다.

A. Manifold Restore Mixing

Manifold Mixup에서 영감을 받은 MRM은 신경망의 특정 중간 레이어 $k$에서 원래 단백질 데이터($GP$)와 증강된 단백질 데이터($GP_a$)의 은닉 표현을 융합합니다.

• 과정: 네트워크는 $GP$와$GP_a$를 레이어$k$까지 처리하여 표현$R_k(GP)$와$R_k(GP_a)$를 생성합니다.
• 혼합(Mixing): 선형 보간을 통해 혼합된 표현 $\tilde{R}_k$가 생성됩니다:
  $$\tilde{R}_k = \lambda \cdot R_k(GP_a) + (1 - \lambda) \cdot R_k(GP)$$
  여기서 $\lambda \sim \text{Beta}(\alpha_1, \alpha_2)$입니다.
• 복원 메커니즘(Restoration Mechanism): 표현 수준에서 혼합함으로써, 이 방법은 증강 과정에서 도입된 다양한 변형은 유지하면서도, $GP$의 온전한 구조적 특징을 재도입함으로써$GP_a$의 증강 중에 손실된 구조 정보를 "복원"하는 것을 목표로 합니다. 이후 혼합된 표현은 손실을 계산하기 위해 남은 레이어들을 통과합니다.

B. 샘플 난이도 스케줄러 (Sample Difficulty Scheduler)

표준 Mixup은 고정된 Beta 분포(일반적으로 $\alpha_1 = \alpha_2 < 1$)를 사용하여, 샘플이 원래 데이터나 증강된 데이터 중 어느 한쪽에 의해 완전히 지배되는 "독점 효과(monopoly effect)"를 초래할 수 있습니다. MRM은 커리큘럼 학습 접근 방식을 도입합니다:

• 초기 상태: 훈련은 $\alpha_1 > \alpha_2 > 1$로 시작하여, 혼합이 원래 데이터(쉬운 샘플) 쪽으로 편향되도록 합니다.
• 진행: 훈련 에포크(epoch)가 경과함에 따라, $\alpha_1$은 감소하고 $\alpha_2$는 증가하도록 매개변수를 조정하여, 더 많은 증강 데이터(어려운 샘플)를 포함하도록 분포를 이동시킵니다.
• 목표: 이를 통해 모델에게 쉬운 샘플(구조 보존형)에서 어려운 샘플(변형 풍부형)로 이어지는 점진적인 학습 곡선을 제공합니다.

C. 2단계 정규화 훈련 (Two-Stage Regularized Training)

모델 파라미터가 무작위로 초기화되었을 때 표현을 혼합함으로써 발생하는 수렴 불안정성을 방지하기 위해, MRM은 2단계 전략을 채택합니다.

1. 1단계: 모델을 오직 원래 데이터로만 훈련하여 안정적인 베이스라인 표현을 구축합니다.
2. 2단계: MRM 모듈을 활성화합니다. 전체 손실은 원래 데이터에 대한 경험적 위험 최소화(Empirical Risk Minimization, ERM)와 혼합된 샘플에 대한 근방 위험 최소화(Vicinal Risk Minimization, VRM)를 결합합니다:
   $$L_{second} = L_{CLS}(Net(GP), y) + \gamma L_{CLS}(\tilde{R}_f, y)$$
   이는 모델이 분포 내 적합도(in-distribution fitting accuracy)를 유지하면서 복원된 표현의 근방 공간(vicinal space)을 탐색할 수 있도록 보장합니다.

3. 주요 기여

1. DA 결함에 대한 실증적 분석: 본 논문은 표준 단백질 DA 연산이 1D-3D 구조적 대응 관계를 방해하여 다운스트림 태스크에서 성능 저하를 초래한다는 체계적인 분석과 실증적 근거를 제공합니다.
2. MRM 프레임워크: 잠재 공간에서 원래 표현과 증강된 표현을 융합함으로써 증강 중에 손실된 구조 정보를 복원하는 Manifold Restore Mixing을 제 제안합니다.
3. 훈련 강화: 훈련을 안정화하고 구조적 충실도와 변형 사이의 균형을 최적화하기 위해 샘플 난이도 스케줄러와 2단계 정규화 훈련 전략을 개발했습니다.
4. 포괄적인 검증: 다양한 PRL 백본(예: GearNet, CDConv)과 다양한 다운스트림 태스크(Fold Classification, Enzyme Reaction, Gene Ontology, EC Number Prediction)에 걸친 광범ening한 실험을 통해 기존 DA 방법 및 mixup 변형들보다 우수한 성능을 입증했습니다.

4. 실험 결과

저자들은 네 가지 주요 태스크(단백질 폴드 분류, 효소 반응 분류, 유전자 온톨로지(GO) 용어 예측, 효소 번호(EC Number) 예측)에 대해 MRM을 평가했습니다.

• 성능 향상: MRM은 베이스라인 모델 및 다른 증강 전략보다 일관되적으로 높은 성능을 보였습니다. 예를 들어, Fold Classification 태스크에서 MRM은 GearNet-Edge의 정확도를 3.1%, CDConv의 정확도를 4.1% 향상시켰습니다. Gene Ontology 예측에서는 전통적인 증강 방법이 실패하거나 성능을 저하시킨 모든 하위 태스크(BP, MF, CC)에서 개선을 달성했습니다.
• Mixup 변형과의 비교: 단백한 단백질 데이터에 표준 Manifold Mixup이나 NFM을 직접 적용하면 성능 저하가 발생했습니다. 이러한 방법들의 "Two-Stage"(TS) 변형들은 약간의 개선을 보였으나 여전히 MRM보다 뒤처졌으며, 이는 MRM의 특정한 "복원(restore)" 메커니즘이 매우 중요하다는 것을 확인시켜 줍니다.
• 단백질 LLM과의 비교: 대규모 사전 학습을 사용하지 않았음에도 불구하고, MRM이 강화된 모델은 Fold 및 Enzyme Reaction 태스크에서 사전 학습된 단백질 대규모 언어 모델(예: ESM-2, ProtBERT)과 대등하거나 경쟁력 있는 결과를 얻었으며, 평가된 모든 방법 중 가장 높은 평균 순위를 기록했습니다.
• 일반화 능력: 이 방법은 서로 다른 서열 유사도 컷오프(30% ~ 95%) 전반에서 강건함을 입증하였으며, 이는 훈련 세트와 상동성이 낮은 단백질에 대해서도 강력한 일반화 능력을 갖추고 있음을 나타냅니다.
• 절제 연구(Ablation Studies): 어떤 구성 요소(1D/3D 증강, 2단계 훈련 또는 난이도 스케줄러)라도 제거하면 성능 저하가 발생하여, 전체 프레임워크의 필요성을 입증했습니다.

5. 의의 및 주장

본 논문은 MRM이 단백질 데이터의 구조적 취약성을 해결함으로써 계산 생물학 분야의 데이터 증강을 위한 새로운 패러다임을 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 구조적 무결성 복원: 정확한 단백질 기능 예측에 필요한 필수적인 구조 정보를 희생하지 않으면서도 다양한 증강 데이터를 사용할 수 있게 합니다.
• 격차 해소: (PRL에 너무 파괴적일 수 있는) 전통적인 증강 연산들이 다운스트림 태스크에서 효과적으로 활용될 수 있도록 합니다.
• 사전 학습에 대한 잠재력: 저자들은 "복원 혼합(restore mixing)" 패러다임이 단백질 LLM의 사전 학습 시 자가 지도 신호(self-supervised signals)를 구축하는 데 통찰력을 제공하여, 증강된 뷰(augmented views)로부터 학습하는 능력을 향상시킬 수 있다고 제안합니다.

저자들은 "복원"이 잠재 공간에서 발생하며 좌표 공간에서의 완전한 물리적 실현 가능한 3D 입체 구조를 재구성하는 것은 아니라는 점을 언급하며, 생물학적 해석에 대해 신중한 태도를 유지했습니다. 또한 향-후 연구에서 증강 연산을 동적으로 조정하기 위해 베이지안 최적화(Bayesian optimization)나 강화 학습(reinforcement learning)을 통합할 수 있음을 인정했습니다.