Search

FrameDiff 논문 리뷰

FrameDiff 논문 리뷰

들어가며

RFdiffusion 이후 protein backbone diffusion은 꽤 익숙한 단어가 됐습니다. Target surface 주변에 backbone을 만들고, ProteinMPNN으로 sequence를 붙이고, AF2나 ESMFold로 다시 확인하는 흐름은 이제 여러 design pipeline의 기본 형태처럼 보입니다. 그런데 여기서 한 단계 더 아래로 내려가면 질문이 생깁니다. Protein backbone에 diffusion을 건다는 것은 정확히 어떤 공간에서 diffusion을 돌린다는 뜻일까요.
FrameDiff는 이 질문에 답하려는 논문입니다. 2023년 ICML에 발표된 “SE(3) diffusion model with application to protein backbone generation”은 protein backbone을 residue-level rigid frame들의 집합으로 보고, SE(3)^N 위에서 diffusion을 정의합니다. 논문은 RFdiffusion처럼 binder wet-lab validation을 보여주는 실전 design paper라기보다, protein backbone diffusion을 수학적·공학적으로 어떻게 세울 수 있는지 정리한 method paper에 가깝습니다.
이 차이를 먼저 잡고 들어가야 합니다. FrameDiff는 RFdiffusion의 대체재로 나온 강력한 binder generator가 아닙니다. Pretrained structure prediction network 없이, monomer backbone generation에서 designability, diversity, novelty를 computational proxy로 평가한 early SE(3) backbone diffusion baseline입니다. 그래서 이 글에서는 FrameDiff를 “실험적으로 검증된 protein design platform”이 아니라 “backbone diffusion의 좌표계를 정리한 논문”으로 보겠습니다.

Backbone을 residue frame으로 보기

Protein backbone은 단순한 point cloud가 아닙니다. 각 residue에는 N, Cα, C backbone atom이 있고, 이 세 atom은 local coordinate frame을 정의할 수 있습니다. FrameDiff는 AlphaFold2-style residue frame parameterization을 사용해 각 residue를 rigid transformation `T_n = (R_n, x_n)`으로 표현합니다. 여기서 `R_n`은 rotation, `x_n`은 translation입니다.
즉 생성 대상은 amino-acid sequence가 아닙니다. FrameDiff가 생성하는 것은 residue frame collection, 다시 말해 `SE(3)^N` 위의 object입니다. Oxygen atom 위치는 별도 torsion angle `ψ`로 구성합니다. 이 framing은 protein backbone generation을 단순한 Cα coordinate generation보다 더 구조적으로 다루게 해줍니다.
왜 SE(3)일까요. Protein backbone은 global rotation과 translation이 바뀌어도 같은 structure입니다. 따라서 model은 전체 좌표계가 어디에 놓였는지보다 residue 사이의 상대 geometry를 다뤄야 합니다. FrameDiff는 이 문제를 rotation과 translation을 나눠 다루는 SE(3) diffusion으로 정식화합니다.

Centered SE(3) diffusion

FrameDiff의 technical core는 centered SE(3)^N diffusion입니다. Rotation 쪽에는 SO(3) Brownian motion을 쓰고, translation 쪽에는 Ornstein–Uhlenbeck 또는 VP-SDE-like process를 씁니다. SO(3) transition density는 compact Lie group heat kernel series, 즉 IGSO3 형태로 계산하고, 여기서 rotation score를 얻습니다.
문제는 translation입니다. R3에는 translation-invariant probability measure가 없습니다. 전체 protein을 통째로 어디로 옮겨도 같은 structure로 봐야 하는데, Euclidean translation coordinate 자체에 완전한 global translation invariance를 넣기는 어렵습니다. FrameDiff는 center of mass를 zero로 고정한 centered process를 정의하고, 매 step에서 translation을 recentering합니다.
이 선택은 단순 구현 detail이 아닙니다. Protein backbone diffusion이 global SE(3) symmetry를 어떻게 다룰지에 대한 논문의 핵심 답입니다. 이후 여러 protein structure generator가 SE(3), residue frame, manifold diffusion, flow matching을 이야기할 때 FrameDiff가 배경으로 계속 언급되는 이유도 여기에 있습니다.

FramePred와 structure-specific engineering

FrameDiff의 neural network는 FramePred입니다. AlphaFold2 structure module에서 온 idea를 사용해 node embedding, edge embedding, residue frame을 여러 layer에 걸쳐 업데이트합니다. Invariant Point Attention, residue-level transformer, edge update, BackboneUpdate-style frame update, self-conditioning이 들어갑니다.
모델은 최종적으로 denoised frame `T(0)`와 oxygen placement를 위한 torsion angle `ψ`를 예측합니다. 논문에서 사용한 모델은 약 17.4M parameters로 비교적 작습니다. Fully connected graph를 쓰기 때문에 memory와 computation은 length에 대해 quadratic하게 증가합니다.
흥미로운 점은 loss입니다. Core는 rotation/translation score를 맞추는 denoising score matching loss입니다. 하지만 DSM만으로는 coarse topology는 그럴듯해도 chain break나 steric clash 같은 fine-grained violation이 생겼다고 합니다. 그래서 논문은 backbone atom coordinate MSE인 `Lbb`, 6 Å 이내 local pairwise atomic distance loss인 `L2D` 같은 auxiliary losses를 추가합니다. Auxiliary losses는 small timestep에서만 적용됩니다.
이 대목은 FrameDiff를 읽을 때 중요합니다. 이 논문은 “순수한 이론적 SE(3) diffusion만으로 모든 것이 해결됐다”는 이야기가 아닙니다. SE(3) score matching, self-conditioning, backbone geometry auxiliary loss, sampling recipe가 함께 작동한 engineering result입니다.

Sampling recipe와 designability-diversity tradeoff

FrameDiff는 reverse diffusion sampling 중 final step에서 backbone이 destabilize되는 문제가 있어 trajectory를 `ε > 0`에서 truncate합니다. 또 RFdiffusion처럼 reverse noise scale `ζ`를 줄이는 sampling도 탐색합니다. 이 sampler choice는 결과에 큰 영향을 줍니다.
Table 1은 noise scale과 step 수, ProteinMPNN sequence 수에 따른 designability와 diversity를 보여줍니다. 기본적으로 `ζ = 1.0`, 500 steps, 8 sequences에서는 `scTM > 0.5`가 49%, `scRMSD < 2 Å`가 11%, diversity가 0.75입니다. Noise scale을 `ζ = 0.5`로 낮추면 `scTM > 0.5`는 74%, `scRMSD < 2 Å`는 23%로 올라가지만 diversity는 0.56으로 내려갑니다. `ζ = 0.1`에서는 `scRMSD < 2 Å`가 28%까지 올라가고 diversity는 0.53입니다.
ProteinMPNN sequence 수를 늘리는 것도 영향을 줍니다. `ζ = 0.1`, 500 steps, 100 sequences에서는 `scTM > 0.5`가 84%, `scRMSD < 2 Å`가 40%까지 올라갑니다. 하지만 이건 backbone 하나당 더 많은 sequence를 설계하고 더 많이 refolding해 보는 setting입니다. 따라서 designability가 올라간다고 해서 generator 자체만 좋아졌다고 해석하면 곤란합니다. Evaluation compute와 filtering 기회가 함께 늘어난 결과입니다.
이 tradeoff는 지금도 protein generator benchmark를 읽을 때 중요합니다. Noise를 줄이면 더 깔끔한 sample이 나오지만 diversity가 줄 수 있습니다. 더 많은 sequence를 붙이면 self-consistency 성공률이 오르지만 compute와 search가 늘어납니다. Headline designability만 보면 이런 차이가 가려집니다.

Evaluation: ProteinMPNN + ESMFold self-consistency

FrameDiff의 designability는 실험 결과가 아닙니다. Generated backbone에 ProteinMPNN으로 sequence를 붙이고, ESMFold로 다시 structure를 예측한 뒤, generated backbone과 predicted backbone이 맞는지 봅니다. 논문은 `scTM`과 `scRMSD`를 사용합니다. `scTM`은 self-consistency TM-score이고, `scRMSD`는 self-consistency Cα-RMSD입니다.
Chroma 같은 논문은 `scTM > 0.5`를 designable criterion으로 사용했지만, FrameDiff는 `scRMSD < 2 Å`가 특히 long backbone에서 더 엄격한 기준이라고 강조합니다. TM-score는 global topology similarity에는 유용하지만, 긴 backbone에서는 local mismatch가 있어도 높은 값을 받을 수 있습니다. 반면 scRMSD < 2 Å는 더 stringent한 recapitulation 기준입니다.
여기서 말하는 designability는 “ProteinMPNN이 붙인 sequence를 ESMFold가 generated backbone과 비슷하게 접는다”는 뜻입니다. Foldability 가능성을 지지하는 proxy일 수는 있지만, 실제 발현, folding, stability, function을 보여주지는 않습니다. 이 구분은 FrameDiff 리뷰 전체에서 유지하는 편이 좋습니다.

Ablation: 이론과 실무 사이

Table 2 ablation은 FrameDiff의 성능이 어떤 구성요소에 기대는지 보여줍니다. Self-conditioning, local distance loss, backbone loss, DSM rotation loss를 모두 쓴 setting이 가장 좋습니다. 이 조합에서 `scTM > 0.5`는 49%입니다. DSM 대신 Frobenius rotation loss를 쓰면 39%로 떨어지고, auxiliary loss를 줄이면 42%, 22%, 16%로 내려갑니다. DSM rotation loss 없이 일부 loss만 쓰는 경우는 0%로 보고됩니다.
이 결과는 두 가지를 말합니다. 첫째, SO(3) score matching을 제대로 다루는 것이 실제 성능에 중요합니다. 둘째, protein backbone generation에서는 manifold diffusion의 원칙만큼이나 local geometry를 잡아주는 auxiliary loss가 중요합니다. FrameDiff는 이론 논문처럼 보이지만, 좋은 backbone을 만들기 위해 꽤 구체적인 structure-specific engineering을 사용합니다.

Diversity와 novelty

FrameDiff는 diversity를 MaxCluster clustering으로 평가합니다. Generated backbones를 structural cluster로 묶고, unique clusters / samples 비율을 diversity로 봅니다. Novelty는 FoldSeek으로 PDB에서 가장 가까운 chain을 찾고, highest TM-score를 `pdbTM`으로 둡니다. 논문은 `scRMSD < 2 Å`이면서 `pdbTM < 0.6`인 sample을 novel and designable candidate로 봅니다.
이 novelty claim도 proxy입니다. PDB nearest-neighbor와 낮은 TM-score를 갖는다는 것은 기존 PDB 구조와 다르게 보인다는 뜻이지, 실제로 새로운 fold가 안정적으로 존재한다는 뜻은 아닙니다. 논문도 Discussion에서 novel protein generation claim에는 experimental characterization이 필요하다고 조심스럽게 말합니다.
따라서 FrameDiff의 novelty는 “PDB search 기준으로 novel-like한 computational candidate” 정도로 두면 충분합니다. Chroma처럼 X-ray structure를 얻은 것도 아니고, RFdiffusion처럼 binder wet-lab validation이 있는 것도 아닙니다. 이 제한이 논문의 약점이라기보다, FrameDiff의 evidence layer를 정확히 표시하는 기준입니다.

RFdiffusion, Chroma, Proteus와의 위치

FrameDiff는 RFdiffusion과 자주 비교되지만, 두 논문의 역할은 다릅니다. RFdiffusion은 pretrained RoseTTAFold/RF lineage를 활용해 motif scaffolding, symmetric design, target-conditioned binder design까지 보여준 practical milestone입니다. BLI와 cryo-EM validation도 포함되어 있어 binder design evidence가 강합니다.
FrameDiff는 pretraining-free SE(3) backbone diffusion baseline입니다. RFdiffusion보다 designability는 낮지만, model size가 작고, pretrained structure predictor에 기대지 않으며, residue-frame diffusion을 원칙적으로 정의했다는 장점이 있습니다. 논문도 Chroma, RFdiffusion과의 비교가 training data, model scale, evaluation setup 차이 때문에 엄밀한 head-to-head가 어렵다고 경고합니다.
Proteus는 이후 FrameDiff류 pretraining-free line에서 designability와 efficiency를 끌어올리려는 후속 branch로 볼 수 있습니다. Proteus는 FrameDiff loss를 기반으로 하면서 architecture와 data, sampling efficiency를 개선합니다. 이런 계보에서 FrameDiff는 최고 성능 모델이라기보다, “SE(3) backbone diffusion을 이렇게 세울 수 있다”는 기준점을 제공한 논문입니다.

평가: backbone diffusion의 좌표계를 만든 논문

FrameDiff의 가치는 flashy한 benchmark보다 formulation에 있습니다. Protein backbone을 residue frames on SE(3)로 보고, SO(3) Brownian motion과 centered translation process를 결합하고, equivariant score network로 denoising하는 길을 제시했습니다. Protein backbone generation을 Euclidean point cloud sampling보다 더 자연스러운 geometric object 위에서 다룬 셈입니다.
동시에 이 논문은 protein generator 평가를 읽는 법도 가르쳐줍니다. `scTM > 0.5`, `scRMSD < 2 Å`, diversity, novelty, ProteinMPNN sequence count, reduced noise sampling 같은 요소가 headline 성능을 크게 바꿉니다. FrameDiff의 결과는 in silico self-consistency benchmark이지 wet-lab validation이 아닙니다. Monomer unconditional backbone generation 논문이지 binder/interface design 논문도 아닙니다.
그래서 FrameDiff는 protein design 시리즈에서 배경 infrastructure에 가까운 글입니다. RFdiffusion이 practical binder generation의 전환점이라면, FrameDiff는 backbone diffusion을 어떤 수학적 공간에서 정의할지 보여준 method anchor입니다. 이후 Chroma, MultiFlow, Proteus, Proteina 계열을 읽을 때 FrameDiff는 계속 돌아오게 되는 좌표계입니다.

참고

Yim, J. et al. “SE(3) diffusion model with application to protein backbone generation”, ICML/PMLR, 2023.
arXiv: https://arxiv.org/abs/2302.02277
Code: https://github.com/jasonkyuyim/se3_diffusion
주요 비교 축: RFdiffusion, Chroma, Proteus, MultiFlow, ProteinMPNN.