AI바라기의 인공지능
ATOMSURF: SURFACE REPRESENTATION FOR LEARNING ON PROTEIN STRUCTURES 본문
용어 설명
- Surface representation: 단백질의 내부 원자 연결 구조 대신, 외부 표면의 기하학적 형태와 생화학적 특성을 3D 메쉬(mesh)로 나타낸 모델링 방식.
- DiffusionNet: 3D 표면 위에서 열이 퍼져나가는 방식(Heat equation)을 학습하여, 해상도나 형태 변화에 강건하게 기하학적 특징을 추출하는 state-of-the-art surface encoder.
- Bipartite graph (이분 그래프): 두 종류의 노드 그룹(여기서는 graph node와 surface vertex) 간에만 엣지(edge)가 존재하도록 연결한 구조.
- Coarsening: 3D 메쉬의 정점(vertex) 수를 뭉쳐서 줄여 해상도를 낮추는 기법으로 연산 효율성을 높이기 위해 사용됨.
- Message Passing: 네트워크 내에서 연결된 노드들끼리 정보를 주고받으며 자신의 feature를 업데이트하는 과정.
Purpose of the Paper
- 기존의 한계: 단백질 구조 학습 분야에서 graph나 sequence 기반 모델은 활발히 연구되었으나, 단백질 상호작용 분석에 직관적인 Surface representation 기반 모델들은 공정한 benchmark 비교가 부족했음. 또한 기존 연구들은 surface 정보를 단독으로만 사용하거나, 기껏해야 모델의 가장 마지막 단계에서 graph 정보와 단순 결합(global pooling)하는 데 그쳤음.
- 새로운 접근 방식: 본 논문은 state-of-the-art surface encoder를 단백질 데이터의 특성에 맞게 직접 개량하여 공정한 성능 평가를 진행함. 나아가, 내부 원자 구조(graph)와 외부 표면 형태(surface)가 모델의 모든 layer에서 세밀하게 정보를 주고받는 혁신적인 Hybrid representation learning 구조를 제안하여 기존 독립적 모델링의 한계를 극복하고자 함.
Key Contributions
- DiffusionNet의 단백질 데이터 최적화 (Novelty: 생물학적 사전 지식 반영)
- 기존 DiffusionNet은 학습 안정성을 위해 모든 3D 객체의 크기를 동일하게 정규화(scale normalization)했으나, 단백질 상호작용에서는 실제 크기(scale) 정보가 필수적임.
- 정규화 없이도 다양한 크기의 단백질을 학습할 수 있도록, 결합 부위의 실제 스케일(약 10 Å)에 맞춘 diffusion time 초기화와 Batch Normalization을 새롭게 도입하여 학습 안정성을 확보함.
- 순수 Surface 모델의 한계 규명 (Novelty: 최초의 공정한 비교 실험)
- Atom3D benchmark 환경에서 동일한 parameter와 input feature 조건으로 공정한 비교를 최초로 수행.
- 단백질 결합 등 표면이 중요한 task에서조차 surface 단독 모델이 graph 모델이나 SOTA를 능가하지 못한다는 점을 실험적으로 명확히 밝혀냄.
- Node-wise 기반의 Bipartite Message Passing 구조 제안 (Novelty: 전 계층에 걸친 공간적 융합)
- 기존의 단순 결합(early/late fusion) 방식에서 벗어나, 3D 공간상에서 물리적으로 인접한 graph node와 surface vertex 간에 Bipartite graph를 구성함.
- 이를 통해 네트워크의 특정 계층이 아닌 모든 layer에서 양방향으로 feature sharing이 일어나는 아키텍처를 고안하여 두 representation 간의 강력한 시너지를 이끌어냄.
Experimental Highlights
- 실험 설정: Atom3D benchmark의 주요 task(PIP, MSP, PSR) 및 실제 응용에 가까운 결합 부위 예측 dataset(MaSIF-ligand, PINDER) 활용. 비교 baseline으로는 ProNet, GVP, GearNet-ESM 및 기존 surface 모델 적용.
- State-of-the-art 성능 달성:
- 제안한 AtomSurf 모델은 parameter 수가 기존 최고 성능 모델들 대비 절반 이하(600k)임에도 압도적인 SOTA 달성.
- 단백질 상호작용 예측(PIP)에서 90.9 AuROC, 돌연변이 안정성 예측(MSP)에서 71.6 AuROC 기록.
- 리간드 결합 선호도 예측(MaSIF-ligand)에서 88% Balanced Accuracy로 최고 성능 경신.
- Coarsening을 통한 연산 효율성 증명: 메모리와 시간 소모가 컸던 surface 데이터에 mesh coarsening을 적용하여 병목 지점인 정점 수를 대폭 감소시킴. 그 결과 기존의 효율적인 graph encoder들보다 더 빠른 학습 및 추론 속도를 확보함.
Limitations and Future Work
- Limitations: Surface representation 특성상 graph 기반 방식보다 본질적으로 메모리 사용량이 매우 큼. 특히 단백질마다 형태와 정점(vertex) 수의 편차가 극심하여, 운이 나쁘게 거대한 단백질들이 하나의 batch에 몰릴 경우 훈련 중 메모리 초과(out-of-memory)가 발생함. 이로 인해 모델 학습 시 큰 batch size를 사용할 수 없다는 치명적인 제약이 있음.
- Future Work: Surface network의 메모리 footprint를 근본적으로 줄이는 파이프라인 최적화가 필수적임. 향후 이 메모리 제약이 해결되면 generative modeling이나 복잡한 구조 생물학 task 등 실제 응용(real-life scenarios)으로 확장하여, 신약 개발 등에 즉각적으로 투입할 수 있는 강력한 multi-modal 솔루션으로 발전시킬 수 있음.
Overall Summary
이 논문은 단백질을 분석할 때 외부 표면(surface)과 내부 뼈대(graph) 정보를 네트워크의 모든 layer에서 유기적으로 결합하는 AtomSurf라는 새롭고 강력한 hybrid architecture를 제안했습니다. 기존처럼 두 정보를 독립적으로 처리하거나 단순 통합하는 방식의 한계를 극복하고, 3D 공간의 근접성을 바탕으로 각 노드와 정점들이 긴밀하게 정보를 교환하도록 설계했습니다. 그 결과, 파라미터 수를 대폭 줄였음에도 불구하고 주요 단백질 상호작용 benchmark에서 압도적인 state-of-the-art를 달성하며, 향후 3D 구조 기반 신약 설계 및 단백질 디자인 분야의 핵심 기반 모델로 자리 잡을 수 있는 높은 잠재력을 입증했습니다.
쉬운 설명
건물의 구조적 안정성과 쓰임새를 분석할 때, 기존의 AI 방식이 '건물 설계도(내부 뼈대 = graph)'만 보거나 '건물 외관 사진(외부 표면 = surface)'만 각각 따로 보고 판단하는 것이었다면, 이 논문의 모델은 마치 엑스레이 안경을 쓴 건축가와 같습니다. 즉, "외벽의 이 오목한 굴곡(surface)은 내부에 이 철근(graph) 구조가 지탱하고 있기 때문이구나!"라고 실시간으로 연결 지어 생각하며(모든 layer에서의 양방향 정보 교환), 건물의 취약점이나 다른 건물과 맞닿을 수 있는 결합 부위를 훨씬 빠르고 정확하게 예측해 내는 획기적인 모델입니다.
ABSTRACT
protein data에서의 learning을 위한 다양한 representations를 평가하고 비교하는 데 있어 상당한 진전이 있었지만, surface-based learning approaches의 역할은 여전히 잘 이해되지 않고 있습니다. 특히, 가용한 최고의 surface-based learning methods와 graphs와 같은 대안적인 representations 간의 직접적이고 공정한 benchmark 비교가 부족한 실정입니다. 더욱이, 소수의 기존 surface-based approaches는 surface information을 독립적으로 사용하거나, 기껏해야 surface 및 graph-based architectures 사이에서 global pooling을 수행하는 데 그칩니다.
본 연구에서는 먼저 protein learning tasks를 위해 state-of-the-art surface encoder를 적용함으로써 이러한 격차를 해소합니다. 그런 다음 Atom3D benchmark 내에서 도출된 method와 대안적인 approaches를 직접적이고 공정하게 비교하여, 순수한 surface-based learning이 가지는 한계를 강조합니다. 마지막으로, 우리는 모든 layers에 걸쳐 nodes 및 vertices 수준에서 graphs와 surface representations 간의 learned feature sharing을 가능하게 하는 integrated approach를 제안합니다.
우리는 결과적인 architecture가 엄격한 benchmark protocol을 준수하면서 Atom3D benchmark의 모든 tasks에서 state-of-the-art results를 달성할 뿐만 아니라, 더 넓게는 binding site identification 및 binding pocket classification에서도 우수한 성과를 달성함을 입증합니다. 또한, 우리는 coarsened surfaces를 사용하고 efficiency를 위해 우리의 approach를 optimize하여, 기존 기술들과 비교했을 때 training 및 inference time 측면에서 우리의 도구가 경쟁력을 갖추도록 만듭니다.
1 INTRODUCTION
구조 생물정보학 데이터는 전례 없는 속도로 가용해지고 있습니다. 특히 극저온 전자 현미경(cryo-EM)의 발전은 더 큰 시스템과 더 나은 해상도뿐만 아니라 실험적으로 파생된 구조들을 그 어느 때보다 더 많이 생성하는 결과를 가져왔습니다. AlphaFold의 개발과 많은 후속 연구들은 단백질 구조를 풍부하게 사용할 수 있게 만들었으며, Protein Data Bank (PDB)에 100만 개 이상의 고품질 예측이 있고 ESM Metagenomic Atlas (ESMatlas)에는 6억 개 이상이 존재합니다. 따라서 이러한 구조적 데이터를 활용하여 구조 생물정보학 및 신약 설계 분야를 발전시키는 데 도움을 줄 수 있는 machine learning 기술에 대한 수요가 증가하고 있습니다.
단백질 구조는 원자 좌표뿐만 아니라 기하학적 구조에 의존하는 복잡한 생화학적 상호작용 모두로 특징지어지는 복잡한 객체입니다. learning pipeline에서 사용되기 위해서는 단백질 구조를 잘 정의된 수학적 객체로 변환하는 초기 모델링 단계가 필요합니다. 서로 다른 수학적 representations는 서로 다른 구조적 및 생물학적 priors를 인코딩합니다. 예를 들어, point cloud representation은 화학적 상호작용에 의해 유도된 연결성을 무시하지만, 데이터에 대한 가장 일반적인 기하학적 설명을 가능하게 합니다. 단백질 표면(Protein surfaces)은 단백질 내부의 세밀한 정보를 외부 표면에 대한 정확한 묘사와 교환하는 것을 선택합니다. 이 representation은 차폐 효과(screening effect)로 인해 표면의 특성에 주로 의존하는 활성 상호작용 부위를 연구하는 데 특별한 관심사가 될 것으로 생각됩니다. 그러나 표면 용어에 의해 지배되는 상호작용의 경우에도, 내부에 대한 지식은 표면의 안정성을 인코딩할 수 있습니다. 이러한 서로 다른 representations는 Supplementary Figure 1에 설명되어 있습니다.
각 데이터 유형을 처리하기 위해 전문화된 architectures가 개발됨에 따라, representation의 선택은 learning-based methods의 맥락에서 특히 두드러집니다. 접근 방식의 범위는 geometric deep learning 분야 내에서 연구되며, graphs, point clouds, surfaces, 데이터의 group symmetry를 존중하는 equivariant methods, equivariant message passing 등 서로 다른 데이터 유형을 처리하기 위한 전문화된 방법들이 개발되었습니다.
몇몇 선구적인 연구들은 3D convolutional networks, equivariant convolutional networks, sequence, surfaces, graphs 그리고 equivariant discrete networks를 사용하여 구조 생물학 데이터 representations에 geometric deep learning을 적용했습니다. 그 뒤를 이어 그들이 사용하는 수학적 representation을 기반으로 전통적으로 분류되는 몇몇 다른 연구들이 나타났습니다. 또한 단백질 특성이 network에 내장되도록(baked into) 단백질 구조를 처리하기 위해 특정 목적(ad-hock)으로 개발된 일부 방법들도 있었습니다.
이러한 맥락에서, Atom3d의 획기적인 연구는 잘 정의된 protocol 내에서 서로 다른 representations와 learning approaches 모두에 걸쳐 공정한 비교를 목표로 합니다. 특히 이 benchmark는 3차원 분자 구조를 위한 9개의 benchmark tasks 세트를 포함하며, 테스트된 모든 methods에 걸쳐 사용될 일관된 input features 및 parameter count 세트를 설정합니다. 저자들은 또한 제안된 tasks에 대해 3D grids, graphs, 그리고 equivariant networks 기반의 neural networks를 평가하여 서로 다른 representations를 비교합니다.
단백질에 대해 단일 representation을 사용하는 것을 넘어서, 여러 수학적 객체로서 단백질의 동시 representation은 유망합니다. 실제로 서로 다른 representations는 데이터의 서로 다른 생물학적 priors를 인코딩하고 서로 다른 계산적 이점을 제시합니다. 잘 연구된 조합은 구조의 graph representation과 함께 sequence 정보를 사용하는 것입니다. 예를 들어 저자들은 sequence를 인코딩하는 추가적인 edge types로 graph를 풍부하게 만듭니다. graph에 sequence 정보를 포함하는 또 다른 방법은 sequence embeddings를 사용하는 것이며, 특히 protein language models에서 파생되어 방대한 양의 가용한 sequence 데이터로부터 이점을 얻는 것을 사용하는 것입니다. 마지막으로, 일부 접근법들은 protein language models의 training에 단백질 구조에서 파생된 정보를 포함합니다.
1 INTRODUCTION 핵심 정리 노트 (AI 연구자용)
단백질 구조 데이터(AlphaFold, cryo-EM 등)의 폭발적 증가로 인해, 이를 효과적으로 활용할 machine learning 기술의 수요가 급증하고 있는 배경 속에서 본 논문이 제기하는 핵심 문제의식과 연구 동향은 다음과 같습니다.
- Representation에 따른 구조적/생물학적 priors의 차이: 단백질을 수학적 객체로 모델링할 때 어떤 방식을 택하느냐에 따라 모델이 학습하는 정보가 달라집니다.
- Point cloud: 화학적 상호작용에 의한 연결성을 무시하지만, 가장 범용적인 기하학적 형태를 제공합니다.
- Surfaces: 단백질 내부의 세밀한 정보를 희생하는 대신 외부 표면을 정확히 묘사합니다. 이는 주로 표면 특성에 의존하는 활성 상호작용 부위 연구에 유리하지만, 표면의 안정성을 결정짓는 내부 구조 정보가 누락된다는 한계가 있습니다.
- Geometric deep learning의 적용과 한계: 그래프, point clouds, surfaces, equivariant methods 등 각 representation에 특화된 architectures가 개발되어 구조 생물학에 적용되어 왔습니다. 그러나 서로 다른 모델과 데이터 표현 방식 간의 공정한 비교가 어려웠으며, 이를 해결하기 위해 Atom3D와 같은 표준화된 benchmark가 등장했습니다.
- Multi-representation의 부상: 단일 representation의 한계를 극복하기 위해 여러 수학적 객체를 동시에 사용하는 접근법이 유망한 방향으로 제시되고 있습니다. 기존 연구들은 주로 graph representation에 sequence 정보(예: protein language models의 sequence embeddings)를 결합하는 방식을 성공적으로 활용해 왔으며, 이는 본 논문이 새로운 representation 결합(abstract에서 언급된 surface와 graph의 결합)을 시도하는 논리적 기반이 됩니다.
쉬운 설명 :
단백질 구조 데이터가 알파폴드 같은 기술 덕분에 엄청나게 많아졌고, 이제 이걸 AI로 잘 분석하는 게 아주 중요해졌습니다. 그런데 단백질을 AI 모델에 입력하려면 먼저 수학적인 형태(representation)로 바꿔줘야 합니다.
이때 단백질을 어떻게 표현하느냐에 따라 AI가 받아들이는 정보가 확 달라집니다. 예를 들어 단백질을 겉면의 '표면(Surfaces)' 모양으로만 표현하면, 단백질이 다른 물질과 어떻게 결합하는지 파악하기는 아주 좋습니다. 하지만 단백질의 뼈대를 이루는 내부 정보는 잃어버리게 되죠. 반대로 내부 원자들을 점이나 그래프로 다 찍어주면 세밀하지만 계산이 복잡해집니다.
그래서 지금까지는 표면 전용 AI, 그래프 전용 AI 등 각자 자기 분야의 AI 모델(architectures)들만 따로 발전해 왔고, 서로 공정하게 성능을 겨루기가 힘들었습니다. (그래서 Atom3D 같은 통일된 시험 무대가 나왔습니다.)
결론적으로 이 도입부가 말하고 싶은 핵심은 **"단백질을 한 가지 형태로만 보지 말고, 여러 형태의 장점을 섞어서 같이 보면 훨씬 좋지 않을까?"**라는 점입니다. 기존에는 그래프 모양에 아미노산 서열(sequence) 정보를 섞어 쓰는 방식이 유행했는데, 이 논문은 이러한 흐름을 바탕으로 새로운 시도(표면+그래프)를 하겠다는 밑밥을 깔고 있는 것입니다.
2 MOTIVATION AND CONTRIBUTION
protein data에서의 learning을 위해 다양한 representations를 비교하는 데 있어 최근의 진전이 있었음에도 불구하고, surface-based representations가 여러 응용 분야에서 유망한 결과를 보여주었음에도 상대적으로 덜 주목받았습니다. surface representation에 기반한 approaches는 일반적으로 초기 MaSIF 논문의 검증을 따랐으며, 따라서 확립된 단일 benchmark의 맥락에서 다른 representations와 직접적으로 비교된 적이 없습니다. 동시에 DiffusionNet과 같은 강력한 surface-based encoders가 최근 geometry processing/computer graphics 문헌에서 제안되었으며, 이는 이전 연구의 기반이 된 초기 Geodesic-CNN based techniques를 robustness와 accuracy 측면에서 크게 능가합니다. 불행히도 현재 가용한 최고의 surface-based encoders가 protein analysis tasks(예: Atom3d benchmark)에서 다른 learning-based paradigms와 어떻게 비교되는지는 현재 알려져 있지 않습니다.
우리는 먼저 현재의 state-of-the-art surface-based learning architecture를 protein analysis tasks에 맞게 조정함으로써 이러한 격차를 메웁니다. 그런 다음 고정된 input features와 parameter count를 가지고 benchmark 프로토콜을 준수하면서, 순수한 surface-based learning method에 대한 최초의 공정하고 포괄적인 비교를 수행합니다. 이 분석의 핵심 발견은 surface-based encoders가 competitive하지만 state-of-the-art 결과를 제공하지는 않는다는 것입니다.
그런 다음 우리는 protein analysis를 위한 surface-based learning이 다른 representations의 정보와 complementary information을 제공할 수 있는지 탐구하는 데 집중합니다. 최근 문헌에서 관련 노력이 이루어졌습니다. 한 연구에서는 surface의 implicit representation을 pretraining objective로 사용할 것을 제안합니다. 다른 연구에서는 먼저 surface 속성을 인코딩하고 이를 graph nodes의 초기 embeddings로 사용할 것을 제안한 반면, 또 다른 연구들에서는 surface-based 및 graph-based model에 의해 만들어진 예측을 평균화했습니다. 그럼에도 불구하고 이러한 노력들은 surface와 기타(예: graph-based) learning을 개별적으로 고려하고 global manner(early or late fusion)로만 결과를 집계합니다. 대신 우리는 features가 중간 layers 내에서도 surface와 graph representation 사이에서 공유되고 전달되는 integrated approach를 생성함으로써 performance를 크게 향상시키고 결과를 개선할 수 있음을 보여줍니다. 결정적으로, graph nodes와 surface vertices 사이에 존재하는 근접성 기반의 자연스러운 spatial relations를 활용함으로써, 우리는 node-wise feature sharing이 두 representations 사이에 synergy를 창출함을 보여줍니다. 나아가, text encoders의 embeddings를 input features로 사용하고 신중하고 효율적인 architecture 설계와 결합함으로써 광범위한 tasks에서 전례 없는 state-of-the-art 결과를 달성할 수 있음을 입증합니다.
요약하자면, 우리의 주요 contributions는 다음과 같습니다:
- protein analysis tasks의 맥락에서 그 한계(instabilities 및 scale independence 포함) 중 일부를 해결하는 최근의 state-of-the-art DiffusionNet architecture의 조정된 설계.
- 확립된 benchmark 내에서 graphs 또는 grids와 같은 대안적인 representations에 대한 surface-based learning의 최초의 포괄적인 비교.
- jointly하게 학습되는 surface 및 graph-based encoders 사이의 모든 학습된 layers에 걸친 node-wise feature-sharing을 기반으로 하는 새로운 integrated approach.
- residue graphs 및 coarsened meshes를 사용하여 다양하고 도전적인 시나리오에서의 state-of-the-art 결과 및 향상된 computational throughput.
논문의 나머지 부분은 다음과 같이 구성됩니다: Section 3.1에서는 우리가 사용하는 surface representation과 이를 처리하는 데 사용되는 specialized networks를 제시합니다. Section 3.2에서는 다양한 scales에서 학습하는 surface networks의 과제를 강조하고 이러한 문제를 완화하기 위한 솔루션을 제시합니다. Section 3.3에서는 두 representations의 힘을 활용하여 unified architecture 내에서 graph와 surface 정보를 시너지 효과가 나도록 통합할 것을 제안합니다. Section 3.4에서는 선택된 architecture에 관한 세부 정보를 제공하고 Section 3.5에서는 그것의 computational aspects를 분석합니다.
2 MOTIVATION AND CONTRIBUTION 핵심 정리 노트 (AI 연구자용)
1. 연구 배경 및 기존의 한계 (Motivation)
- 표면(Surface) 데이터의 평가 부재: MaSIF 등 surface-based representation이 유망함을 보였으나, 확립된 benchmark(예: Atom3D)에서 graph 등 타 representation과 공정하게 비교된 적이 없음.
- 최신 기술 적용의 부재: Computer graphics 분야에서 강력한 성능을 입증한 최신 surface-based encoder(예: DiffusionNet)가 단백질 분석 task에 어떻게 적용되고 어느 정도의 성능을 내는지 미지수임.
- 기존 Multi-representation 융합의 한계: Surface와 Graph를 함께 사용하려는 시도들은 있었으나, 단순히 초기 embedding으로 사용하거나 마지막 예측값을 평균 내는 수준(early or late fusion / global aggregation)에 그침.
2. 핵심 발견 및 방법론 제안 (Methodology & Discovery)
- Pure Surface Model의 한계 확인: SOTA surface encoder를 단백질에 맞게 조정하여 최초로 공정 비교를 수행한 결과, competitive하지만 SOTA 달성에는 미치지 못함을 확인.
- Deep Integration & Node-wise Feature Sharing 제안: 단순한 초기/후기 결합을 넘어, 학습되는 모든 중간 layers(middle layers)에서 surface의 vertices와 graph의 nodes 간에 공간적 근접성(spatial proximity)을 기반으로 직접적인 feature sharing이 일어나는 통합 아키텍처(integrated approach)를 구축함.
3. 주요 기여 (Key Contributions)
- 단백질 분석 맥락에서 발생하는 불안정성과 scale independence 문제를 해결하기 위해 최신 DiffusionNet architecture의 구조를 커스텀(adapted design).
- 확립된 benchmark 내에서 순수 surface-based learning과 대안적 representation(graphs, grids 등) 간의 최초의 포괄적이고 공정한 성능 비교 제공.
- 모든 layer에 걸쳐 surface와 graph-based encoder가 공동 학습(jointly learned)되며 상호 작용하는 새로운 Node-wise feature-sharing 통합 모델 제안.
- Residue graphs와 coarsened meshes, Text encoders의 embeddings를 결합하여 다양한 task에서 새로운 SOTA 달성 및 연산 효율성(computational throughput) 대폭 향상.
쉬운 설명 :
기존의 단백질 AI 연구들은 단백질의 '겉모습(표면)'을 분석하는 AI와 '내부 뼈대(그래프)'를 분석하는 AI가 따로 놀았습니다. 두 개를 같이 써보려는 시도도 있었지만, 각자 독립적으로 계산을 다 끝낸 다음 맨 마지막에 나온 결과물만 합치거나(Late fusion), 아예 맨 처음에만 정보를 살짝 넘겨주는(Early fusion) 수준이었습니다. 비유하자면, 두 명의 전문가가 각자 방에 들어가서 따로 프로젝트를 다 끝낸 뒤에 보고서만 합친 격이죠.
이 논문이 기존 연구들과 확연히 다르고 성능이 폭발적으로 좋아진 이유는, **'모든 작업 단계(모든 layer)마다 겉모습 담당 AI와 뼈대 담당 AI가 실시간으로 정보를 주고받게 만들었기 때문(Node-wise feature sharing)'**입니다.
실제 3D 공간상에서 단백질 표면의 특정 지점과 내부 뼈대의 특정 위치가 가깝게 붙어있다는 점을 이용해, 학습하는 내내 서로의 특징(feature)을 교환하며 시너지를 일으키게 구조를 짠 것입니다. 거기에 더해 최신 AI 기술들을 단백질 구조에 맞게 튜닝하고 최적화하여, 연산 속도도 빠르면서 역대 최고 성능(SOTA)을 달성했다는 것이 이 섹션의 핵심입니다.
3 METHODS
3.1 SURFACE REPRESENTATION LEARNING
우리의 첫 번째 목적은 protein data에서의 learning을 위한 현존하는 최고의 surface encoders의 유용성을 연구하는 것입니다. 단백질 $P$의 surface representation $S_P$를 생성하기 위해 우리는 MSMS와 Section C.1에 자세히 설명된 mesh coarsening 및 cleaning 단계에 의존합니다. 그런 다음 우리의 surface-based learning method의 기반으로 DiffusionNet approach를 사용합니다. 이 method는 다양한 shape analysis tasks 세트 전반에 걸쳐 매우 robust하고 effective한 것으로 입증되었습니다 (예: (Attaiki et al., 2021; Cao & Bernard, 2022; Sun et al., 2023; Li et al., 2022) 등). 특히, DiffusionNet은 서로 다른 mesh structures에서 instabilities를 유발할 수 있는 local patch parametrizations의 사용을 피하고, information sharing을 위해 learned diffusion을 사용함으로써 long range 및 multi-scale information propagation을 가능하게 합니다. DiffusionNet의 수학적 기반은 열의 확산, 혹은 동등하게 시간에 따른 표면 위 Brownian motion의 거동을 시뮬레이션하는 열방정식(heat equation)입니다. surface $S$에 대해, $f_t : S \rightarrow \mathbb{R}$를 시간 $t$에서 $S$ 위 열 분포를 정의하는 함수로, $\Delta_S$를 surface의 Laplace-Beltrami operator로 둡니다. 열방정식은 Laplace-Beltrami를 대각화하여 푸는 아래의 선형 미분 방정식(equation 1)입니다. 실제로는, $\Delta_S$의 가장 작은 $k=128$개의 고윳값을 대각 행렬 $\Lambda \in \mathbb{R}^{k \times k}$에, 해당하는 고유 벡터를 $\Phi \in \mathbb{R}^{n \times k}$에, 그리고 vertex area weights를 $M \in \mathbb{R}^{n \times n}$에 저장합니다. 열방정식의 spectrally truncated solution은 $f_t = \Phi e^{-\Lambda t}(\Phi^T M)f_0$ 로 주어집니다.
(Heat equation)
DiffusionNet에서 이 방정식은 $f_0$로 feature map을 사용하고 task-specific manner로 diffusion time을 학습하여 surface에서 information propagation을 수행하는 데 사용됩니다. 이 mechanism은 dense linear algebra operations에 의존하여 $f$와 $t$ 모두에 대한 간단한 differentiation을 제공합니다. diffusion layers는 $f_t$, 그것들의 spatial gradients 그리고 standard, point-wise MLPs를 기반으로 한 features와 결합됩니다. 이는 task-specific manner로 surface의 multi-scale geometric details를 포착할 수 있는 architecture로 이어집니다.
3.2 ADAPTING TO PROTEIN SURFACES OF DIVERSE SCALES
아래 Section 4에서 입증하듯이, 우리의 첫 번째 경험적 관찰은 수정 없이 DiffusionNet architecture를 단백질 datasets에 직접 적용할 경우 비교적 좋지 않은 performance를 낸다는 것입니다. 우리는 그 주된 이유가 초기 DiffusionNet architecture가 관련된 거의 등거리의(near-isometric) 모양들(예: 다른 포즈의 사람)을 포함하는 응용 분야를 목표로 했기 때문이라고 생각합니다. 핵심적인 기술적 문제는 DiffusionNet을 사용하는 기존 approaches의 대부분이 모든 모양을 균일한 표면적(uniform surface area)으로 normalize한다는 것입니다. 이 단계는 모양들이 동일한 global scale을 갖도록 보장하여 learning을 안정화시킵니다. 불행히도 protein analysis의 맥락에서, ligand-binding preference determination과 같은 tasks는 단백질과 리간드의 상대적인 크기에 결정적으로 의존하므로 단백질의 scale에 대한 고려가 필수적이며, global scale normalization은 이 귀중한 정보를 잃게 만들 것입니다.
반면에 DiffusionNet의 receptive field의 효능은 각 diffusion layer 내에서 학습된 diffusion times에 달려있습니다. scale normalization 없이 입력되는 모양들의 크기 변화는 네트워크에서 학습된 receptive field의 불일치로 이어질 수 있습니다. 이는 다음의 잘 알려진 명제에 의해 설명되며, 그 증명은 완전성을 위해 supplementary material에 제공됩니다:
Proposition 3.1. $X$를 shape라 하고, $Y = \alpha X$를 요소 $\alpha > 0$만큼 scaled된 버전이라 하자. 시간 $t$ 이후에 점 $x$에서 시작하는 Brownian motion에 대한 expected geodesic distance를 $E_\cdot(t, x)$로 표기할 때, 다음이 성립한다: $E_Y(t, x) = \alpha E_X\left(\frac{t}{\alpha^2}, x\right)$.
중요하게도, 이 결과는 diffusion의 시간 매개변수가 scale에 따라 조정되어야 함을 시사하는 반면, DiffusionNet에서 학습된 시간 매개변수는 shape independent합니다. 우리는 또한 다른 scales를 가질 수 있는 단백질과 같은 non-isometric surfaces를 포함하는 시나리오에서, 학습된 diffusion이 특히 생물학적 고려 없이 학습될 경우 매우 다양한 모양들 전반에 걸쳐 일반화 성능이 저하되고(generalize poorly), training 중 instabilities를 유발할 수 있음에 주목합니다.
이러한 문제를 해결하기 위해, 우리는 원래의 DiffusionNet framework를 두 가지 방식으로 강화합니다. 첫째, batch 지원을 가능하게 하고(원래 모델은 batch 크기가 1로 제한됨) learning을 안정화하기 위해 각 diffusion layer 뒤에 Batch Normalization layer를 통합합니다. 둘째, spatial scales와 관련된 biological priors를 통합하여 최적화 프로세스를 촉진합니다. 결과적으로 우리는 약 10의 diffusion times가 결합 부위의 spatial scale과 일치하는 약 10 Å의 receptive fields를 초래한다고 결정했습니다(Supplementary Figure 2 참조). 단백질 구조의 내재적인 multi-scale nature에서 영감을 받아, 우리는 비교적 높은 분산으로 특징지어지는 정규 분포 $t \sim N(10, 5)$에서 샘플을 추출하기로 선택했습니다. 그런 다음 이 샘플들의 절댓값을 우리의 diffusion timescales의 초기 값으로 활용했습니다. Supplementary Figure 3에 설명된 것처럼 large scale과 large variance 모두 training 중에 유지되어 효율적인 multi-scale 및 long-distance message passing을 가능하게 합니다. DiffusionNet 구현에 대한 이러한 향상 기능은 training process의 instabilities를 완화하며(Supplementary Figure 2에서 볼 수 있듯) 제공된 code repository와 pip 패키지로 사용 가능합니다.
3.3 HYBRID REPRESENTATION LEARNING
앞서 언급했듯이, protein learning을 위한 다른 representations와 비교하여 surface-based learning의 효능을 평가하는 것을 넘어, 우리는 서로 다른 representations를 unified framework에 통합하여 각각의 고유한 강점을 활용하는 이점을 탐구합니다. 직관적으로 surface representations는 단백질 상호작용과 관련된 tasks에 중요한 복잡한 기하학적 세부사항을 포착할 수 있는 반면, graph representations는 표면 역학 및 상호작용 능력에 간접적으로 영향을 미치는 단백질 내부의 특정한 원자적 상호작용을 상세히 나타냅니다. 나아가 이러한 representations는 learning에 대한 보완적인 접근 방식을 촉진합니다: 그래프를 통한 local message passing과 학습된 diffusion을 통한 표면의 global information dissemination. 이러한 고려 사항에서 영감을 얻어, 우리는 graph와 surface-based representations 간의 feature sharing을 가능하게 하는 method를 제안합니다. 아래에서 강조되듯이, 이전의 관련 approaches와 달리 우리는 network의 모든 학습된 layers에 걸쳐 두 representations 간의 통신을 가능하게 합니다.
우리의 hybrid representation의 기초로서 surface $S_P$에 추가하여 graph representation $G_P = (V_g, E_g)$를 구성합니다. 우리는 Atom3d benchmark 내에서 사용되는 것과 일치하는, 노드가 원자인 그래프를 사용하거나 residue 수준에서 정의된 그래프를 사용합니다. residue-level graph는 node features로 사용되는 ESM-650M sequence embeddings로 풍부해집니다.
hybrid approach를 구성하기 위해 우리는 이분 그래프(bipartite graph) $G = (V, E)$를 구축하는데, 여기서 $V = V_g \cup V_s$는 각각 graph nodes와 surface vertices를 나타냅니다. surface의 각 vertex에 대해 우리는 그래프에서 가장 가까운 16개의 이웃(nearest neighbors)을 찾고 이분 그래프에 해당하는 bidirectional edges를 추가합니다. 우리는 Section C.1에서 atomic, residue 및 bipartite graphs의 구성과 features에 대한 보다 자세한 설명을 제공합니다.
이제 $S_P$, $G_P$ 및 $G$를 사용하여 단백질을 인코딩하는 block operations를 정의합니다. surfaces와 graphs의 encoders를 각각 $s_\theta$와 $g_\theta$로 표시하고, input features의 집합을 $X = {x_n, n \in V}$로 표시합니다. 상응하는 encoded features는 $H = {h_n, n \in V}$이며, nodes $n \in V_s$에 대해서는 $h_n = s_\theta(x_n)$이고 nodes $n \in V_g$에 대해서는 $h_n = g_\theta(x_n)$입니다. 우리의 일반적인 methodology는 이분 그래프 $G$ 위에 message-passing neural networks인 $MP_\theta$를 통합하여, layer $l$에서 $X^{l+1} = MP_\theta^l(H^l)$를 얻도록 합니다. 구별되는 세트인 $\theta_{sg}$와 $\theta_{gs}$를 채택함으로써, architecture는 표면에서 그래프로, 그리고 그 반대로 이동하는 메시지를 처리합니다. 이러한 block operations는 Figure 1에 표시된 것처럼 stack될 수 있습니다. 우리는 우리의 feature sharing이 local(node) 수준에서 발생하며 3D 공간의 proximity relations에 의해 가능해짐을 강조합니다. 또한 우리의 hybrid approach는 모든 network layers 전반에 걸쳐 정보 공유를 가능하게 하면서 두 representations를 jointly하게 훈련하며, 이는 성공에 있어 매우 중요합니다.
3.4 PROPOSED ARCHITECTURES
우리의 framework는 diffusion operation 뒤에 지정된 너비의 두 개의 hidden layers를 가진 pointwise neural network로 구성된 surface encoding blocks인 $s_\theta$를 통합합니다. 우리가 제안하는 첫 번째 network인 Surface Diff는 오직 이러한 surface blocks에만 기반합니다. 이는 분리되어 사용될 때 surface representation의 관련성을 평가하는 데 사용됩니다. Surface Diff는 DiffusionNet에 기반하지만 Section 3.2에서 언급된 우리의 수정 사항을 통합한다는 점에 유의하십시오. hybrid approach를 사용하는 다른 모든 methods에 대해, 두 encoding blocks $s_\theta$와 $g_\theta$의 너비는 일관되게 동일한 값으로 설정되었습니다.
우리는 Atom3d benchmark 프로토콜을 따름으로써 공정한 방식으로 representations를 비교하는 것을 목표로 하는 모델인 AtomSurf-bench를 소개합니다. 그것의 graph encoder $g_\theta$는 Graph Convolutional Networks (GCN)로 구성되며 Batch Normalization operations와 얽혀있고, 이분 그래프 위에서의 message-passing은 Graph Attention Layer입니다. Atom3d benchmark 표준을 따라, AtomSurf-bench는 200k의 learnable parameters를 가지며 surface input features를 사용하지 않고 atomic-level graph representation에서 atom type만을 고려합니다.
게다가 최근 제안된 ProNet encoder와 함께 앞서 언급한 residue graph를 활용하는 AtomSurf를 소개합니다. residue graphs에서 작동함에도 불구하고, ProNet은 완전성(completeness)을 허용하며(ComeNet에서 소개된 바와 같이) 형태를 고유하고 동변적으로(equivariantly) 식별하는 상대적인 local coordinate systems를 기반으로 각 residue를 구성하는 원자의 기하학적 형태에 대한 featurization을 추가합니다. 그런 다음 이러한 features는 구면 조화 함수(spherical harmonics)를 사용하여 임베드되고, ProNet은 learning을 수행하기 위해 GraphConv에서 소개된 message passing에 의존합니다. 또한 우리는 앞서 언급한 이분 그래프 features와 GVP encoder를 사용하여 우리의 message passing을 수행합니다.
몇 가지 기존 approaches를 특수한 경우로 포함하여, 다양한 bipartite message passing networks와 조직이 가능합니다. 우리는 Section C에 구현 세부 정보를 제공합니다. Section 4.5에서 이러한 구성들을 평가하고 최종 모델에 대한 몇 가지 ablation studies를 수행합니다.
3.5 COMPUTATIONAL ENHANCEMENTS
Surface-methods는 전통적으로 compute-expensive methods로 여겨져 왔으며, 이는 그들의 내재적 복잡성을 회피(side-step)하려는 approaches의 동기가 되었습니다. Section D.1에 제시된 complexity analysis를 통해, 우리는 vertex의 수가 DiffusionNet runtime에 있어 중요하다는 것을 발견했고, meshes를 coarsening함으로써 이를 해결했습니다. 이 coarse surface regimen에서 graph encoding with ProNet은 computational bottleneck이 되는데, 이는 분리되어 사용될 때보다 더 작습니다. 따라서 우리는 우리의 method가 coarse surface regimen에서 이 효율적인 graph encoder보다 더 빠르다고 주장합니다.
그러나 저장 공간과 메모리 관점에서 볼 때, 특히 확률성(stochasticity) 때문에 surface-related operations의 메모리 발자국이 graphs에서 기인하는 것들을 지배합니다. 우리는 이것이 I/O 측면보다는 batch size 측면에서 제한적이라는 것을 발견했습니다. 우리는 메모리 문제를 완화하는 dynamic batching 절차를 구현했습니다. 추가 세부 정보는 Section D.2에 제공됩니다.
surface networks의 메모리 발자국을 줄일 방법을 찾는 것은 여전히 중요한 방향으로 남아 있습니다
3 METHODS 핵심 정리 노트 (AI 연구자용)
이 논문의 방법론 핵심은 1) SOTA 표면 모델(DiffusionNet)의 생물학적 스케일 문제 해결, 그리고 2) 모든 layer에 걸친 Surface-Graph 간의 Deep Fusion으로 요약할 수 있습니다.
1. Surface Encoder 최적화: 단백질 스케일(Scale) 문제 해결
- 기반 모델: 열방정식(Heat equation, $\frac{\partial f}{\partial t} = \Delta_S f$)을 기반으로 multi-scale geometric feature를 추출하는 DiffusionNet 채택.
- 한계 극복 (Scale Independence 문제): 기존 DiffusionNet은 모든 shape의 표면적을 균일하게 normalize하여 학습을 안정화시켰으나, 단백질은 리간드와의 상대적 크기(물리적 스케일)가 매우 중요함. 이를 normalize하면 핵심 정보가 소실됨.
- 해결책 (Biological Prior 주입): * Scale normalization을 제거하고 Batch Normalization을 추가하여 학습 안정성 확보.
- 단백질 결합 부위의 실제 물리적 크기(~10 Å)에 맞추어, 확산 시간(diffusion times)의 초깃값을 $t \sim \mathcal{N}(10, 5)$의 정규분포에서 샘플링. 이를 통해 네트워크가 생물학적으로 유의미한 receptive field를 갖도록 강제함.
2. Hybrid Representation Learning (Deep Integration)
- Bipartite Graph 구성: Surface vertices($V_s$)와 Graph nodes($V_g$)를 결합한 이분 그래프(bipartite graph) 구축. 표면의 각 vertex마다 3D 공간상 가장 가까운 16개의 graph node 이웃을 찾아 bidirectional edge로 연결.
- 모든 Layer에서의 Message Passing: $X^{l+1} = \text{MP}_{\theta}^{l}(H^l)$ 구조를 통해, 단순히 입출력단에서 합치는 것(early/late fusion)이 아니라 네트워크의 모든 중간 layer에서 표면과 그래프 간의 node-wise feature sharing을 수행하여 시너지 극대화.
3. 제안하는 주요 Architectures
- Surface_Diff: 순수 surface-based baseline (수정된 DiffusionNet 적용).
- AtomSurf-bench: Atom3D benchmark의 엄격한 제약(200k 파라미터 제한, atom type 외 추가 feature 사용 불가)을 맞추기 위한 모델. GCN과 Graph Attention Layer를 활용.
- AtomSurf (최종 모델): Residue graph 기반. ESM-650M sequence embedding을 node feature로 사용하고, 기하학적 3D 구조를 캡처하는 ProNet을 graph encoder로, GVP를 이분 그래프의 message passing encoder로 채택하여 성능을 극대화함.
4. 연산 최적화 (Computational Enhancements)
- 표면 연산의 bottleneck인 vertex 수를 줄이기 위해 mesh coarsening 기법 적용.
- Surface network의 높은 stochasticity로 인한 메모리 부족 문제를 해결하기 위해 Dynamic batching 도입.
쉬운 설명 :
이 섹션은 **"어떤 AI 기술을 어떻게 조립하고 개조해서 단백질 분석의 끝판왕을 만들었는가?"**에 대한 설계도입니다. 핵심은 크게 두 가지 개조 작업입니다.
첫 번째 개조: 표면 분석 AI에게 '진짜 크기' 알려주기
연구진은 3D 그래픽 분야에서 아주 성능이 좋은 'DiffusionNet'이라는 모델을 가져왔습니다. 이 모델은 물체 표면에 열이 퍼져나가는 방식을 계산해서 형태를 파악합니다. 그런데 기존의 3D AI들은 쥐든 코끼리든 다 똑같은 크기로 화면에 꽉 차게 늘리거나 줄여서(정규화) 학습하는 버릇이 있습니다. 그래야 AI가 덜 헷갈리니까요.
하지만 단백질과 약물이 결합하려면 **'실제 물리적인 크기(Scale)'**가 아귀에 딱 맞아야 합니다. 그래서 연구진은 크기를 강제로 맞추는 기능을 빼버렸습니다. 대신, 결합 부위의 실제 물리적 크기(약 10 옹스트롬)에 딱 맞게 열이 퍼져나가도록 AI의 '시간 세팅'을 조작(정규분포에서 샘플링)해 주었습니다. 생물학적 지식을 AI에 직접 주입한 것이죠.
두 번째 개조: 겉(표면)과 속(그래프)을 잇는 수많은 다리 놓기
단백질의 '표면'과 내부 원자들의 '그래프'를 함께 학습시키기 위해, 공간상에서 서로 가까이 있는 표면의 점과 내부의 점을 찾아 '다리(Bipartite graph)'를 무려 16개씩 연결했습니다.
그리고 학습을 다 끝내고 결과를 합치는 게 아니라, AI가 생각을 거치는 모든 단계(Layer)마다 이 다리를 통해 서로의 정보를 실시간으로 주고받게 만들었습니다. 표면 AI가 "여기가 좀 튀어나왔어" 하면 그래프 AI가 "아, 내부에 이 원자가 밀어내고 있네" 하고 바로바로 소통하며 학습하는 완벽한 하이브리드 팀플을 구현한 것입니다.
마지막으로, Atom3D 시험용 규격을 맞춘 가벼운 모델과, 현존하는 강력한 언어 모델(ESM) 지식까지 전부 끌어다 쓴 풀스펙 모델(AtomSurf)을 각각 만들고, 연산이 너무 무거워지지 않게 덜 중요한 디테일을 쳐내는 최적화 작업까지 마쳤다는 내용입니다.
.