Generalizing GNNs with Tokenized Mixture of Experts

기술 요약: 토큰화된 전문가 혼합(Mixture of Experts)을 통한 GNN의 일반화

1. 문제 정의: 동결된 배포(Frozen Deployment)의 "불가능한 삼각형"

본 논문은 그래프 신경망(GNN) 배포 시 발생하는 결정적인 한계를 다룹니다. 프로덕션 환경에서 GNN은 일반적으로 릴리스 사이에 파라미터 업데이트 없이 "동결된 스냅샷(frozen snapshots)" 형태로 배포됩니다. 이 동결된 기간 동안, 단일 고정 모델은 다음 세 가지 상충하는 목표를 동시에 만족해야 하며, 이를 **"불가능한 삼각형(impossible triangle)"**이라 명명합니다:

1. 적합도(Fit): 깨끗한 인-디스트리ب션(in-distribution) 데이터에 대한 높은 성능.
2. 일반화(Generalization): 분포 변화(예: 노드 차수 또는 호모필리 변화)에 대한 강건성.
3. 안정성(Stability): 입력 섭동(예: 피처 마스킹 또는 엣지 삭제)에 대한 저항성.

저자들은 표준 GNN이 이러한 목표들의 균형을 맞추는 데 실패한다고 주장하는데, 이는 모델이 **단일 고정 추론 규칙(static inference)**을 채택하기 때문입니다. 이론적으로, 안정성(섭동에 대한 불감성)을 강제하는 것은 종종 고주파 신호를 억제하는 것을 요구하며, 이러한 신호는 태스크 관련성이 높은 경우가 많아 적합도와 다양한 환경으로의 일반화를 저해합니다. 반대로, 적합도와 일반화에 최적화하는 것은 모델을 노이즈에 민감하게 만듭니다.

인스턴스 조건부 계산(Instance-Conditional Computation, ICC)—서로 다른 입력을 서로 다른 메커니즘을 통하도록 라우팅하는 방식—은 가용 가능한 계산 경로의 집합을 확장함으로써 잠재적인 해결책을 제공하지만, 기존의 ICC 방식은 새로운 취약성을 도입합니다. 즉, 라우팅 결정이 분포 변화에 따라 표류(drift)할 수 있으며, 라우팅의 변동이 다운스트림 레이어의 섭동 민감도를 증폭시킬 수 있습니다.

2. 방법론: STEM-GNN

커버리지(일반화)와 안정성 사이의 긴장을 해결하기 위해, 저자들은 견고한 ICC를 구현하기 위해 긴밀하게 결합된 세 가지 구성 요소를 갖춘 사전 학습 후 미세 조정 프레임워크인 STEM-GNN(Stable TokEnized Mixture-of-Experts GNN)을 제안합니다.

A. 커버리지 확장을 위한 MoE 인코더
단일 메시지 패싱 규칙 대신, STEM-GNN은 전문가 혼합(Mixture-of-Experts, MoE) 인코더를 사용합니다.

• 메커니즘: 각 노드는 경량 라우터를 통해 공유된 전문가 연산자의 입력 의존적 조합을 통해 라우팅됩니다.
• 목표: 이를 통해 단일 고정 파라미터 세트가 다양한 유효 메커니즘의 집합을 실현할 수 있게 하여, 이질적인 테스트 조건(예: 다양한 차수 또는 호모필리 수준을 가진 그래프)에 대한 커버리지를 개선합니다.

B. 표현 안정화를 위한 VQ 토큰화
라우팅 결정과 표현 표류(representation drift)가 예측 헤드로 전파되어 발생하는 불안정성을 방지하기 위해, 저자들은 벡터 양자화(Vector-Quantized, VQ) 토큰 인터페이스를 도입합니다.

• 메커니즘: 연속적인 인코더 출력은 예측 헤드로 들어가기 전, 최근접 이웃 매칭을 통해 고정된 코드북으로 이산화됩니다.
• 목표: 양자화 경계를 넘지 않는 인코더 출력의 작은 섭동(입력 노이즈나 라우팅 변동으로 인한)은 예측 헤드의 입력에 변화를 주지 않습니다(zero change). 이는 표현 표류를 흡수하는 안정적인 인터페이스를 생성합니다.

C. 민감도 제어를 위한 리프시츠(Lipschitz) 정규화 예측 헤드
남아있는 변동(예: 토큰 전환이 발생하는 경우)의 영향을 제한하기 위해 예측 헤드가 정규화됩니다.

• 메커ㅿ즘: 선형 예측 헤드의 가중치에 프로베니우스 노름(Frobenius norm) 페널티를 적용하여 리프시츠 상수를 제한합니다.
• 목표: 이는 잔여 입력 변동(이산적 토큰 전환 포함)이 최종 출력으로 증폭되는 것을 제한하여, 모델이 업스트림의 변화에 과도하게 반응하지 않도록 보장합니다.

3. 이론적 기여

본 논문은 **정적 추론(Static Inference)**과 **인스턴스 조건부 계산(ICC)**을 대조하는 공식적인 이론적 분석을 제공합니다.

• 정적 추론의 긴장: 저자들은 리프시츠 안정성 제약 하에서 정적 추론이 구조적인 "증인 텐션(witness tension)"에 직면함을 증명합니다. 안정성을 유지하려면 모델은 고주파 성분에 대한 의존도를 제한해야 하며, 이는 최악의 환경 리스크에 대한 비제로 하한선(floor)을 유발합니다.
• ICC 리스크 분해: ICC의 경우, 저자들은 리스크를 두 가지 별개의 레버로 분해합니다:
  1. 커버리지 vs 선택(Coverage vs. Selection): OOD 리스크는 가용 메커니즘이 얼마나 다양한 조건을 커버하는지와 라우터가 그중 얼마나 정확하게 선택하는지로 구분됩니다.
  2. 민감도 vs 표류(Sensitivity vs. Drift): 안정성 리스크는 고정된 경로의 민감도와 라우팅 표류에 의해 발생하는 증폭으로 구분됩니다.
     이 분석은 STEM-GNN의 설계를 정당화합니다: MoE는 커버리지를 해결하고, VQ는 표류/안정성을 해결하며, 리프시츠 정규화는 민감도 증폭을 해결합니다.

4. 실험 결과

저자들은 노드 분류, 링크 예측, 그래프 분류를 아우르는 8개의 벤치마크(Cora, PubMed, Wiki-CS, WN18RR, FB15K-237, HIV, PCBA 포함)에서 STEM-GNN을 평가했습니다.

• 클린 성능(Clean Performance): STEM-GNN은 모든 벤치마크에서 평균 정확도 80.79%를 달로 달성하며, GFT, GIANT, GraphMAE와 같은 강력한 베이스라인들을 능가했습니다. 또한 다른 방법론들에 비해 랜덤 시드에 대한 민감도가 낮음을 보여주었습니다.
• 섭동 강건성(Perturbation Robustness): 추론 시점의 섭동(피처 마스킹 및 엣지 삭제) 상황에서, STEM-GNN은 베이스라인들보다 더 완만하게 성능이 저하되며, 노이즈 심도가 높아짐에 따라 더 높은 정확도를 유지합니다.
• OOD 일반화: 차수 변화(degree shifts) 및 호모필리 변화(homophily shifts) 하에서, STEM-GNN은 특히 도전적인 고차수 또는 고호모필리 시나리오에서 OOD 버킷에 대해 최상의 평균 성능을 달성했습니다.
• 삼중 목적 달성(Tri-Objective Balance): Fit, OOD-worst, Perturb-mean에 대한 통합 평가에서, STEM-GNN은 모든 노드 벤치마크에서 최상의 평균 점수를 기록했습니다. 결정적으로, 하나를 위해 다른 하나를 희생하는 것이 아니라, 세 축 모두를 베이스라인 대비 동시에 개선했습니다.
• 절제 연구(Ablation Studies): 세 가지 구성 요소(MoE, VQ, 또는 리프시츠 정규화) 중 하나라도 제거하면 삼중 목적 평균이 하락하며, 이는 각 구성 요소가 중복되는 것이 아니라 상호 보완적임을 확인시켜 줍니다.

5. 의의 및 주장

본 논문은 STEM-GNN이 동결된 배포 환경에서의 강건한 그래프 일반화를, 단순히 훈련 목적 함수뿐만 아니라 테스트 시점의 계산 메커니즘에 의해 지배되는 삼중 목적 텐션으로 성공적으로 재정의했다고 주장합니다.

주요 핵심 주장은 다음과 같습니다:

• 통합 프레임워크: 그래프 학습에서 메커니즘 커버리지 확장과 명시적인 배포 시점 안정성 제어(토큰화 및 리프시츠 제약 활용)를 결로 결합한 최초의 방법론입니다.
• 이론적 통찰: 정적 추론이 본질적으로 안정성-일반화 트레이드오프에 직면해 있는 반면, ICC는 이 긴장을 완화할 수 있는 새로운 설계 레버를 제공한다는 공식적인 증명을 제공합니다.
• 실질적 효용성: 본 프레임워크는 강력한 클린 성능을 유지하면서도 분포 변화와 섭동에 대한 강건성 사이의 균형을 크게 개선하여, 테스트 시점에 모델을 업데이트할 수 없는 실제 세계의 동결된 배포 시나리오에 적합합니다.

저자들은 이 접근 방식이 리프시츠 안정성 제약과 유한한 라우팅 메커니즘 집합에 의존하고 있음을 언급하며, 이러한 가정을 완화하는 것이 향후 연구 방향임을 밝히고 있습니다. 재현성을 돕기 위해 코드와 데이터는 공개되어 있습니다.