AI바라기의 인공지능

Protein : 논문 리뷰 : Protein Autoregressive Modeling via Multiscale Structure Generation 본문

논문리뷰

Protein : 논문 리뷰 : Protein Autoregressive Modeling via Multiscale Structure Generation

AI바라기 2026. 3. 23. 16:51


용어 설명

  • Exposure bias: Autoregressive 모델의 고질적 문제. 학습 시에는 완벽한 실제 데이터(ground truth)를 입력받아 다음을 예측하지만, 실제 추론(inference) 시에는 모델 자신이 예측한 이전 결과물(노이즈가 섞인 값)을 바탕으로 다음을 예측하므로 단계가 지날수록 오류가 누적되어 품질이 무너지는 현상입니다.
  • Next-scale prediction: 일반적인 텍스트 모델의 단방향 순차 예측(Next-token)과 달리, 전체 3D 구조의 뼈대(거친 형태, Coarse)를 먼저 예측하고 점차 세밀한 디테일(Fine)로 해상도를 높여가며 계층적으로 예측하는 방식입니다.
  • Noisy Context Learning (NCL) & Scheduled Sampling (SS): Exposure bias를 완화하기 위한 학습 기법. 입력 조건에 의도적으로 노이즈를 섞거나(NCL), 학습 중 일정 확률로 모델 자신의 예측값을 다음 스텝의 입력으로 대체(SS)하여 모델이 불안정한 조건 속에서도 올바른 구조를 생성하도록 면역력을 키웁니다.
  • Flow-based decoder (Flow matching): 단순한 노이즈 분포를 실제 데이터 분포로 변환하는 연속적인 흐름(흐름 방정식)을 학습하는 생성 모델 기법입니다. 이 논문에서는 이산화된 토큰 대신 연속적인 3D 좌표를 직접 생성하는 역할을 합니다.
  • FPSD (Frechet Protein Structure Distance): 생성된 단백질 구조와 실제 자연계 단백질 구조 간의 분포적 차이(품질 및 다양성)를 한 번에 측정하는 지표입니다. 이미지 생성 모델을 평가하는 FID와 유사한 역할을 합니다.
  • Motif scaffolding: 단백질의 핵심 기능을 담당하는 특정 부위(Motif)의 3D 좌표를 그대로 유지한 채, 이 부위를 안정적으로 지탱할 수 있는 나머지 뼈대 구조(Scaffold)를 새롭게 설계하는 작업입니다.

Purpose of the Paper

  • 기존 연구의 한계: 단백질 3D 구조 생성은 주로 Diffusion 기반 모델들이 주도하고 있었습니다. 반면 LLM에서 파괴적인 성능을 입증한 Autoregressive (AR) 모델을 단백질 구조에 적용하려는 시도는 실패를 거듭했습니다. 3D 좌표를 이산화(Tokenize)하면서 미세한 원자 디테일이 훼손되었고, 특히 AR의 '왼쪽에서 오른쪽으로' 예측하는 단방향(unidirectional) 가정이 서열상 멀리 떨어져 있어도 3D 공간에서는 서로 강하게 결합하는 단백질 특유의 양방향 구조적 의존성을 완전히 무시했기 때문입니다.
  • 새로운 접근 방식: 이 논문은 AR 모델의 장점(확장성, Zero-shot 등)을 살리면서 단방향의 한계를 깨부수는 최초의 다중 스케일 Autoregressive 프레임워크인 PAR (Protein AutoRegressive) 를 제안합니다. 토큰화 없이 연속적인 3D 공간에서 직접 구조를 모델링하며, 순차적 예측이 아닌 거친 구조에서 세밀한 구조로 넘어가는 Next-scale prediction을 통해 단백질 내의 복잡한 물리적 상호작용을 온전히 보존하고자 했습니다.

Key Contributions

  • Multi-scale Next-scale AR Architecture: (Novelty) 기존의 단방향 시퀀스 모델링에서 벗어나, 데이터를 여러 해상도로 다운샘플링하여 Coarse-to-fine 계층 구조를 만들었습니다. 하나의 스케일 내에서는 전체 서열을 동시에 다루므로 단백질의 양방향성을 완벽히 포착합니다.
  • Continuous Direct Modeling (이산화 제거): (Novelty) VQ-VAE와 같은 이산화 도구를 버렸습니다. AR Transformer가 이전 스케일의 정보를 바탕으로 스케일별 조건(Conditioning)을 만들면, 이를 Flow-based backbone decoder가 넘겨받아 연속적인 공간에서 직접 단백질의 Ca(알파 탄소) 3D 좌표를 생성합니다. 이로 인해 구조적 충실도(fidelity)가 비약적으로 상승했습니다.
  • Exposure Bias의 성공적 완화: (Contribution) 지상 관측치(ground truth)를 기반으로 학습하는 AR 프레임워크가 3D 좌표 생성에서 겪는 치명적인 구조 붕괴 현상(Exposure bias)을 식별하고, Noisy Context Learning과 Scheduled Sampling을 결합하여 이 문제를 효과적으로 억제함으로써 강건한 구조 생성을 이뤄냈습니다.

Experimental Highlights

  • SOTA 달성 및 Scaling Behavior 입증: Unconditional backbone generation 작업에서 PAR는 PDB 데이터셋 기준 FPSD 161.0, Designability 96.6%를 달성하여 기존 순수 Diffusion 베이스라인(Proteina 등)을 능가하는 최고의 데이터 분포 학습 능력을 보여주었습니다. 또한 모델 크기(60M -> 400M)와 연산량을 늘릴수록 생성 품질이 뚜렷하게 향상되는 AR 특유의 우수한 Scaling behavior를 확인했습니다.
  • 강력한 Zero-Shot Task Generalization: 별도의 Fine-tuning 전혀 없이, 사용자가 제시한 단 16개의 점(Point)만으로 전체 단백질의 레이아웃을 통제하거나, 기존 Motif 좌표를 완벽히 유지한 채 전혀 새로운 뼈대(Scaffold)를 생성하는 복잡한 조건부 생성 작업을 Zero-shot으로 성공해냈습니다.
  • SDE/ODE 오케스트레이션을 통한 추론 속도 혁신: 다중 스케일 구조의 이점을 활용하여, 뼈대를 잡는 첫 번째 Coarse 스케일에서는 자유도가 높은 SDE 샘플링을 적용하고, 디테일을 채우는 나머지 Fine 스케일에서는 계산이 빠른 ODE 샘플링을 적용했습니다. 결과적으로 생성 품질(Designability 97%)을 유지하면서도 단일 스케일 모델 대비 2.5배 빠른 획기적인 추론 속도를 달성했습니다.

Limitations and Future Work

  • 한계점 (Limitations):
    • 현재 모델은 뼈대(Backbone Ca) 원자 생성에만 집중되어 있어, 실제 단백질 기능 발현에 중요한 측쇄(Side chain)를 포함한 All-atom 모델링은 불가능합니다.
    • Flow decoder와 달리, 조건을 생성하는 AR Transformer 모듈의 파라미터를 크게 키우면 오히려 성능이 정체되는 현상이 발견되었습니다. 이는 한정된 데이터(588K) 상황에서 모델이 완벽한 정답 컨텍스트에 지나치게 Overfitting 되어 앞서 언급한 Exposure bias가 다시 심화되기 때문입니다.
  • 향후 연구 방향 (Future Work):
    • 현재의 다중 스케일 프레임워크가 유연한 Zero-shot 프롬프트 설계를 지원하므로, 이를 확장하여 구조뿐만 아니라 Side chain까지 완벽히 생성하는 All-atom AR 모델링으로 발전시킬 계획입니다.
    • 단백질 구조를 거칠게 Downsample한 뒤 PAR로 다시 Upsample하는 본연의 동작 방식을 응용하여, 단백질이 살아 움직이는 동적 변화(Conformational dynamics)를 분자 동역학 시뮬레이션 없이 Zero-shot으로 모델링하는 새로운 가능성을 탐구할 예정입니다.

Overall Summary
이 논문은 기존 언어 모델 방식의 단일 방향 순차 예측과 토큰 이산화로 인해 3D 단백질 모델링에서 외면받았던 Autoregressive 모델의 치명적 한계를 Next-scale prediction 방식과 연속적 공간 직접 생성(Flow matching)을 통해 완벽히 극복한 최초의 다중 스케일 PAR (Protein AutoRegressive) 프레임워크를 제안했습니다. 고질적인 Exposure bias를 노이즈 학습 기법으로 억제함으로써 기존 SOTA 모델을 뛰어넘는 구조 생성 퀄리티를 달성했으며, 동시에 별도의 재학습 없이 사용자의 의도대로 구조를 짜맞추는 강력한 Zero-shot 제어 능력과 효율적인 추론 속도를 입증했습니다. 이 성과는 대규모 언어 모델에서 증명된 AR 아키텍처의 무한한 확장성과 범용성을 3D 바이오 분자 설계의 영역으로 성공적으로 이식하여, 향후 복잡한 맞춤형 단백질 설계와 구조 동역학 분석의 새로운 패러다임을 열어줄 핵심적인 연구입니다.


쉬운 설명
기존의 Autoregressive(AR) 단백질 모델은 그림을 그릴 때 왼쪽 위 픽셀부터 오른쪽 아래 픽셀까지 한 줄씩 순서대로 꼼꼼히 점을 찍어가는 방식이었습니다. 이러다 보니 나중에 그릴 픽셀이 처음에 찍어둔 픽셀과 전체적인 조화를 이루지 못하고 엉뚱한 모양이 되는 문제가 있었습니다(단방향 예측의 한계).
이 논문에서 만든 PAR 모델은 실제 미술가들이 그림을 그리는 방식을 사용합니다. 처음에는 전체적인 스케치(거친 뼈대 레이아웃)를 한 번에 쓱쓱 잡아내고, 그 다음 단계에서 점차 세부적인 명암과 디테일(미세한 원자 위치)을 덧칠해 나갑니다(Coarse-to-fine). 밑그림을 그릴 때는 창의력을 발휘해 다양하게 그리고(SDE 샘플링), 디테일을 팔 때는 실수 없이 빠르고 정교하게 칠하기 때문에(ODE 샘플링), 기존 방식보다 훨씬 빠르면서도 완벽한 퀄리티의 3D 단백질 조각상을 완성해 냅니다.

 

 

 

Abstract

우리는 coarse-to-fine next-scale prediction을 통해 protein backbone generation을 위한 최초의 multi-scale autoregressive framework인 protein autoregressive modeling (PAR)을 제시합니다.

단백질의 계층적 특성을 이용하여, PAR은 조각상을 조각하는 것을 모방하는 structures를 generates하여, coarse topology를 형성하고 scales에 걸쳐 구조적 세부 사항을 refining합니다.

이를 달성하기 위해, PAR은 세 가지 핵심 components로 구성됩니다: (i) training 동안 여러 scales에 걸쳐 protein structures를 represent하는 multi-scale downsampling operations; (ii) multi-scale information을 encodes하고 structure generation을 guide하기 위해 conditional embeddings를 produces하는 autoregressive transformer; (iii) 이러한 embeddings에 conditioned되어 backbone atoms를 generates하는 flow-based backbone decoder.

더욱이, autoregressive models는 training 및 generation procedure mismatch로 인해 발생하는 exposure bias를 겪으며, 이는 structure generation quality를 크게 degrades합니다. 우리는 noisy context learning 및 scheduled sampling을 adopting함으로써 이 문제를 효과적으로 완화하여, robust backbone generation을 가능하게 합니다.

특히, PAR은 강력한 zero-shot generalization을 보여주며, fine-tuning을 요구하지 않고 유연한 human-prompted conditional generation 및 motif scaffolding을 지원합니다.

unconditional generation benchmark에서, PAR은 단백질 distributions를 효과적으로 learns하고 높은 design quality의 backbones를 produces하며, 유리한 scaling behavior를 보여줍니다.

종합하면, 이러한 특성들은 PAR을 protein structure generation을 위한 유망한 framework로 자리 잡게 합니다.

 

 

1 Introduction

더보기

단백질의 deep generative modeling은 원하는 기능과 특성을 가진 새로운 구조를 design하고 model하는 방법으로 부상했으며, 생의학 및 나노 기술 분야에서 광범위하게 응용되고 있습니다. 널리 채택되는 접근법은 단백질 기능을 지배하는 3차원 단백질 구조의 distribution을 직접 model하는 것입니다. 일반적으로 structure generative models는 sequences나 측쇄 없이 단백질 backbones를 produce합니다. 이 분야의 이전 연구는 SE(3) backbone frame representations를 predict하는 방법과 단순성과 scalability를 위해 Cα 좌표와 같은 원자를 직접 model하는 방법으로 크게 분류할 수 있습니다. 그러나 이러한 모든 연구는 diffusion models와 그 변형(예: flow matching)을 기반으로 합니다.

반면에, autoregressive (AR) modeling은 large language models를 위한 강력한 paradigm으로 부상했습니다. AR models는 이전 token을 기반으로 각 token의 확률을 model하기 위해 next-token prediction을 채택하며, scalability 및 보이지 않는 작업에 대한 zero-shot generalization과 같은 놀라운 경험적 동작을 보여줍니다.

다른 도메인에서의 성공에도 불구하고, AR modeling은 backbone modeling에서 거의 주목받지 못했습니다. 우리는 두 가지 주요 이유를 식별했습니다. (i) AR models를 3차원의 원자 위치와 같은 연속적인 데이터로 확장하는 것은 종종 데이터 discretization에 의존하는데, 이는 단백질의 structural fidelity와 세밀한 디테일을 감소시켜 generative performance를 제한할 수 있습니다. (ii) 단백질 잔기는 강한 양방향 의존성(bidirectional dependencies)을 보입니다. sequence 상에서 멀리 떨어진 잔기들이 공간적으로 가까울 수 있으며 수소 결합이나 소수성 접촉을 형성할 수 있습니다. 이러한 상호 의존성은 표준 AR models의 단방향(unidirectional) 가정과 충돌하며, 따라서 autoregressive structure generation에 대한 이전 시도들의 quality를 제한합니다. 따라서 자연스러운 질문이 생깁니다: 단백질 backbone design에 AR modeling을 적용할 수 있을까요?

본 논문에서 우리는 위 질문에 긍정적으로 답하며, protein backbone generation을 위한 AR models의 힘을 활용하기 위해 Protein AutoRegressive framework인 PAR을 제안합니다. 우리는 단백질의 계층적 특성에서 착안했습니다. 단백질 구조는 coarse 3D topology 및 3차 접힘 배열, 국소 2차 구조에서부터 가장 미세한 원자 좌표에 이르기까지 다양한 granularity의 scales에 걸쳐 있습니다. 따라서 PAR은 next-scale prediction을 통해 multi-scale autoregressive framework를 채택하여, 이전의 더 거친 scales에 conditioned된 상태에서 각 scale을 predict합니다. image generation의 발전에서 영감을 받은 이 전략은 AR models가 image synthesis에서 강력한 diffusion models를 최초로 능가할 수 있게 했으며, 나아가 multimodal LLMs가 통합된 text 및 image generation framework를 달성할 수 있도록 합니다.

이 multi-scale framework를 바탕으로, PAR은 세 가지 핵심 components를 포함합니다 (Fig. 1). multi-scale downsampling은 training 동안 structural context 및 targets 역할을 할 coarse-to-fine structural representations를 생성합니다. non-equivariant attention layers의 스택인 AR transformer는 Li et al.을 따라 scale-wise conditional embedding을 produce하기 위해 모든 이전 scales를 encodes합니다. flow-based backbone decoder는 Cα backbone 원자를 직접 model하기 위해 이 embedding에 conditioned됩니다. 결과적으로 PAR은 단백질 구조의 discretization과 잔기 단위의 unidirectional autoregressive ordering을 모두 피하며, 이를 통해 structural fidelity와 generative quality를 손상시키는 앞서 언급한 두 가지 한계를 극복합니다. 더욱이, ground-truth structural context에 대해 training할 때, AR models는 exposure bias를 겪게 되며, 이는 우리의 사전 연구에서 structure generation quality를 크게 떨어뜨리는 주요 과제입니다. 우리는 noisy context learning과 scheduled sampling을 통해 이러한 문제를 효과적으로 완화하여, model이 손상된 context로부터 학습할 수 있도록 합니다.

이 multi-scale 접근 방식은 몇 가지 주목할 만한 model behaviors를 도입합니다. PAR은 전역적인 topology를 확립하고 refinements를 수행하여 backbones를 generates하며, 이는 조각상을 점진적으로 걸작으로 조각하는 것과 유사합니다 (Fig. 2). unconditional generation의 경우, PAR은 유리한 scaling behavior를 보여주며, Fréchet Protein Structure Distance (FPSD)와 같은 distributional metrics에서 경쟁력 있는 결과를 도출합니다. 단일 scale에서 작동하는 diffusion models와 달리, PAR은 다양한 granularities의 입력을 유연하게 처리하므로 prompt-based generation 및 motif scaffolding과 같은 작업에서 zero-shot generalization을 보여줍니다. multi-scale 공식화는 PAR이 sampling strategies를 조율할 수 있게 하여, single-scale baselines에 비해 2.5배 빠른 sampling 속도 향상을 달성합니다. 마지막으로, PAR은 단일 scale로 제한될 때 flow-based models를 특수한 경우로 통합하는 더 일반적인 framework를 제공하며, 따라서 self-conditioning과 같은 flow-based models의 기술들과 호환성을 유지합니다.

Main contributions:

  • (i) 우리는 기존 AR 방법들의 주요 한계를 해결하는, protein backbone generation을 위한 최초의 multi-scale AR model인 PAR을 제시합니다.
  • (ii) PAR은 multi-scale downsampling, AR transformer, 그리고 flow-based decoder로 구성되어, discretization 손실을 피하면서 Cα 원자를 직접 model합니다.
  • (iii) 우리는 noisy context learning과 scheduled sampling을 통해 exposure bias를 완화하여, structure generation을 효과적으로 개선합니다.
  • (iv) 우리의 model은 거친 backbone topology를 형성하고 이를 점진적으로 refines하는 해석 가능한 generation 과정을 보여줍니다.
  • (v) Benchmarking 결과는 PAR이 단백질 데이터 distributions를 효과적으로 포착하여, training 연산량에 따라 더욱 확장되는 PDB dataset에 대해 161.0의 FPSD 점수를 달성함을 보여줍니다.
  • (vi) PAR은 효율적인 sampling 및 zero-shot generalization 잠재력을 보여주며, 이는 AR large language models의 다재다능함을 반영합니다.

 

 

 

 

더보기

Figure 1 (PAR 아키텍처 개요도)

이 그림은 PAR 모델이 단백질 구조를 어떻게 학습하고 만들어내는지를 보여주는 전체 설계도입니다. 핵심 부품 3가지가 어떻게 협력하는지 나타냅니다.

  • Multi-scale Downsample (왼쪽 아래): 복잡한 원본 단백질 구조($x$)를 가장 거친 형태부터 세밀한 형태까지 여러 단계의 해상도($x^1, x^2, \dots$)로 단순화(다운샘플링)하여 학습의 기준점을 만듭니다.
  • AR Transformer (왼쪽 위 파란색): 이 모델의 '두뇌' 역할을 합니다. 이전 해상도의 구조 정보를 보고 "다음 단계에서는 어떤 디테일이 추가되어야 하는지"를 예측하여 가이드라인(조건부 임베딩 $z^1, z^2, \dots$)을 제시합니다.
  • Flow-based Decoder (오른쪽 분홍색): AR Transformer가 준 가이드라인($z$)을 바탕으로 실제로 원자들의 위치를 깎아내는 '손' 역할을 합니다. 노이즈가 낀 거친 상태($x_t$)에서 점진적으로 선명하고 정확한 3D 좌표($x$)를 복원(디노이징)해 냅니다.

 

 

 

 

 

더보기

Figure 2 (스케일별 구조 생성 시각화)

이 그림은 PAR 모델이 실제로 단백질을 어떻게 생성하는지를 5단계(Scale 1 ~ 5)에 걸쳐 보여주는 결과물입니다. 단백질 길이($L$)에 상관없이 앞서 언급된 **"조각상을 조각하는 과정"**을 시각적으로 완벽하게 증명하고 있습니다.

  • Scale 1 (초기 단계): 디테일은 전혀 없고, 단백질이 전체적으로 어떤 공간을 차지하며 구부러질지 대략적인 전체 뼈대(Global structural layout)만 굵은 선으로 듬성듬성 잡혀 있습니다.
  • Scale 2 ~ 4 (중간 단계): 해상도가 높아지면서 뼈대가 점점 얇아지고, 단백질이 어떻게 꼬이고 접힐지에 대한 구체적인 윤곽이 드러나기 시작합니다.
  • Scale 5 (최종 단계): 알파 나선(스프링 모양)이나 베타 시트(화살표/병풍 모양) 같은 단백질 특유의 세밀하고 복잡한 2차 구조(Details)가 완벽하게 완성된 것을 볼 수 있습니다.

 

 

 

 

 

 

 

[논문 요약 노트] PAR: Protein AutoRegressive Modeling

1. 기존 연구의 한계점 (Why not AR for Proteins?)

  • 현재 단백질 백본 생성 모델은 주로 Diffusion 모델(및 Flow matching)에 편중되어 있습니다.
  • LLM에서 강력한 성능을 입증한 AR(Autoregressive) 모델을 3D 연속 데이터인 단백질에 적용 시 두 가지 치명적인 병목이 발생합니다.
    1. Discretization Loss: 연속적인 3D 원자 좌표를 이산화(Discretization)하는 과정에서 단백질의 구조적 충실도(fidelity)와 세부 정보가 손실됩니다.
    2. Bidirectional Dependency: 단백질은 서열상 멀리 떨어져 있어도 3D 공간에서 접히며 상호작용(수소결합 등)을 합니다. 이는 AR 모델의 단방향(Unidirectional) 생성 가정과 정면으로 충돌합니다.

2. 핵심 제안: PAR (Protein AutoRegressive Framework)

  • 단백질의 계층적 특성에 착안하여, 서열 기반의 좌우 생성이 아닌 Multi-scale Next-scale Prediction 프레임워크를 도입했습니다.
  • 3대 핵심 컴포넌트 구성:
    1. Multi-scale Downsampling: Coarse-to-fine 구조적 표현을 생성하여 학습 시 타겟 및 구조적 컨텍스트로 활용합니다.
    2. AR Transformer: 이전의 모든 해상도(scale) 정보를 인코딩하여 다음 스케일 생성을 안내하는 Scale-wise Conditional Embedding을 생성합니다.
    3. Flow-based Backbone Decoder: 위 임베딩을 조건으로 받아 Cα 백본 원자의 연속적인 좌표를 직접 모델링합니다. (이를 통해 구조 이산화 과정을 완전히 생략합니다.)

3. Exposure Bias (노출 편향) 해결

  • AR 모델 특성상 완벽한 정답(Ground-truth) 컨텍스트로만 학습하면, 실제 생성 단계에서 발생하는 미세한 오류가 누적되어 품질이 급락하는 노출 편향 문제가 발생합니다.
  • 이를 해결하기 위해 Noisy Context LearningScheduled Sampling을 도입, 모델이 의도적으로 손상된 컨텍스트에서도 올바른 구조를 복원하고 생성하는 법을 학습하게 하여 강건성을 확보했습니다.

4. 주요 성과 및 모델의 강점

  • 해석 가능한 생성 과정: 전체적인 위상(Topology)을 먼저 잡고 세부 구조를 점진적으로 정제하는, 마치 조각상을 깎는 듯한 과정을 거칩니다.
  • Zero-shot Generalization: 단일 스케일로 작동하는 Diffusion과 달리 다양한 해상도 입력을 유연하게 처리하므로, 파인튜닝 없이 프롬프트 기반 생성 및 모티프 스캐폴딩(Motif Scaffolding)이 가능합니다.
  • 효율성 및 확장성: 단일 스케일 베이스라인 대비 2.5배 빠른 샘플링 속도를 달성했으며, 연산량이 증가할수록 성능이 향상되는 우수한 Scaling Behavior를 입증했습니다.

쉬운 설명 : 1 Introduction

이 논문의 도입부가 제안하는 핵심 아이디어를 직관적으로 이해하려면 **"거대한 대리석 조각상 만들기"**를 떠올리시면 됩니다.

어떤 문제가 있었나요? 최근 챗GPT 같은 인공지능은 앞의 단어들을 보고 다음 단어를 예측하며 글을 써 내려가는 방식(AR 모델)에 아주 능숙합니다. 하지만 이 방식을 3차원으로 복잡하게 꼬여있는 단백질을 만드는 데 그대로 적용하면 문제가 생깁니다. 단백질은 1차원적인 글과 달리, 저 멀리 있는 부분끼리 맞닿아 상호작용하기 때문입니다. 게다가 연속적인 3차원 위치를 억지로 텍스트처럼 쪼개어 예측하려다 보니 정밀도가 뚝 떨어지는 한계가 있었습니다.

PAR 모델은 어떻게 해결했나요? 연구진은 단백질을 '왼쪽에서 오른쪽으로' 한 조각씩 이어 붙이는 대신, '전체적인 윤곽에서 점진적으로 세밀한 디테일로' 해상도를 높여가며 예측하는 획기적인 방식을 제안했습니다.

  1. 먼저 단백질의 대략적인 뼈대(거친 형태)를 잡습니다.
  2. 그 다음 단계의 AI가 이전의 거친 형태를 참고해서 조금 더 세밀한 굴곡을 만듭니다.
  3. 마지막으로 원자 단위의 정확한 위치를 부드럽게(연속적으로) 찍어냅니다. 이때 기존 방식의 단점이었던 '위치 쪼개기(이산화)'를 피하기 위해, 원자 위치를 직접 다루는 Flow-based 모델이라는 다른 방식을 결합했습니다.

현실적인 훈련 방법 도입: AI를 훈련시킬 때 항상 '완벽한 중간 과정'만 보여주면, 실전에서 자기가 조금만 실수를 해도 당황해서 전체 구조를 망쳐버립니다. 그래서 연구진은 훈련할 때 일부러 '조금 망가진(Noisy) 상태'를 보여주고 수습하는 연습을 시켰습니다. 그 결과, 실전에서도 훨씬 튼튼하고 완벽한 단백질 구조를 만들어낼 수 있게 되었습니다.

 

 

 

2 Background and Related Work

더보기

Flow and diffusion-based structure generative models.

Flow-based 및 diffusion methods는 prior distribution의 samples를 target data distribution으로 transforming하여 작동하며, protein backbone generation에 널리 적용되어 왔습니다. 이러한 methods는 frame-based Riemannian manifold representation을 사용하여 잔기당(per-residue) rotations 및 translations를 predict하거나, $C\alpha$ positions와 같은 원자 좌표를 직접 model하며, 일부 접근 방식은 측쇄(side chains)를 포함한 완전한 원자 수준의 단백질을 generating하기도 합니다. structure tokens에 대해 trained된 Discrete diffusion methods는 종종 structural fidelity를 감소시키고 generation quality를 제한합니다. single-scale인 대부분의 diffusion approaches와 달리, PAR은 짧은 것에서 긴 것으로의 parameterized upsampling autoregressive process를 사용하여 여러 scales에 걸쳐 protein structures를 models하며, 이는 다양한 structural granularities의 유연한 처리와 prompt-based generation과 같은 tasks로의 zero-shot generalization을 허용합니다. 또한, PAR은 단일 scale로 제한될 때 자연스럽게 flow-based model로 축소되므로 더 일반적인 framework를 제공합니다.

Autoregressive modeling.

Autoregressive (AR) modeling은 강력한 scalability와 zero-shot generalization 덕분에 natural language processing 및 computer vision을 주도해 왔습니다. 이 접근 방식은 unidirectional sequence에서 이전 token을 기반으로 다음 token의 distribution을 predicts하는 next-token prediction에 의존합니다. 그러나 autoregressive models를 image generation과 같은 continuous domains에 적용하는 것은 종종 VQVAE와 같은 tokenizers를 수반하는데, 이는 transformer training을 위해 데이터를 discretizes하며 fine-grained details를 버릴 수 있습니다. 최근 Li et al.은 diffusion network(예: 작은 MLP)를 위한 conditioning을 produces하는 AR model을 사용하여 image latents를 model함으로써, continuous-valued space에서 AR models의 operations를 열었습니다.

또한 데이터 특성을 보존하는 적절한 autoregressive orders를 정의하는 것이 중요합니다. next-token prediction은 본질적으로 2D image feature map을 1D sequence로 flattening하여 spatial locality를 버리기 때문에, VAR은 next-scale prediction을 도입했습니다. multi-scale VQVAE를 활용하여, image feature map은 spatial 및 bidirectional correlations를 보존하는 $n$개의 multi-scale token maps로 quantized됩니다.

우리가 아는 한, autoregressive modeling은 다른 domains에서의 성공에도 불구하고 protein structure generation에는 널리 적용되지 않았습니다. 유일한 예외는 causal transformer를 사용하여 structure tokens를 models하는 Gaujac et al.입니다. 이와 대조적으로, 우리는 flow-based backbone decoder를 사용하여 continuous backbone space에서 직접 작동하는 multi-scale autoregressive framework를 design하며, 이를 통해 protein structures의 bidirectional biophysical relations를 존중하면서 discrete token maps의 한계를 해결합니다.

 

 

[논문 요약 노트] 2 Background and Related Work

1. 기존 구조 생성 모델 (Flow and diffusion-based)과의 차별점

  • 기존 방식의 한계: 대부분의 Diffusion 및 Flow-based 모델은 **단일 해상도(Single-scale)**에서 작동합니다. 일부 Discrete diffusion 방식은 구조를 이산화(token)하여 structural fidelity 저하를 유발합니다.
  • PAR의 핵심 차별화: 짧은 것에서 긴 것으로 이어지는 Multi-scale upsampling autoregressive process를 채택했습니다.
    • 다양한 구조적 granularity를 유연하게 처리할 수 있습니다.
    • Prompt-based generation 같은 task에서 Zero-shot generalization이 가능합니다.
    • 단일 스케일로 작동하도록 제한할 경우 일반적인 Flow-based model로 환원되는, 더 포괄적이고 일반화된 framework를 제공합니다.

2. 기존 AR (Autoregressive) 모델링을 단백질에 적용할 때의 병목과 해결

  • Next-token prediction의 한계: 기존 AR은 1D sequence 기반의 단방향 생성 방식입니다. 이를 2D/3D 데이터에 적용하려면 1D로 평탄화(Flattening)해야 하는데, 이 과정에서 공간적, 양방향적 상관관계(Spatial and bidirectional correlations)가 훼손됩니다.
  • Tokenizer(Discretization)의 한계: 연속적인 데이터를 AR transformer로 학습하기 위해 VQVAE 등으로 강제 이산화(Quantization)하면 fine-grained details가 유실됩니다.
  • PAR의 도입 아이디어 (최신 비전 연구 차용):
    • Next-scale prediction (VAR 참조): 1D 토큰 순서가 아닌, 다중 해상도 맵을 점진적으로 예측하여 양방향 상관관계를 보존합니다.
    • Continuous AR (Li et al. 참조): AR 모델이 이산 토큰을 예측하는 대신, Diffusion/Flow 네트워크의 조건(conditioning)을 생성하게 하여 **연속적인 공간(Continuous-valued space)**에서 작동하도록 만듭니다.

3. 단백질 구조 생성에서의 PAR의 독창적 위치

  • 단백질 분야에서 기존에 AR을 시도한 연구(Gaujac et al.)는 Causal transformer와 Discrete structure tokens에 머물렀습니다.
  • 반면 PAR은 단백질의 물리화학적 양방향성(Bidirectional biophysical relations)을 완벽히 존중하면서, Discrete token map의 한계를 극복하기 위해 Continuous backbone space에서 직접 C$\alpha$ 좌표를 다루는 Flow-based decoder를 결합한 최초의 Multi-scale AR 모델입니다.

쉬운 설명 : 2 Background and Related Work

이 섹션은 "왜 지금까지 아무도 단백질 생성에 AR(챗GPT 같은 자동회귀 모델)을 제대로 쓰지 못했는가?" 그리고 "PAR은 남들이 실패한 지점을 어떻게 뚫어냈는가?"를 설명하는 배경지식 파트입니다.

1. 기존 단백질 AI들 (Diffusion 모델 등)

기존 방식들은 주로 한 번에 전체 단백질의 형태를 빚어내는 방식(Single-scale)이었습니다. 성능은 좋지만 유연성이 떨어져서, 응용 작업(예: 특정 조건만 주고 만들어보라는 프롬프트 기반 생성)을 시키려면 처음부터 다시 학습을 시켜야 하는 한계가 있었습니다.

2. 챗GPT 방식(AR)의 치명적인 단점 2가지

자연어 처리를 휩쓴 AR 방식을 단백질 구조에 그대로 가져오면 두 가지 큰 사고가 발생합니다.

  • 공간감 상실: 글은 "나는-밥을-먹는다"처럼 한 줄(1D)로 이어지지만, 단백질은 3차원 공간에서 앞부분과 뒷부분이 구부러져서 서로 달라붙는 등 '양방향'으로 영향을 주고받습니다. 이걸 글쓰기처럼 강제로 한 줄로 쭉 펴서(Flattening) 예측하게 하면 단백질 고유의 입체적인 얽힘 구조를 깡그리 잊어버리게 됩니다.
  • 디테일 뭉개짐: AR 모델을 학습시키려면 부드럽고 연속적인 3차원 좌표를 마치 레고 블록이나 모자이크 픽셀처럼 뚝뚝 끊어지는 디지털 데이터(Token)로 변환(이산화)해야 했습니다. 이 과정에서 미세한 원자 단위의 곡선 디테일이 뭉텅이로 날아가 버렸습니다.

3. PAR의 똑똑한 해결책

PAR은 이미지 생성 분야의 최신 천재적인 아이디어들을 단백질에 맞게 조합했습니다.

글쓰기처럼 왼쪽에서 오른쪽으로 1차원적으로 예측하는 대신, '흐릿한 전체 뼈대 $\rightarrow$ 선명한 세부 구조' 순서로 해상도를 높여가며 예측(Next-scale prediction)하도록 바꿨습니다. 덕분에 단백질의 복잡한 얽힘 구조(양방향성)를 지켜냈습니다. 또한, 좌표를 억지로 디지털 블록(토큰)으로 쪼개지 않고 연속적인 3D 공간을 그대로 다룰 수 있는 수학적 기법(Flow-based)을 결합해서, 아주 미세하고 부드러운 단백질 구조의 디테일까지 살려낸 것입니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

3 Protein Autoregressive Modeling

더보기

이 섹션에서는 protein backbone generation을 위한 multi-scale autoregressive (AR) framework인 PAR을 소개합니다.

공식적으로, 우리는 L개의 잔기를 가진 단백질 백본 C$\alpha$ 구조 $x \in \mathbb{R}^{L \times 3}$를 다음과 같이 autoregressive 방식으로 모델링하고자 합니다.

$$p_\theta(x) = \mathbb{E}_{X \sim q_{decompose}(\cdot|x)} p_\theta(X = \{x^1, \dots, x^n\}) = \mathbb{E}_{X \sim q_{decompose}(\cdot|x)} \prod_{i=1}^n p_\theta(x^i | X_{<i})$$

(1)

여기서 $q_{decompose}(\cdot|x)$는 단백질 구조 $x$에 대한 autoregressive order의 분해를 $x^n = x$ $n$개의 스케일 $X = {x^1, \dots, x^n}$로 정의하며, $p_\theta(x^i | X_{<i})$는 scale-wise autoregression을 통해 $x$를 generate하도록 학습하는 목표 PAR model입니다.

이 공식(식 1)에서 $q_{decompose}$와 $p_\theta$의 design space는 유연할 수 있습니다. 앞서 1장에서 논의했듯이, 우리의 목표는 AR modeling이 spatial dependencies를 보존하고 discretization을 피할 수 있게 하는 것임을 상기하십시오. 이를 위해, 3.1절에서 우리는 hierarchical down-sampling(식 2)을 통해 여러 스케일에서 단백질 백본을 나타내는 multi-scale protein downsampling (Fig. 1)을 통해 non-parametric하고 deterministic한 $q_{decompose}$를 고안하여 structural context와 training targets를 제공합니다. 3.2절에서는 PAR $p_\theta$를 next-scale prediction을 통한 backbone autoregressive upsampling process로 parameterize하고, continuous space에서 직접 C$\alpha$ modeling을 달성합니다 (식 3). 이는 두 가지 핵심 components로 구성됩니다: (i) generation을 가이드하기 위해 이전 스케일들의 정보를 바탕으로 scale-wise conditional embeddings를 생성하는 autoregressive transformer (Fig. 1) (식 4); 그리고 (ii) 학습된 임베딩에 조건화되어 C$\alpha$ 백본 좌표를 샘플링하는 flow-based backbone decoder (Fig. 1) (식 5).

마지막으로 3.3절에서는 exposure bias를 완화하기 위한 전용 전략을 다룹니다. exposure bias는 실제 정답 데이터로 학습하는 것과 모델 예측으로 추론하는 것 사이의 불일치로, AR models에서 오류가 누적되고 generation quality를 저하시키는 원인이 됩니다. 이러한 components가 모여 PAR이 coarse-to-fine 방식으로 견고하게 protein backbones를 generate할 수 있게 합니다.

3.1 Multi-scale Protein Downsampling

우리는 PAR을 위한 training context 및 targets 역할을 할 hierarchical downsampling을 통해 단백질 구조의 multi-scale representations를 구성합니다 (Fig. 1). 단백질 구조 $x \in \mathbb{R}^{L \times 3}$가 주어지면, $x$를 점진적으로 $n$개의 스케일로 다운샘플링하여 coarse-to-fine 스케일의 계층 구조를 생성합니다.

$$q_{decompose} : x \mapsto X = \{x^1, x^2, \dots, x^n\} = \{\text{Down}(x, \text{size}(1)), \text{Down}(x, \text{size}(2)), \dots, x\}$$

(2)

여기서 $\text{Down}(x, \text{size}(i)) \in \mathbb{R}^{\text{size}(i) \times 3}$은 시퀀스 차원을 따라 $x$를 보간(interpolate)하는 다운샘플링 연산을 나타내며, 대략적인 구조적 레이아웃을 제공하는 size($i$) 개의 3D 중심점(centroids)을 도출합니다. $q_{decompose}$가 모든 $x$에 대해 deterministic mapping으로 설계되었기 때문에, 식 1의 likelihood는 주변화(marginalization) 없이 단순화될 수 있습니다: $p_\theta(x) = \prod_{i=1}^n p_\theta(x^i | X_{<i})$. 우리는 이 다운샘플링 전략이 pairwise spatial relationships를 적절하게 보존함을 부록 C.8에서 보여줍니다.

Scale configurations. $S = {\text{size}(1), \dots, \text{size}(n)}$은 두 가지 방식으로 정의할 수 있습니다. 길이(length)로 정의할 때 스케일은 하이퍼파라미터로 선택됩니다 (예: $S = \{64, 128, 256\}$). 이 경우 $L$이 (size($i$), size($i + 1$)] 범위에 있다면, 단백질은 단 $i+1$번의 autoregressive steps만으로 생성될 수 있습니다. 비율(ratio)로 정의할 때 스케일은 단백질 길이에 따라 적응적으로 결정됩니다 (예: $S = \{L/4, L/2, L\}$). 경험적으로, 길이로 스케일을 정의하는 것이 data distributions를 모델링할 때 약간 더 나은 결과를 산출합니다. 우리는 이를 기본 구성으로 채택합니다. 이 설계는 유연한 스케일 구성으로 PAR을 훈련할 수 있게 해줍니다. 다음 섹션에서는 이러한 표현의 계층 구조가 autoregressive transformer 및 backbone decoder를 사용하여 어떻게 모델링되는지 설명합니다.

3.2 Coarse-to-Fine Backbone Autoregressive Modeling

autoregressive order를 정의할 때 데이터에 내재된 종속성(dependencies)을 보존하는 것은 매우 중요하며 generation performance에 영향을 미칩니다. 표준 AR models는 단방향 종속성(unidirectional dependency)을 가정하는데, 이는 단백질 서열 내의 강력한 양방향 상호작용(bidirectional interactions)과 충돌합니다. 예를 들어, 공간적으로 가까운 잔기들은 서열상 멀리 떨어져 있더라도 소수성 접촉이나 수소 결합을 형성할 수 있습니다. PAR은 각 스케일에 걸친 상호 구조적 종속성을 포착하는 next-scale prediction을 통한 multi-scale AR framework로 이 문제를 해결합니다. Li et al.에서 영감을 받아, 우리는 continuous space에서 직접 C$\alpha$ 원자를 모델링할 수 있도록 diffusion/flow-based regression loss가 있는 AR Transformer를 사용할 것을 제안합니다. 즉, likelihood를 다음과 같이 다시 쓸 수 있습니다.

$$p_\theta(X = \{x^1, \dots, x^n\}) = \prod_{i=1}^n p_\theta(x^i | X_{<i}) = \prod_{i=1}^n p_\theta(x^i | z^i = \mathcal{T}_\theta(X_{<i}))$$

(3)

여기서 $\mathcal{T}_\theta$는 scale-wise conditioning $z^i$를 생성하는 AR Transformer이며, $p_\theta(x^i|z^i)$는 flow matching이 있는 flow-based atomic decoder $v_\theta$로 최적화됩니다. 이는 단백질 구조를 토큰으로 이산화(discretizing)하는 것을 방지하여 structural details와 generation fidelity를 보존합니다. 각 구성 요소는 아래에 설명되어 있습니다.

Autoregressive transformer for scale-wise conditioning. autoregressive order를 공식화하기 위해, 우리는 단백질 구조가 대략적인 3차 위상(topology)부터 가장 미세한 원자 좌표까지 다양한 수준의 표현에 걸쳐 있을 수 있다는 단백질의 계층적 특성을 활용합니다. 우리는 이전의 거친 스케일을 기반으로 스케일별 분포를 모델링하기 위해 next-scale prediction을 채택하며, 이는 각 스케일에 걸쳐 잔기의 양방향 종속성(bidirectional dependencies)이 모델링되도록 보장합니다. 우리는 non-equivariant transformer인 autoregressive model $\mathcal{T}_\theta$ (Fig. 1)를 훈련하여 이전 스케일 $X_{<i} = {x^1, \dots, x^{i-1}}$에 따라 스케일 $i$에 대한 scale-wise conditioning embedding $z^i$를 생성합니다.

$$z^i = \mathcal{T}_\theta(X_{<i}) = \mathcal{T}_\theta\Big(bos, \text{Up}(x^1, \text{size}(2)), \dots, \text{Up}(x^{i-1}, \text{size}(i))\Big)$$

(4)

여기서 $bos \in \mathbb{R}^{\text{size}(1) \times 3}$은 학습 가능한 임베딩이며, $\text{Up}(x^{i-1}, \text{size}(i))$는 $x^{i-1}$을 size($i$) 3D 포인트로 보간(interpolates)합니다. 모든 입력은 $\mathcal{T}_\theta$에 주입되기 전에 시퀀스 차원을 따라 연결(concatenated)됩니다. 그런 다음 임베딩 $z^i$는 아래에 자세히 설명된 대로 백본 좌표 $x^i$를 예측하기 위해 flow matching decoder를 조건화(condition)하는 데 사용됩니다.

Flow-based atomic decoder. 우리는 PAR이 C$\alpha$ 위치 $x$를 직접 모델링할 수 있도록 하며, 여기서 $p_\theta(x|z^i)$는 표준 정규 분포를 타겟 데이터 분포로 매핑하는 flow matching [FM] 방식을 사용하는 atomic decoder $v_\theta$에 의해 매개변수화(parameterized)됩니다. 우리는 각 스케일 $i$에서 AR Transformer $\mathcal{T}_\theta$에 의해 예측된 scale-wise conditioning $z^i$ $v_\theta$를 조건화합니다 (Fig. 1). 훈련 중에 우리는 노이즈 $\epsilon^i \sim \mathcal{N}(0, I)$와 시간 변수 $t^i \in [0, 1]$를 샘플링하고, 보간된(interpolated) 샘플을 $x_{t^i}^i = t^i \cdot x^i + (1 - t^i) \cdot \epsilon^i$로 계산합니다. 이와 같이, 우리는 FM objective를 사용하여 $v_\theta$ $\mathcal{T}_\theta$를 공동으로 훈련할 수 있습니다.

$$\mathcal{L}(\theta) = \mathbb{E}_{x \sim p_D} \Bigg[ \frac{1}{n} \sum_{i=1}^n \frac{1}{\text{size}(i)} \mathbb{E}_{t^i \sim p(t^i), \epsilon^i \sim \mathcal{N}(0,I)} \bigg\| v_\theta(x_{t^i}^i, t^i, z^i) - (x^i - \epsilon^i) \bigg\|^2 \Bigg]$$

(5)

여기서 $p_D(x)$는 훈련 데이터 분포를 나타내고 $p(t)$는 Geffner et al.의 t-sampling 분포를 나타냅니다. conditioning embedding $z^i$는 adaptive layer norms를 통해 atomic decoder 네트워크 $v_\theta$에 주입됩니다. 우리는 모델이 서로 다른 스케일을 식별하는 것을 돕고 self-conditioning 입력을 추가 조건으로 통합하기 위해 $z^i$ 옆에 학습 가능한 스케일 임베딩을 추가로 연결(concatenate)하지만, 단순화를 위해 방정식에서는 생략했습니다. 스케일 $i$에서 positional encoding $p^i$의 인덱스를 공식화하기 위해, 우리는 구간 [1, L]에서 균일하게 size($i$) 개의 숫자를 샘플링합니다 (즉, $p^i = \text{linspace}(1, L, \text{size}(i))$). 거친 스케일에서는 인접 인덱스 사이의 넓은 간격이 모델로 하여금 전체적인 구조적 레이아웃을 캡처하도록 장려하는 반면, 더 미세한 스케일에서는 조밀한 인덱스가 모델이 국소적인 디테일에 집중할 수 있게 합니다. 더 자세한 내용은 부록 A.1을 참조하십시오.

학습된 flow network $v_\theta$를 활용하면 상미분 방정식(ODE) $dx_t = v_\theta(x_t, t)dt$를 통해 각 스케일에서 샘플링을 수행할 수 있으며, 단순화를 위해 스케일 윗첨자 $i$와 조건 $z$는 생략되었습니다. 더욱이, 우리는 샘플링을 위한 확률 미분 방정식(SDE)을 다음과 같이 정의할 수 있습니다.

$$dx_t = v_\theta(x_t, t) dt + g(t) s_\theta(x_t, t) dt + \sqrt{2g(t)\gamma} dW_t$$

(6)

여기서 $g(t)$는 score function $s_\theta(x_t, t)$와 노이즈 항에 대한 시간에 따른 스케일링 함수이고, $\gamma$는 노이즈 스케일링 매개변수이며, $W_t$는 표준 위너 과정(Wiener process)입니다. 시간 $t$에서의 노이즈 데이터 분포의 로그 확률 기울기(gradient)로 정의되는 score function은 $s_\theta(x_t, t) = \frac{v_\theta(x_t, t) - x_t}{1 - t}$로 계산될 수 있습니다.

Multi-scale structure generation. 추론 시, autoregressive transformer는 먼저 가장 거친 스케일에서 $z^1$을 생성하며, 이는 식 6의 ODE 또는 SDE 샘플링을 통해 $x^1$을 생성하도록 flow matching decoder를 조건화합니다. 우리는 $\text{Up}(x^1, \text{size}(2))$를 사용하여 $x^1$을 업샘플링하고 다음 스케일 임베딩 $z^2$를 예측하기 위해 이를 다시 autoregressive transformer로 보냅니다. 이 coarse-to-fine 과정은 flow-matching 모델이 완전한 해상도의 백본 $x$를 생성할 때까지 $n$번 반복됩니다. 효율성을 위해 autoregressive 과정 내내 KV cache가 적용됩니다.

3.3 Mitigating Exposure Bias

AR models 훈련은 일반적으로 학습을 안정화하기 위해 정답(ground-truth) 데이터가 컨텍스트로 제공되는 teacher forcing을 사용합니다. 그러나 추론하는 동안 모델은 자체 예측에 조건화되어 exposure bias로 알려진 훈련-추론 불일치를 만듭니다. 그런 다음 오류가 autoregressive steps에 걸쳐 누적되어 출력 품질이 저하될 수 있습니다. 우리의 사전 연구는 teacher forcing이 생성된 구조의 designability를 크게 감소시킨다는 것을 보여줍니다. 이를 완화하기 위해, 우리는 언어 및 이미지 AR modeling의 기술인 Noisy Context Learning (NCL)과 Scheduled Sampling (SS)을 PAR에 적용합니다.

Noisy context learning. 우리는 훈련 중 정답 이전 스케일 입력에 노이즈를 추가하는, 노이즈가 있는 컨텍스트로 PAR을 훈련합니다. 이는 모델이 완벽하게 정확한 컨텍스트에 의존하지 않고 스케일별 분포를 학습하도록 장려하여 robustness를 향상시킵니다. 우리는 $n$개의 노이즈 가중치 ${w_{ncl}^1, \dots, w_{ncl}^n} \in [0, 1]$을 무작위로 샘플링하고 $n$개의 노이즈 샘플 ${\epsilon_{ncl}^1, \dots, \epsilon_{ncl}^n} \sim \mathcal{N}(0, I)$을 추출합니다. 각 입력 컨텍스트 $x^i$ $x_{ncl}^i = w_{ncl}^i \cdot x^i + (1 - w_{ncl}^i) \cdot \epsilon_{ncl}^i$로 손상됩니다. 이러한 섭동(perturbation)은 훈련 중에만 입력 컨텍스트에 적용되며, 이는 식 4의 autoregressive step을 $z^i = \mathcal{T}\theta\Big(bos, \text{Up}(x{ncl}^1, \text{size}(2)), \dots, \text{Up}(x_{ncl}^{i-1}, \text{size}(i))\Big)$로 업데이트합니다.

Scheduled sampling. 훈련 중에, 우리는 스케일에 걸쳐 순방향(forward) 과정을 반복적으로 실행함으로써 scheduled sampling을 사용합니다. $i$번째 스케일에서, flow-based backbone decoder는 깨끗한 데이터 $x_{pred}^i = x_t^i + (1 - t^i)v_\theta(x_t^i, t^i, z^i)$를 예측합니다. 0.5의 확률로, 우리는 이후 스케일에서 정답 컨텍스트 $x^i$를 이 예측된 $x_{pred}^i$로 대체합니다. 이는 모델을 자체 출력에 노출시켜 훈련-테스트 격차(train-test gap)를 줄입니다. 특히, 우리는 모델이 예측한 컨텍스트 $x_{pred}^i$에 노이즈를 추가함으로써 이 기술과 noisy context learning을 결합할 수 있었습니다.

 

 

[논문 요약 노트] 3 Protein Autoregressive Modeling

1. Multi-scale AR Framework Formulation (수식적 정의)

  • 단백질 백본 $C\alpha$ 구조 $x$$n$개의 해상도(scale)로 분해하여 점진적으로 생성하는 방식을 수식화했습니다.
  • $$p_\theta(x) = \mathbb{E}_{X \sim q_{decompose}(\cdot|x)} \prod_{i=1}^n p_\theta(x^i | X_{<i})$$
  • 이 프레임워크는 확정적(Deterministic)인 분해 과정($q_{decompose}$)과, 이전 스케일들을 조건으로 다음 스케일을 예측하는 학습 모델($p_\theta$)로 구성됩니다.

2. 3.1 Multi-scale Protein Downsampling (정답 컨텍스트 생성)

  • 핵심: $q_{decompose}$를 학습 가능한 파라미터 없이 Sequence dimension을 따라 보간(Interpolation)하는 연산으로 정의했습니다.
  • 효과: 이산화(Discretization) 없이 각 스케일마다 3D Centroids를 추출하여 거친(Coarse) 구조적 레이아웃을 생성합니다.
  • Scale Configurations: 스케일을 단백질 길이의 비율(Ratio)로 정하기보다, 고정된 길이(Length, 예: 64, 128, 256)로 하이퍼파라미터화하는 것이 데이터 분포 모델링에 더 효과적임을 실험적으로 확인했습니다.

3. 3.2 Coarse-to-Fine Backbone Autoregressive Modeling (모델 아키텍처)

  • 기존 AR 모델의 단방향 토큰 예측(Next-token prediction)이 가진 한계를 벗어나고자, Next-scale prediction을 도입하여 단백질 서열의 양방향 종속성(Bidirectional dependencies)을 스케일 단위로 포착합니다.
  • AR Transformer ($\mathcal{T}_\theta$): 이전까지 생성된 모든 거친 스케일의 구조($X_{<i}$)를 업샘플링하여 입력받고, 다음 스케일 생성을 안내하는 조건부 임베딩(Scale-wise conditioning embedding, $z^i$)을 계산합니다.
  • Flow-based Atomic Decoder ($v_\theta$): $\mathcal{T}_\theta$가 만든 $z^i$에 조건화되어, Continuous space에서 타겟 데이터 분포로 향하는 Flow Matching(FM) objective를 수행합니다. 이를 통해 이산화 과정 없이 $C\alpha$ 좌표를 직접 모델링하며, 추론 시에는 식에 명시된 ODE 또는 SDE를 통해 샘플링을 진행합니다.

4. 3.3 Mitigating Exposure Bias (노출 편향 완화 전략)

  • Teacher forcing으로만 훈련된 AR 모델은 자기 자신의 예측 오류가 누적되는 환경(추론 단계)에 취약해져 Generation quality가 급락합니다. 이를 해결하기 위해 두 가지 기법을 도입했습니다.
  • Noisy Context Learning (NCL): 훈련 시 주어지는 이전 스케일의 정답(Ground-truth) 컨텍스트에 의도적으로 노이즈($\epsilon$)를 추가합니다. 모델이 완벽하지 않은 컨텍스트에서도 올바른 스케일별 분포를 복원하도록 강제합니다.
  • Scheduled Sampling (SS): 훈련 중 50%의 확률로, 정답 컨텍스트 대신 모델 스스로가 Flow-based decoder를 통해 방금 예측한 결과물($x^i_{pred}$)을 다음 스케일의 컨텍스트로 사용합니다. NCL과 결합하여 예측 컨텍스트에 노이즈를 추가할 수도 있습니다.

쉬운 설명 : 3 Protein Autoregressive Modeling

이 섹션은 PAR 모델이 실제로 **"어떤 수학적 원리와 네트워크 구조를 가지고 단백질을 만들어내는지"**를 설명하는 논문의 심장부입니다.

1. 스케치 준비하기 (Downsampling)

그림을 배울 때 처음부터 눈썹 한 올을 그리는 사람은 없습니다. 전체적인 동그라미(머리)와 막대기(몸통)로 뼈대를 잡고, 점점 디테일을 추가하죠. 논문에서도 진짜 단백질 구조(정답)를 여러 단계로 흐릿하게 뭉개서 '1단계 스케치', '2단계 스케치', '완성작' 형태의 정답지들을 미리 만들어 둡니다.

2. 뇌와 손의 협업 (AR Transformer & Flow Decoder)

모델은 크게 **'설계자(뇌)'**와 **'조각가(손)'**로 나뉩니다.

  • 설계자 (Transformer): 1단계 스케치를 보고 "음, 2단계에서는 이 부분이 헬릭스(나선) 모양으로 접혀야겠군" 하고 다음 단계의 밑그림(임베딩 $z$)을 구상합니다.
  • 조각가 (Flow Decoder): 설계자의 밑그림을 전달받아, 실제로 3차원 공간상에서 원자들의 x, y, z 좌표를 정밀하게 깎아냅니다. 여기서 토큰(레고 블록)처럼 뚝뚝 끊어지는 방식을 안 쓰고, 부드러운 찰흙을 빚어내듯 연속적인 좌표를 맞추는 최신 기법(Flow Matching)을 썼기 때문에 훨씬 정교한 모양이 나옵니다.

3. 실전 모의고사 훈련 (Exposure Bias 완화)

보통 AI를 학습시킬 때는 항상 "완벽한 이전 단계 정답"만 보여주고 다음을 예측하게 합니다. 하지만 실전(추론)에서는 AI가 스스로 만든 불완전한 결과를 보고 다음을 만들어야 하니 당황해서 구조가 와르르 무너지는 현상(노출 편향)이 생깁니다.

그래서 연구진은 **"모래주머니 훈련법"**을 썼습니다. 학습할 때 일부러 이전 단계 스케치에 노이즈(잡음)를 잔뜩 끼워 넣거나(NCL), 정답을 뺏고 AI 본인이 방금 만든 삐뚤빼뚤한 스케치를 던져주며(SS) 다음 단계를 그리게 했습니다. 이런 스파르타식 훈련 덕분에, 실전에서 자기가 조금 실수를 해도 훌륭하게 단백질 구조를 완성할 수 있는 강인한 모델이 탄생했습니다.

 

 

 

 

 

 

 

 

 



 

 

 

더보기


실제 단백질의 3D 원자 좌표 데이터(x)를 준비

L x  3좌표

이걸 다운 샘플링을 해버림 L방향으로

64 128 256 이게 정답임. 길이 안되면 걍 거기서 끝
그리고 단계별로 노이즈를 기록해둠.


첫 번째 스케일은 이전 정보가 없으므로 학습 가능한 bos 토큰을 입력 z1 출력

x1에 노이즈를 섞은 현재 좌표, t, 방금만든 z1을 입력해서
x1으로 가기위한 방향 학습.

그리고
bos 토큰이랑 Up(x1_ncl, 128) 을 이어서 z2 뽑고 
디코더에서 다시 x2를 보내기 위한 두번째 방향 학습



**z1**은 bos만 보고 만들었기에, "그냥 아무 단백질이나 일단 64개 점으로 기본 덩어리를 잡아라"라는 막연한 지시

**z2**는 bos에 더해 1단계의 형태($x^1$)를 힌트로 보고 만들었기에, "방금 만든 64개 뼈대 모양을 유지하면서, 128개 점으로 더 세밀하게 깎아라"라는 구체적인 지시


이런 식으로 3단계(256)까지 가면, 이전 단계의 z값들이 계속 누적되면서 모델이 점점 더 정교한 단백질 구조를 완성하게됨

이를 실제 만들어둔 정답과 함께 loss로 보냄. 




별점 3.5점 / 5점

단백질 3D 좌표를 억지로 쪼개는 Tokenization의 한계를 피하기 위해, 비전 분야의 Coarse-to-fine 기법과 Flow Matching을 영리하게 조립하여 연속적인 3D 공간에서 뼈대를 안정적으로 깎아낸 엔지니어링 감각은 훌륭함

무거운 3D 물리 법칙 연산을 회피하고자 시퀀스 길이를 선형 보간으로 듬성듬성 솎아내는 '단순한 꼼수'를 택했는데, 이는 미세한 원자 구조가 생명인 단백질에서 정보 손실을 운에 맡긴 치명적인 타협