다중스케일 그래프 토크나이저와 작업 적응 라우팅

초록

본 논문은 그래프 기반 파운데이션 모델을 위한 계층적 양자화 토크나이저를 제안한다. 노드 임베딩과 퍼스널라이즈드 페이지랭크(PPR) 확장 특징을 각각 잔차 벡터 양자화(RVQ)로 다중 레벨 토큰화하고, 작업별 라우팅 모듈이 깊이별 가중치를 학습한다. 이후 이중 교차‑어텐션으로 로컬 토큰과 전역 토큰을 정렬·융합해 하나의 토큰 시퀀스로 만든 뒤, 기존 백본 인코더를 그대로 사용해 노드 분류와 링크 예측에서 기존 양자화 기반 방법보다 일관된 성능 향상을 보인다.

상세 분석

이 연구는 그래프 데이터의 다중스케일 특성을 토큰 수준에서 포착하려는 시도다. 기존 RVQ 기반 토크나이저는 여러 레벨의 코드북을 만들지만, 최종 임베딩을 만들 때는 단순히 모든 레벨을 동등하게 결합하거나 고정된 규칙을 적용한다. 이는 특정 작업(예: 노드 분류)에서는 세밀한 잔차 정보가, 다른 작업(예: 링크 예측)에서는 거친 구조 정보가 더 중요할 수 있다는 점을 반영하지 못한다. 논문은 이를 해결하기 위해 두 가지 핵심 메커니즘을 도입한다. 첫째, 작업‑적응 양자화 라우팅(TAQR) 은 각 RVQ 레벨에 스칼라 가중치를 할당한다. 레벨별 토큰을 평균 풀링해 요약 벡터를 만든 뒤, 두‑계층 MLP와 온도‑조절 소프트맥스를 통해 라우팅 확률을 계산한다. 이렇게 얻은 가중치는 학습 중에 자동으로 조정되어, 작업에 가장 유리한 스케일을 강조한다. 둘째, 이중 교차‑어텐션 은 로컬 토큰(원본 GCN 임베딩 기반)과 전역 토큰(PPR‑확장 임베딩 기반) 사이에 상호‑주의적 어텐션을 수행한다. 각각을 쿼리·키·밸류로 변환해 서로를 컨텍스트로 삼음으로써, 로컬 의미와 장거리 연결성을 동시에 보존한다. 이 과정은 게이트 메커니즘과 결합돼, 최종 토큰 시퀀스가 작업‑특화 가중치를 반영하도록 한다.

기술적으로는 먼저 GCN(또는 샘플링 기반 GCN)으로 노드 임베딩 H를 얻고, PPR 매트릭스 P를 이용해 H에 K 단계 확산을 적용해 H_PPR을 만든다. 두 임베딩 집합에 대해 동일한 M‑레벨 RVQ를 수행해 각각 C_node와 C_PPR을 생성한다. 이후 TAQR이 각 레벨 m에 대해 w(m) 를 계산하고, 가중합 C(m) = w(m)·(C_node(m)+C_PPR(m)) 형태로 토큰을 합산한다. 마지막으로 이중 교차‑어텐션을 적용해 정렬된 토큰을 얻고, 이를 그대로 Transformer‑기반 혹은 GNN‑기반 파운데이션 모델에 입력한다.

실험에서는 Cora, PubMed, ogbn‑arxiv 등 노드 분류 벤치마크와 ogbl‑collab, ogbl‑citation 등 링크 예측 벤치마크에서 동일한 컴퓨팅 비용 하에 기존 RVQ·GQT·VQ‑GNN 등 강력한 베이스라인을 능가한다. 특히, 라우팅 가중치가 작업에 따라 어떻게 변하는지 시각화한 결과는 TAQR이 실제로 코스톤(거친) 레벨을 강조하거나 세밀한 잔차 레벨을 강조하는 것을 확인시켜준다. Ablation 실험에서는 (1) 라우팅 없이 단순 평균, (2) 단일 교차‑어텐션, (3) PPR 없이 로컬 토큰만 사용했을 때 성능 저하가 나타나며, 제안된 세 요소가 모두 기여함을 입증한다.

이 논문은 그래프 토크나이저 설계에 있어 다중스케일 토큰화, 작업‑조건부 라우팅, 이중 교차‑어텐션이라는 세 축을 제시함으로써, 파운데이션 모델이 다양한 그래프 작업에 일관된 토큰 인터페이스를 공유하도록 만든다. 또한 라우팅 가중치는 해석 가능성을 제공해, 어느 스케일이 특정 작업에 유리한지 직관적으로 파악할 수 있다. 향후 연구에서는 라우팅을 메타‑학습이나 강화학습으로 확장하거나, 더 복잡한 구조(예: 하이퍼그래프)에도 적용하는 방향이 기대된다.