Scaling Weisfeiler-Leman Expressiveness Analysis to Massive Graphs with GPUs

기술 요약: GPU를 이용한 대규모 그래프 대상 Weisfeiler–Leman 표현력 분석의 확장

문제 정의
Weisfeiler–Leman (1-WL) 테스트와 그와 관련된 안정적 채색(stable coloring)은 그래프 신경망(GNN)의 표현력을 이해하는 데 있어 핵심적인 역할을 한다. 1-WL 안정적 채색은 메시 전달 GNN이 그래프 구조를 구별할 수 있는 능력의 상한선을 제공한다. 즉, 동일한 안정적 색상을 받는 노드들은 이러한 네트워크에 의해 구별 불가능하다. 그러나 실제 세계의 거대 그래프에 대해 이 안정적 채색을 계산하는 것은 두 가지 결정적인 병목 현상에 직면한다:

1. 순차적 실행: 분할 정제(partition refinement)에 기반한 고전적 알고리즘(예: $O(m \log n)$ 스플리터 기반 접근 방식)은 본질적으로 순차적이며, 현대 GPU 아키텍처의 대규모 병렬성을 활용할 수 없다.
2. 전역 메모리 요구량: 이러한 알고리즘은 전체 그래프가 동시에 메모리에 상주해야 한다. 이러한 "전역적" 특성은 단일 장치나 심지어 대규모 CPU 클러스터의 용량을 초외하는 웹 규모 그래프(예: 수백억 개의 엣지를 가진 그래프)의 분석을 불가능하게 만든다.

문제를 병렬화하거나 선형화하려는 기존의 시도들은 종종 부동 소수점 산술에 의존하여 대규모 환경에서 수치적 불안정성을 초래하거나, 장치 메모리보다 큰 그래프의 메모리 제약 문제를 해결하지 못한다.

방법론
저자들은 1-WL 안정적 채색의 선형 대수적 해석을 활용하여, 계산 확장성과 메모리 확장성을 분리하는 두 갈래의 접근 방식을 제안한다.

1. 무작위 선형 대수적 정제 (계산 확장성)
순차적인 분할 정제 대신, 저자들은 등적 분할(equitable partition)의 불변성을 이용하여 문제를 재정의한다: 분할은 임의의 블록 상수 벡터(block-constant vector)의 이미지가 아드자센시 행렬(adjacency matrix) 하에서도 블록 상수로 유지될 때 안정적이다.

• 알고리즘: 이 방법은 현재 분할에 대해 대칭인 무작위 벡터를 샘플링하고, 이를 아드자센시 행렬과 곱한 후($z = Aw$), 결과 벡터의 대칭성을 바탕으로 분할을 정제한다.
• 산술 도메인: 부동 소수점 거듭제곱 반복법(power iteration)이 대규모 그래프에서 발생시킬 수 있는 수치적 불안정성(잘못된 병합 유발)을 피하기 위해, 알고리즘은 전적으로 부호 없는 64비트 정수 산술인 모듈러 링 $\mathbb{Z}/2^{64}\mathbb{Z}$ 내에서 작동한다.
• 이론적 보장: 아드자센시 행렬의 행 합(row sum)이 $2^{64}/2$로 제한된다는 구조적 제약을 활용하여, 저자들은 확률적 정확성 보장을 도출한다. 무작위 프로브(probe)는 최대 $10^{-9}$의 실패 확률로 미세한 정제 누락를 감지하며, 이는 필드(field)를 필요로 하는 고전적인 슈츠-지펠(Schwartz–Zippel) 보조정리를 우회한다.
• GPU 매핑: 핵심 연산은 단일 희소 행렬-벡터 곱셈(SpMV)과 사전식 정렬(lexicographic sort)로 축소되며, 이는 GPU 하드웨어에서 매우 효율적으로 수행된다.

2. 정확성을 보존하는 배칭 (메모리 확장성)
장치 메모리에 담기지 않는 그래프를 처리하기 위해, 저자들은 엣지 집합을 서로소 부분 집합으로 분해하면서도 전역 안정적 채색의 정확성을 유지하는 배칭 기법을 도입한다.

• 내부 노드 vs 경계 노드: 주어진 엣지 배치(batch)에 대해, 노드는 "내부(inner)" 노드(모든 나가는 엣지가 배치 내에 존재) 또는 "경계(boundary)" 노드(나가는 엣지가 배치 경계를 가로지름)로 분류된다.
• 국소적 정제: 알고리즘은 각 배치를 독립적으로 처리한다. 결정적으로, 각 배치의 초기 분할은 모든 내부 노드를 하나의 블록으로, 모든 경계 노드를 싱글톤(singleton) 블록으로 배치한다. 이는 정제가 해당 배치 내에서 완전한 이웃 정보를 가진 노드들만을 병합하도록 보장한다.
• 반복적 수축: 국소적 정제 결과는 전역 분할로 병합된다. 그런 다음 알고리즘은 quotient graph(노드가 현재 분할의 블록을 나타냄)를 구성하고 이 과정을 반복한다.
• 종료 조건: 그래프가 단일 배치에 들어갈 때까지(최적의 안정적 채색을 얻음) 또는 고정점(fixed point)에 도달할 때까지 반복한다. 만약 메모리 용량에 도달하기 전에 고정점에 도달하면, 결과는 가장 거친(coarsest) 분할의 정제된 유효한 안정적 채색이 된다.

핵심 기여

1. 무작위 정제 알고리즘: 정밀도 손실이 발생하는 부동 소수점 방식 대신, 엄격한 확률적 보장($<10^{-9}$ 오차)을 갖는 정수 산술 기반의 몬테카를로 프레임워크를 통한 1-WL 안정적 채색 구현.
2. 배칭 기법: 독립적인 엣지 부분 집합을 처리하고 quotient graph를 통해 결과를 병합함으로써, 장치 메모리를 초과하는 그래프에서도 1-WL 계산을 가능하게 하는 분해 전략.
3. GPU 구현: 이러한 알고리즘을 SpMV, 정렬, 리덕션(reduction)과 같은 GPU 프리미티브에 직접 매핑하여 CUDA 기반 구현을 달성하고, CPU 베이스라인 대비 엄청난 속도 향상을 실현.

실험 결과
저자들은 11개의 중규모 그래프와 3개의 웹 규모 그래프(33~48 tỷ 엣지)를 대상으로 접근 방식을 평가했다.

• 정확성: 전체 그래프가 메모리에 들어가는 중규모 그래프($N_b=1$)에서, 무작위 알고리즘은 결정론적 베이스라인과 일치하는 정확한 최적 안정적 채색을 복구했다.
• 수치적 안정성: 제안된 정수 산술 방식은 대규모 그래프(예: it-2004 데이터셋에서 800,000개 이상의 블록 차이 발생)에서 이전의 부동 소수점 거듭제곱 반복법이 잘못 병합했던 노드들을 성공적으로 구별해냈다.
• 배칭 품질: 메모리 제약을 시뮬레이션하기 위해 여러 배치($N_b=2, 4$)를 사용했을 때, 결과 분할은 경험적으로 최적 안정적 채색의 5% 이내에 있었다. 일부 데이터셋의 경우 그 격차는 0이었다.
• 성능:
  • 속도: 단일 GPU 구현은 고전적인 단일 스레드 CPU 분할 정제 대비 최대 138배, 멀티코어 CPU 베이스라인 대비 최대 449배의 속도 향상을 달升했다.
  • 확장성: 300억 개 이상의 엣지를 가진 웹 규모 그래프에서, 모든 CPU 베이스라인(멀티코어 심볼릭 리덕션 포함)은 타임아웃이 발생하거나 메모리 부족으로 중단되었다. 제안된 방법은 이러한 인스턴스들에 대해 성공적으로 안정적 채색을 계산했으며, 약 39~46%의 블록 감소를 달성했다.

의의 및 주장
본 논문은 300억 개 이상의 엣지를 가진 웹 규모 그래프에서 1-WL 안정적 채색을 성공적으로 계산한 첫 번째 사례라고 주장하며, 이는 메모리와 계산 능력의 한계로 인해 이전에는 불가능했던 작업이었다. 본 연구의 의의는 다음과 같다:

• 표현력 분석 가능화: 소규모 또는 중규모 그래프에 국한되지 않고, 실제 거대 데이터셋에 대해 GNN의 표현력을 경험적으로 분석할 수 있게 한다.
• 아키텍처 적합성: 전통적으로 순차적이고 전역적인 문제로 간-주되었던 1-WL 문제를, GPU 아키텍처에 최적화된 선형 대수적 공식과 배칭 전략을 사용하여 효과적으로 병렬화 및 분산할 수 있음을 입증한다.
• 강건성: 정수 산술을 통해 수치적으로 안정적이며, 배칭을 통해 메모리 확장이 가능함을 보여줌으로써 기존 접근 방식의 실패 모드를 극복한다.

저자들은 이 방법이 이전에는 도달할 수 없었던 규모에서 1-WL 표현력 분석을 가능하게 한다고 결론지으며, 안정적 채색과 등적 분할 사이의 대응 관계가 향후 다운스트림 그래프 ML 작업 및 비선형 동역학 시스템의 모델 축소 등에 적용될 수 있는 잠재적 경로를 제공한다고 언급했다.