SaDiT: 단백질 구조 설계의 새로운 지평, 토큰화와 디퓨전의 융합

초록

SaDiT는 단백질의 3차원 구조를 나타내는 ‘구조적 토큰’을 활용해 생성 과정을 압축하고, Diffusion Transformer 아키텍처와 IPA Token Cache 메커니즘을 결합해 기존 최고 모델 대비 월등한 속도와 구조적 타당성을 달성한 단백질 백본 생성 프레임워크입니다.

상세 분석

SaDiT 논문의 기술적 핵심은 ‘고차원의 연속적 좌표 공간’에서 ‘이산적(discrete)이고 압축된 구조 토큰 공간’으로 생성 문제의 패러다임을 전환한 데 있습니다. 기존 RFDiffusion이나 FrameDiff와 같은 모델들은 단백질의 원자 좌표나 프레임을 직접 노이즈 제거하는 방식으로, 수백 단계의 반복적 샘플링이 필요해 계산 부하가 컸습니다. SaDiT는 이 문제를 세 가지 혁신으로 해결합니다.

첫째, SaProt 구조 토큰화를 채택합니다. SaProt 인코더는 연속적인 단백질 백본 지오메트리를 유한한 크기의 ‘구조 단어’ 코드북으로 양자화합니다. 이는 마치 자연어 처리에서 문장을 단어 토큰으로 나누는 것과 유사한, 구조 영역의 ‘압축’입니다. 이 토큰 공간은 물리적으로 가능한 구조들로 정규화되어 있어, 생성 과정 초기부터 무의미한 구조 탐색을 줄여주고, 최종 디코딩 시에도 물리적 제약 조건을 위반할 가능성을 낮춥니다.

둘째, Diffusion Transformer(DiT) 백본을 도입합니다. 기존 U-Net이나 그래프 신경망 대신 트랜스포머를 생성 모델의 핵심으로 사용해, 단백질 내 장거리 상호작용(Long-range dependency)을 효과적으로 모델링합니다. 여기에 Adaptive Layer Norm(adaLN)을 통해 디퓨전 타임스텝 정보를 주입하여, 노이즈가 많은 초기 상태와 구조가 결정되는 후기 상태에서 다른 특징 추출 전략을 적용할 수 있게 합니다.

셋째, IPA Token Cache 메커니즘이라는 독창적인 최적화 기술을 제안합니다. 단백질 구조 생성에 필수적인 Invariant Point Attention(IPA)은 SE(3) 등가성을 보장하지만 계산 비용이 높습니다. SaDiT는 반복 샘플링 과정에서 토큰들의 공간적 유사도(affinity)가 수렴함에 따라, 이전 스텝에서 계산된 어텐션 상태를 캐싱하고 재활용합니다. 이는 생성 후반부로 갈수록 계산량을 극적으로 줄여, 샘플링 복잡도를 거의 선형에 가깝게 만드는 결정적 요소입니다.

이론적으로 SaDiT는 SE(3) 등가성을 엄격히 유지합니다. 인코더가 상대적 거리와 방향만을 사용하므로 토큰 공간 자체가 SE(3)-불변(invariant)이고, 디코더가 상대적 변환(ΔT)을 예측하여 전체 구조를 재구성하기 때문에 최종 출력은 전역적 회전/평행이동에 대해 등가적(equivariant)입니다. 이는 생성된 구조의 생물학적 타당성을 보장하는 근간이 됩니다.

실험 결과, SaDiT는 무조건 생성 및 폴드 클래스 조건 생성 과제에서 RFDiffusion, Proteína보다 빠른 샘플링 속도와 동등하거나 더 높은 ‘설계 가능성(designability)‘을 보였습니다. 이는 단백질 AI 생성 모델이 정확도뿐만 아니라 실용적인 계산 효율성의 단계로 진입했음을 의미하며, 대규모 가상 스크리닝과 실시간 설계 응용에 대한 길을 열었다는 점에서 파급력이 큽니다.