다양한 도메인에서 전이 가능한 그래프 표현을 학습하는 범용 그래프 트랜스포머 GraphFM

초록

GraphFM은 Perceiver 기반 인코더와 학습 가능한 잠재 토큰을 활용해 서로 다른 크기와 토폴로지를 가진 그래프들을 하나의 공유 잠재 공간으로 압축한다. 152개의 실·합성 그래프(노드 740만, 엣지 1억 9천만)를 사전학습한 뒤, 고정된 라이트웨이트 MLP 혹은 노드 디코더를 통해 다양한 노드 분류·회귀 태스크에 빠르게 적응한다. 대규모·다양한 데이터와 모델 규모 확대가 미보인 그래프에 대한 일반화 성능을 꾸준히 향상시키며, 기존 GNN·그래프 트랜스포머 대비 하이퍼파라미터 민감도가 낮다.

상세 분석

GraphFM은 그래프 데이터를 시퀀스 토큰으로 변환한 뒤, 고정된 K(=512)개의 학습 가능한 잠재 토큰 Z₀에 교차‑어텐션을 적용한다. 이 과정에서 각 그래프의 노드 임베딩 Xᵍ는 Q·K·V 선형 변환을 거쳐 Z₀와 상호작용하고, 이후 L개의 자체‑어텐션 블록이 잠재 공간에서 연산을 수행한다. 복잡도는 O(K·Nᵍ + L·K²)로, Nᵍ≫K인 대규모 그래프에서도 효율적이다. 잠재 토큰은 모든 도메인에 걸쳐 공유되므로, 서로 다른 그래프 구조와 피처 스키마가 동일한 “어휘”를 통해 교차 전이될 수 있다.

위치 인코딩은 라플라시안 고유벡터를 SignNet으로 정규화해 부호·기저 변환에 불변하도록 설계했으며, 이는 그래프마다 고유한 순서가 없다는 문제를 해결한다. 이렇게 정제된 위치 정보는 잠재 토큰과 결합돼 전역 구조를 압축한다.

노드 디코더는 (1) 자기 토큰, (2) 이웃 토큰(랜덤 워크 샘플링), (3) 잠재 토큰을 순서대로 연결한 시퀀스 Sᵢᵍ를 만든 뒤, 얕은 트랜스포머(M)로 처리한다. 복잡도는 O(Nᵍ·M·(K+T+1)²)이며, T는 샘플링 이웃 수다. 이 설계는 전역 정보를 잠재 토큰에, 지역 정보를 이웃 토큰에 각각 맡겨 효율적인 정보 흐름을 만든다.

스케일링 측면에서 저자는 두 가지 핵심 기법을 제시한다. 첫째, 다양한 크기의 그래프를 패딩 없이 하나의 토큰 시퀀스로 연결하고 FlashAttention을 이용해 가변 길이 시퀀스를 효율적으로 처리한다. 둘째, 그래프 크기에 따라 “뱀‑형”으로 GPU에 할당하는 DistributedSSSampler를 도입해 대규모 분산 학습 시 메모리 불균형을 최소화한다. 이로써 8 GPU 환경에서도 100%에 가까운 메모리 활용률을 달성했다.

실험에서는 모델 파라미터를 389K에서 75M까지, 사전학습 토큰을 20만에서 7.3M까지 확대했을 때, 미보인 그래프에 대한 정확도가 평균 2.1%p 상승함을 확인했다. 특히 합성 그래프(저동질성)와 생물학적 그래프를 포함시켰을 때, 인용 네트워크와 같은 전통적 동질성 그래프에서도 성능이 크게 개선되었다. 파인‑튜닝은 두 가지 방식으로 수행했는데, (a) 라이트웨이트 MLP만 학습하는 MFT는 10~20 스텝 내에 수렴하며, (b) 전체 노드 디코더를 미세조정하는 NFT는 기존 최첨단 그래프 트랜스포머와 동등하거나 우수한 결과를 보였다.

민감도 분석 결과, GraphFM은 학습률, 가중치 감쇠, 레이어 수 등에 대해 비교적 평탄한 성능 곡선을 보이며, GCN이나 NAGphormer와 달리 하이퍼파라미터 튜닝 비용이 크게 감소한다. 전체적으로, 이 논문은 “범용 그래프 전이 학습”이라는 새로운 패러다임을 제시하며, 다양한 도메인에 걸친 대규모 사전학습이 그래프 인공지능의 확장성을 크게 향상시킬 수 있음을 실증한다.