곡률로 GNN 지식을 MLP에 증류해 그래프 없이 추론한다
Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs
무엇인가
GNN은 메시지 패싱에 의존하기 때문에 사기 탐지처럼 지연에 민감한 응용에서 추론이 느리다. GNN-to-MLP 증류는 메시지 패싱 교사의 예측 정확도를 유지하면서 추론 시점에는 그래프가 필요 없는 MLP를 배포하는 방법이다. 기존 방법들은 주로 노드별 예측을 전이하거나 교사 신뢰도로 재가중할 뿐, 학생이 교사의 그래프 유도 기하를 어디에서 보존해야 하는지는 명시하지 않는다. 이 논문은 그 누락이 학생 표현 공간에 두 가지 스펙트럼 실패 모드를 만든다고 주장한다. 희소 그래프에서는 경계 영역 근처에 집중된 고에너지 교사 방향을 놓치는 스펙트럼 과소적합(spectral underfit)이, 밀집 그래프에서는 교사가 aggregation으로 이미 붕괴시킨 잉여 방향을 그대로 유지하는 스펙트럼 과대적합(spectral overfit)이 나타난다.
어떻게 동작하나
분석은 교사와 학생의 커널 불일치에서 출발한다. K층 GNN의 GNTK는 이웃 쌍에 대한 aggregation 항을 포함해 (X, G)에 함께 의존하는 반면, MLP의 NTK는 aggregation이 없는 대응물이라 X에만 의존한다. 그 차이 Δ(K) = K(K) − K_MLP가 메시지 패싱이 만든 그래프 의존 성분이다. 저자들은 교사 고유기저에서 ρ_i = (u_i^T K_MLP u_i) / λ_i^K를 정의해, ρ_i < 1이면 과소적합, ρ_i > 1이면 과대적합으로 부르고 대각 불일치 D = Σ (ρ_i − 1)² λ_i^K로 요약한다. 이상적 목표는 에너지 가중 정렬 손실 L* = Σ λ_i^K (f̂_i^T − f̂_i^S)²인데, 전체 교사 커널을 저장하고 대각화해야 해서 직접 쓸 수 없다. 대신 올리비에-리치 곡률(ORC) κ(i,j) = 1 − W₁(μ_i, μ_j)을 국소 프록시로 쓴다. 낮은 곡률 엣지는 랜덤워크 이웃이 서로 어긋나는 경계 영역을, 높은 곡률 엣지는 aggregation이 표현을 강하게 평활화한 밀집 내부 영역을 표시한다.
무엇과 다른가
제안 방법 G²MLP는 L*를 세 단계로 완화한다. 첫째, 전역 커널 정렬을 국소 이웃 분포 매칭으로 바꾼다. 각 노드 v에 대해 lazy random walk 분포 μ_v를 정의하고, 은닉 특징과 출력 로짓을 이어붙인 표현 r_u에 대한 국소 측도 ν_v^T, ν_v^S를 Wasserstein-2 거리로 맞춘다. 실무에서는 계산 단순화를 위해 균일 1-hop 프록시 q_v를 쓴다. 둘째, 이산 최적수송을 대각 가우시안 W₂ 닫힌 형태로 대체한다. 이웃 측도의 1·2차 모멘트만 맞춰 평균 차이 항과 표준편차 차이 항의 합으로 거리를 계산해, 노드당 비용을 O(d)로 줄인다. 셋째, 균일 노드 가중을 곡률 의존 가중으로 바꾼다. α_v = exp(−γ_z κ̃_v)/정규화 항은 낮은 곡률 경계 노드에서 예측 수준 정렬을 강화하고, β_v = exp(+γ_h κ̃_v)/정규화 항은 높은 곡률 내부 노드에서 표현 수준 이웃 정렬을 강화한다. 정규화는 전체 손실 스케일을 바꾸지 않고 강조만 재분배한다. 최종 손실은 L = λ_CE·L_CE + λ_KD·L_KD + λ_CAW·L_CAW/(sg(L_CAW)+ε_sg)이며, 노드 분류와 Graph Transformer 실험에서는 λ_CE = 0으로 둔다. 반복당 비용은 O((|V|+|E|)·d)이고, 곡률 재가중은 전부 학습 시점에만 적용되므로 배포 모델은 평범한 MLP 그대로다.
어떻게 쓰나
노드 분류 실험은 Cora, Citeseer, Pubmed, A-Computer, A-Photo, ogbn-Arxiv 여섯 개 벤치마크에서 진행했다. 교사는 GraphSAGE를 고정하고 학생은 원본 노드 특징만 먹는 MLP이며, 베이스라인은 vanilla MLP, GLNN, KRD, FF-G2M이다. Transductive 설정에서 G²MLP는 여섯 데이터셋 모두 그래프 없는 최고 정확도를 기록해 최강 베이스라인 대비 0.16~1.66pp 개선했고, 여섯 중 다섯에서 GraphSAGE 교사까지 넘어섰다(Citeseer 최대 +4.02pp). 유일하게 ogbn-Arxiv에서만 교사에 −6.16pp 못 미쳤는데, 저자들은 이 대규모 그래프에서 특징만 쓰는 추론의 알려진 어려움으로 설명한다. Inductive·production 설정에서는 5/6 데이터셋에서 최고를 기록해 0.16~1.60pp 개선했고, 특히 inductive에서 Cora·Citeseer·Pubmed·A-computer에서 2.07~2.69pp가 올라 transductive보다 이득이 컸다. A-computer(inductive)에서 82.67%로 그래프 접근 없이 교사(82.83%)에 사실상 도달했고, Citeseer·Pubmed(production)에서는 교사를 각각 3.61pp, 2.68pp 능가했다.
전제와 한계
Graph Transformer 교사로의 전이도 확인했다. GT, GraphGPS, NAGphormer 세 교사에 대해 MLP 학생이 15개 셀 중 13개에서 교사와 같거나 앞섰고(NAGphormer 5/5, GT 5/5, GraphGPS 3/5), 15/15 셀 전부에서 vanilla-KL GLNN 베이스라인을 앞섰다. 링크 예측에서는 세 데이터셋 모두 AUC와 AP에서 승리했고, Cora에서 KRD 대비 +4.70 AUC, +5.47 AP, Citeseer +0.30/+0.51, Pubmed +0.04/+0.58이었다. 세 데이터셋 모두 SAGE 교사도 넘었다(Cora +2.61, Citeseer +2.41, Pubmed +2.00 AUC). 추론 효율은 3층 GraphSAGE 교사 대비 58배 빨랐고(8.6ms vs 499.7ms), vanilla MLP 대비 교사 이득의 59.1%를 회복해 같은 학생 구조의 GLNN(53.7%)을 앞서고 더 깊은 학생을 쓰는 TINED(58.2%)와 비슷한 수준이다.
실무 관점에서 이 논문이 주는 것은 학습 시점에만 그래프를 쓰고 배포는 피처 벡터→로짓 한 번으로 끝내는 구체적 절차다. 이웃 샘플링이나 그래프 순회가 사라지므로 지연 예산이 빡빡한 서비스에서 GNN 서빙 스택을 통째로 대체할 후보가 된다. 다만 도입 전에 확인할 것이 있다. 곡률은 그래프당 한 번 전처리하면 모든 학습 런과 배포에 걸쳐 상각되지만, 정확한 ORC는 엣지마다 최적수송 문제를 풀어야 해서 그래프가 아주 크면 병목이 된다. 또 대각 가우시안 근사는 이웃 표현 분포가 한 덩어리로 모여 있을 때 가장 믿을 만하므로, 이종성이 강한 그래프에서는 근사 품질을 따로 검증해야 한다. 마지막으로 ogbn-Arxiv처럼 특징만으로는 부족한 대규모 그래프에서는 학생이 교사를 따라가지 못한다는 점을 감안해야 한다.
저자들이 밝힌 한계는 명확하다. L*는 정확한 등가물이 아니라 실용적 근사이며, 세 완화 단계(국소 이웃 매칭, 대각 가우시안 W₂, 곡률 가중)는 의도적으로 계산 가능성을 택한 근사다. 대각 공분산 근사는 다봉(multimodal)이거나 이종성이 강한 이웃에서 정확도가 떨어질 수 있다고 명시했다. 곡률 사전계산은 정확한 ORC 기준 O(|E|·d̄³)의 최적수송 문제를 요구해 중소 그래프에서는 가능하지만 초대규모 그래프로는 확장이 나쁘며, 대안으로 Sinkhorn 정규화(엣지당 O(d̄²))나 조합적 O(1) 대리인 Forman-Ricci 곡률을 제시한다. 또한 위치 인코딩을 학생 입력에 추가하는 NOSMOG나 코드북을 함께 학습하는 VQGraph 같은 변형은 셋업 자체가 달라 공정 비교 대상에서 제외했다.