Graph is a Natural Regularization: Revisiting Vector Quantization for Graph Representation Learning

기술 요약: 그래프는 자연스러운 정규화 도구이다: 그래프 표현 학습을 위한 벡터 양자화의 재고

1. 문제 정의

벡터 양자화(Vector Quantization, VQ)는 비전 및 언어 영역에서 압축된 이산 표현을 학습하는 강력한 기술로 부상했으며, 그래프 구조를 토큰화함으로써 확장 가능한 그래프 파운데이션 모델(GFM)로 나아가는 경로를 제공합니다. 그러나 본 논문은 그래프 도메인에 특화된 근본적이고 아직 충분히 탐구되지 않은 과제인 **코드북 붕괴(codebook collapse)**를 식별합니다.

코드북 붕계는 입력 데이터 포인트의 대다수가 학습 가능한 코드북 내의 극소수 코드워드(codewords)에만 매핑되어, 코드북의 나머지 광범위한 부분이 활용되지 못하는 현상을 의미합니다. 저자들은 이러한 현상이 비전 및 언어 분야에서 개발된 최첨단(SOTA) 완화 전략(예: 지수 이동 평균, 코드북 리셋, 사전 학습된 인코더)을 적용하더라도 그래프 VQ에서 체계적으로 발생함을 관찰했습니다. 이러한 붕괴는 그래프 토큰의 표현력과 일반화 능력을 심각하게 제한하여, 조잡한 표현과 다운스트림 태스크 성능 저하를 초క합니다.

본 논문은 두 가지 관점에서 이러한 붕괴의 근본 원인을 진단합니다:

1. 데이터 관점: 그래프 데이터는 특징 중복성(feature redundancy) 및 **높은 연결 밀도(high connectivity density)**와 같은 고유한 특성을 가집니다. 연구 결과, 특징 중복성(PCA@95%로 측정)과 붕괴의 심각도 사이에는 양의 상관관계가 있으며, 연결 밀도(평균 노드 차수)와 코드북 활용도 사이에는 음의 상관관계가 있음이 밝혀졌습니다. 밀집된 연결성과 중복된 특징은 노드들이 동일한 토큰으로 매핑되려는 경향을 악화시킵니다.
2. 최적화 관점: 결정론적 VQ는 하드 할당(hard assignment)(최근접 이웃 선택)에 의존합니다. 이는 자기 강화 피드백 루프를 생성합니다: 자주 선택되는 코드워드는 더 많은 그래디언트 업데이트를 받게 되어 점점 더 지배적이 되는 반면, 드물게 선택되는 코드워드는 업데이트를 받지 못해 비활성 상태로 남게 됩니다. 이 "코쿤 효과(cocoon effect)"는 인코더를 좁은 토큰 집합에 가두어 코드북의 전체 용량을 탐색하는 것을 방해합니다.

2. 방법론: RGVQ

이러한 문제를 해결하기 위해, 저자들은 그래프 토폴로지와 특징 유사성을 명시적인 정규화 신호로 통합하는 프레임워크인 **RGVQ (Regularized Graph Vector Quantization)**를 제안합니다. 이 방법은 두 가지 핵심 구성 요소로 이루어집니다:

A. Gumbel-Softmax 재매개변수화 (Gumbel-Softmax Reparameterization)

결정론적 하드 할당의 자기 강화 루프를 깨기 위해, RGVQ는 미분 불가능한 argmax 연산을 Gumbel-Softmax 재매개변수화로 대체합니다.

• 모델은 노드 임베딩과 모든 코드워드 사이의 거리에 기반하여 할당 로짓(assignment logits)을 계산합니다.
• 이 로짓들은 Gumbel 노이즈에 의해 섭동(perturb)되고, 온도 제어 소프트맥스(temperature-controlled softmax)를 통과하여 코드북에 대한 미분 가능한 확률 분포를 생성합니다.
• 양자화된 임베딩은 이 확률들에 기반한 모든 코드워드의 가중 평균으로 계산됩니다.
• 효과: 이를 통해 모든 코드워드에 대해 할당 확률에 비례하여 그래디언트가 흐르게 함으로써, 선택된 것뿐만 아니라 비활성 코드워드에도 업데이트가 전달되어 비활성 코드워드가 재활성화될 수 있도록 보장합니다.

B. 구조 인식 대조 정규화 (Structure-Aware Contrastive Regularization)

그래프 데이터의 중복성과 밀집된 연결성의 영향을 완화하기 위해, RGVQ는 이질적인 노드들 간의 과도하게 집중된 토큰 사용을 명시적으로 처벌하는 정규화 항을 도입합니다.

• 양성 세트 ($N_P$): 주어진 앵커 노드에 대해, 양성 세트는 구조적으로 연결된 노드(이웃) 또는 의미적으로 유사한 노드(Top-$K$ 특징 유사 노드)를 포함합니다.
• 음성 세트 ($N_N$): 연결되어 있지 않거나 특징이 유사하지 않은 노드들을 포함합니다.
• 손실 함수: 모델은 InfoNCE 손실을 사용하여 양성 세트에 있는 노드들이 유사한 토큰 할당 분포를 공유하도록 장려하는 동시에, 음성 세트에 있는 노드들이 겹치는 분포를 갖는 것에 대해 페널티를 부여합니다.
• 효과: 이 정규화는 모델이 구조적 또는 의미적으로 구별되는 노드들을 동일한 토큰으로 매핑하는 것을 억제하여, 토큰 다양성을 촉진하고 특징 중복성에 의한 붕괴를 방지합니다.

전체 손실 함수는 표준 재구성 손실(특징 및 링크 재구성), 코드북 손실, 커밋먼트(commitment) 손실, 그리고 새로운 구조 인식 정규화 손로 결합됩니다.

3. 주요 기여

본 논문의 기여는 다음과 같습니다:

1. 체계적인 실증 연구: 그래프 VQ에서의 코드북 붕괴에 대한 최초의 체계적인 벤치마킹을 제공하며, 다른 도메인의 표준 완화 전략들이 그래프 데이터에서는 실패함을 입증했습니다.
2. 붕괴 진단: 코드북 붕괴가 데이터의 고유한 특성(중복성 및 연결성)과 결정론적 하드 할당의 최적화 역학에 의해 유발됨을 식별했습니다.
3. RGVQ 프레임워크: 코드북 활용도를 높이기 위해 Gumbel-Softmax 재매개변수화와 구조 인식 대조 정규화를 결합한 새로운 미분 가능 프레임워크를 제안했습니다.
4. 포괄적인 평가: 다양한 SOTA 그래프 VQ 백본과 다운스트림 태스크에 걸쳐 RGVQ를 검증하여, 코드북 활용도와 태스크 성능 모두에서 일관된 개선을 보여주었습니다.

4. 실험 결과

저자들은 다양한 동질적(homophilous) 및 이질적(heterophilous) 그래프 데이터셋(예: Cora, PubMed, WikiCS, Amazon-Ratings, Roman-Empire)에 대해 광범한 실험을 수행했습니다.

• 코드북 활용도: RGVQ는 코드북 퍼플렉시티(perplexity) 측면에서 바닐라 그래프 VQ 및 그 변형들(EMA, affine parameters, codebook reset, pretrained encoders 포함)보다 유의미하게 우수한 성능을 보였습니다. 예를 들어, PubMed 데이터셋에서 바닐라 VQ는 심각한 붕괴를 나타내는 ~4.14의 퍼플렉시티를 기록한 반면, RGVQ는 이론적 최대치에 근접한 ~319.09를 달축했습니다. RGVQ는 코드북 크기가 증가함에 따라 높은 활용도를 유지했습니다.
• 전이 가능성: 그래프 파운데이션 모델(GFT)에 통합되었을 때, RGVQ는 교차 태스크 및 교차 도메인 벤치마크(노드 분류, 링크 예측, 그래프 분류)에서 SOTA 성능을 달성하며 지도 학습 기반 GNN 및 기타 자기 지도 학습 방법들을 능가했습니다.
• 직렬화(Serialization): 그래프 양자화 트랜스포머(GQT)의 맥락에서, RGVQ가 강화된 모델은 다른 안티-붕괴(anti-collapse) 방법을 사용한 GQT에 비해 노드 분류 정확도를 일관되게 향상시켰습니다.
• 절제 연구(Ablation Studies):
  • Gumbel-Softmax 재매개변수화 또는 구조 인식 정규화 중 하나라도 제거하면 성능이 크게 저하되었으며, 이는 두 구성 요소가 붕패 방지를 위해 상호 의존적임을 확인시켜 줍니다.
  • 방법론은 온도($\tau$)와 대조 샘플 수의 변화에 견고했으며, 낮은 고정 온도가 충분하다는 것을 입증했습니다.
  • RGVQ는 다양한 수준의 그래프 이질성(heterophily)에 대해서도 강건함을 보였습니다.

5. 의의 및 주장

본 논문은 코드북 붕괴가 부수적인 문제가 아니라 이산 그래프 토큰 학습의 근본적인 병목 현상이라고 주장합니다. 그래프 구조 자체를 자연스러운 정규화 도구로 활용함으로써, RGVQ는 그래프 표현의 표현력과 전이 가능성을 높이는 유연한 "플러그인" 솔루션을 제공합니다.

저자들은 본 연구가 새로운 생성 모델 아키텍처 자체를 제안하는 것이 아니라, 그래프 표현의 보다 효과적인 학습을 가능하게 하는 양자화 과정의 방법론적 개선임을 강조합니다. 이 연구의 의의는 학습된 토큰 어휘가 완전히 활용되도록 보장함으로써, 다운스트림 태스크에 필수적인 미세한 구조적 및 의미적 패턴을 포착하고, 더욱 강력한 그래프 파운데이션 모델을 구축할 수 있게 한다는 점에 있습니다. 결론적으로 RGVQ는 토큰 다양성을 유지하면서 효과적으로 확장 가능한 유연한 프레임워크를 제공하여, 현재 그래프 표현 학습의 핵심적인 한계를 해결한다고 밝히고 있습니다.