다중 스케일 단계적 단백질 구조 생성

초록

PAR는 단백질 백본을 다중 스케일로 다운샘플링하고, 각 스케일을 순차적으로 예측하는 자동회귀 트랜스포머와 흐름 기반 디코더를 결합한 프레임워크이다. 노이즈 컨텍스트 학습과 스케줄드 샘플링으로 노출 편향을 완화하고, 조건 없는 생성, 프롬프트 기반 설계, 모티프 스캐폴딩 등에서 뛰어난 제로샷 일반화와 샘플링 효율성을 보인다.

상세 분석

본 논문은 단백질 구조 생성 분야에서 기존의 확산·플로우 모델이 갖는 단일 스케일 처리와 이산화에 따른 세밀도 손실 문제를 근본적으로 해결하고자 한다. 핵심 아이디어는 단백질이 자연스럽게 계층적 구조(전체 토폴로지 → 2차 구조 → 원자 좌표)를 갖는다는 점을 활용해, 입력 백본을 여러 해상도로 다운샘플링하고, 각 해상도를 순차적으로 “다음 스케일 예측” 방식으로 복원하는 자동회귀(AR) 프레임워크를 설계한 것이다.

1. 다중 스케일 다운샘플링

   • 입력 백본 (x \in \mathbb{R}^{L \times 3})을 사전 정의된 크기 집합 (S={size_1,\dots,size_n})에 따라 연속적인 다운샘플링 연산 (Down(\cdot))을 적용한다.
   • 각 스케일 (x_i)는 원본 좌표를 시퀀스 차원에서 선형 보간해 얻은 3D 중심점이며, 이론적으로 쌍(pairwise) 거리와 방향 정보를 보존한다(논문 부록 C.8).
   • 스케일 구성을 “길이 기반”(예: 64,128,256) 혹은 “비율 기반”(예: L/4, L/2, L)으로 선택할 수 있어, 다양한 단백질 길이에 유연하게 적용 가능하다.

2. 자동회귀 트랜스포머

   • 비등가성(non‑equivariant) 트랜스포머 (T_\theta)는 이전까지 생성된 모든 코스 스케일을 입력으로 받아, 현재 스케일 (i)에 대한 조건 임베딩 (z_i)를 출력한다.
   • 입력은 (\text{bos}) 학습 가능한 시작 토큰과, 각 이전 스케일을 현재 스케일 크기로 업샘플링한 좌표 집합을 시퀀스 차원으로 연결(concatenate)한 형태다.
   • 이렇게 함으로써 “단방향” AR의 한계를 극복하고, 서로 멀리 떨어진 잔기라도 공간적으로 인접한 경우를 자연스럽게 모델링한다.

3. 플로우 기반 백본 디코더

   • 조건 임베딩 (z_i)를 적응형 레이어 정규화와 스케일 임베딩을 통해 플로우 네트워크 (v_\theta)에 주입한다.
   • 플로우는 표준 정규분포에서 목표 백본 좌표 분포로 매핑하는 연속적 변환이며, 시간 변수 (t)와 노이즈 (\epsilon)를 이용한 “플로우 매칭”(flow‑matching) 손실을 최소화한다.
   • 학습 시 (x_{i}^{t}=t\cdot x_i+(1-t)\cdot\epsilon) 형태로 노이즈를 섞어, 디코더가 다양한 노이즈 레벨에 강건하도록 만든다.

4. 노출 편향(Exposure Bias) 완화

   • 전통 AR 모델은 학습 시 실제 데이터(ground‑truth) 컨텍스트에 의존하지만, 추론 시 자체 예측값을 사용해 오류가 누적되는 “노출 편향”에 취약하다.
   • 논문은 두 가지 전략을 도입한다. 첫째, 노이즈 컨텍스트 학습으로 입력 컨텍스트에 의도적으로 잡음(또는 변형)을 주입해 모델이 불완전한 상황에서도 복원하도록 훈련한다.
   • 둘째, 스케줄드 샘플링(scheduled sampling)으로 학습 초기에 ground‑truth 비율을 높게 유지하고, 점진적으로 자체 예측을 섞어가며 전이 학습한다. 이로써 추론 시 컨텍스트와 실제 입력 사이의 분포 차이를 크게 감소시킨다.

5. 제로샷 일반화 및 응용

   • PAR는 별도 파인튜닝 없이도 인간이 제공하는 “프롬프트”(예: 특정 토폴로지, 길이, 혹은 부분 모티프)와 결합해 조건부 백본을 생성한다.
   • 특히 모티프 스캐폴딩 실험에서, 주어진 구조적 모티프를 고정하고 나머지 영역을 자동으로 채워넣어 설계 가능한 백본을 빠르게 얻을 수 있다.
   • 이러한 제로샷 능력은 대규모 언어 모델이 보여준 “작업 전이”와 유사하며, 향후 단백질 설계 파이프라인에 즉시 적용 가능성을 시사한다.

6. 성능 및 스케일링

   • 무조건 생성(조건 없는) 벤치마크에서 FPSD(Frechet Protein Structure Distance) 점수 161.0을 기록했으며, 이는 기존 확산 기반 모델 대비 경쟁력 있는 수치다.
   • 학습 compute를 늘릴수록 FPSD가 지속적으로 개선되는 “스케일링 법칙”을 관찰했으며, 이는 모델이 데이터 분포를 점진적으로 더 정확히 학습한다는 증거다.
   • 샘플링 단계에서 다중 스케일 접근법을 활용해, 단일 스케일 플로우 모델 대비 약 2.5배 빠른 속도를 달성했다.

7. 제한점 및 향후 과제

   • 현재는 Cα 원자만을 모델링하므로, 실제 설계에 필요한 side‑chain 배치나 전반적인 에너지 최소화 단계가 별도로 필요하다.
   • 다운샘플링 방식이 선형 보간에 의존하므로, 매우 긴 단백질(수천 잔기)에서는 스케일 선택이 성능에 큰 영향을 미칠 수 있다.
   • 트랜스포머가 비등가성이라 구조 회전·이동에 대한 불변성을 직접 학습해야 하는데, 이는 추가적인 정규화 혹은 equivariant 아키텍처와의 결합을 통해 개선될 여지가 있다.

핵심 기여는 (1) 다중 스케일 자동회귀 프레임워크 설계, (2) 플로우 기반 연속 좌표 디코더와의 결합, (3) 노출 편향 완화 기법 도입, (4) 제로샷 조건부 생성 및 스케일링 효율성 입증이다. 이로써 단백질 구조 생성 분야에서 자동회귀 모델의 새로운 가능성을 열었다.