기술 요약: GraViti – 완화된 치환 불변성을 갖춘 그래프 레벨 변분 오토인코더

문제 정의

그래프 레벨 오토인코딩은 보간, 속성 기반 탐색, 생성 모델링과 같은 작업을 용이하게 하기 위해 전체 그래프를 컴팩트한 잠재 벡터로 매핑하는 것을 목표로 합니다. 노드 레벨 변분 오토인코더 (VAE) 는 잘 정립되어 있지만, 근본적인 한계를 겪습니다. 즉, 잠재 공간이 노드 레벨에서 정의되며 임의의 노드 순서에 의존한다는 점입니다. 결과적으로 두 개의 동형 그래프가 완전히 다른 잠재 벡터 집합으로 매핑될 수 있어, 그래프 간의 의미 있는 거리 정의나 그래프 레벨 매니폴드의 포착이 불가능해집니다.

기존의 그래프 레벨 접근 방식은 트레이드오프에 직면해 있습니다. PIGVAE 와 GRALE 와 같은 방법들은 엄격한 치환 불변성을 강제하려 시도하며, 종종 입력 그래프와 재구성된 그래프를 정렬하기 위해 그래프 매칭, Sinkhorn 알고리즘과 같은 계산 비용이 많이 드는 매칭 절차에 의존합니다. 이러한 정렬 단계는 상당한 계산 오버헤드 (일반적으로 3 차 복잡도) 를 도입하여 대규모 데이터셋에서의 확장성을 제한할 수 있습니다. furthermore, 많은 그래프 레벨 디코더는 고정된 최대 그래프 크기를 부과하거나 "더미 노드" 클래스를 도입하는 패딩 전략에 의존하는데, 이는 더 큰 그래프로의 일반화를 제한하고 분류 작업에 편향을 초래합니다.

방법론: GraViti 아키텍처

GraViti 는 일관된 표준 노드 순서가 존재할 때 치환 불변 재구성에 대한 엄격한 요구 조건을 완화하면서 전체 그래프를 단일 잠재 벡터로 인코딩하도록 설계된 트랜스포머 기반 그래프 레벨 VAE 입니다.

1. 그래프 표현

이 모델은 이산적인 노드 및 엣지 클래스 (예: 원자 유형, 결합 유형) 를 가진 그래프에서 작동합니다.

• 입력: 노드 클래스 행렬 $X$, 엣지 클래스 텐서 $E$, 그리고 노드/엣지 유형 수와 그래프 크기를 인코딩하는 그래프 레벨 통계 정보를 담은 글로벌 특징 벡터 $y$입니다.
• 위치 인코딩: 구조적 정보를 주입하기 위해 모델은 라플라시안 고유벡터와 랜덤 워크 특징을 위치 인코딩 ($P$) 으로 활용합니다.
• 도메인 지식: 분자와 같은 특정 도메인의 경우, 원자 질량이나 전기 음성도와 같은 추가적인 노드별 특징 ($K$) 이 통합됩니다.
• 배치 처리: 그래프는 배치 내 가장 큰 그래프의 크기로 패딩되지만, 모델은 이러한 패딩 요소를 완전히 무시하도록 설계되어 더미 노드 클래스의 필요성을 제거합니다.

2. 인코더

인코더는 입력 트리플릿 $(X, E, y)$을 잠재 분포 파라미터 $(z_\mu, z_\sigma)$로 매핑합니다.

• 임베딩: 노드 클래스, 엣지 클래스, 그리고 글로벌 특징은 임베딩으로 투영됩니다.
• 그래프 트랜스포머 스택: 임베딩은 그래프 트랜스포머 레이어 스택 (DiGress 및 CatFlow 아키텍처 기반) 을 통해 처리됩니다. 이는 인코딩 단계에서 치환 공변성을 보장하면서 노드, 엣지, 그리고 그래프 레벨 특징을 동시에 업데이트할 수 있게 합니다.
• 리드아웃: 전통적인 합/평균 풀링과 달리 GraViti 는 시퀀스 트랜스포머 기반 리드아웃을 사용합니다. 두 개의 병렬 시퀀스 디코더 트랜스포머가 각각 노드 및 엣지 표현을 처리하며, 글로벌 컨텍스트를 주입하기 위해 글로벌 특징 벡터 $y_{rep}$을 타겟으로 사용합니다. 이는 리드아웃 단계에서 출력이 노드 순서에 독립적이도록 보장합니다.

3. 그래프 크기 예측

중요한 혁신은 잠재 표현의 나머지 부분에서 그래프 크기를 명시적으로 분리하는 것입니다.

• 모델은 잠재 공간의 처음 $d_s$ 차원이 그래프의 로그 크기를 인코딩하도록 강제합니다.
• MLP 는 이러한 차원으로부터 로그 크기 $\hat{s}$를 예측합니다.
• 디코딩 중에는 실제 로그 크기 $s$가 다시 인코딩되어 잠재 표현에 재주입되어 올바른 노드 수 생성을 안내하며, 이는 디코더가 마스킹에 의존하지 않고 목표 그래프 크기를 알 수 있게 합니다.

4. 디코더

디코더는 잠재 벡터 $z$로부터 그래프를 재구성합니다.

• 브로드캐스팅: 잠재 벡터는 $N$번 복제됩니다 (여기서 $N$은 예측되거나 실제 크기입니다) 그리고 대칭성을 깨기 위해 사인 위치 인코딩이 추가됩니다.
• 시퀀스 트랜스포머: 시퀀스 트랜스포머 인코더가 이 시퀀스를 처리하여 노드 지향 및 엣지 지향 시퀀스로 분할합니다.
• 정제: 두 번째 그래프 트랜스포머 스택이 노드 및 엣지 표현을 정제합니다.
• 예측 헤드: 최종 MLP 가 노드 유형과 엣지 유형을 분류합니다. 모델은 그래프 크기를 직접 예측하여 "no-edge" 더미 클래스의 필요성을 제거합니다.

5. 학습 목적 함수

모델은 표준 VAE 목적 함수를 사용하여 학습됩니다.

• 재구성 손실: 클래스 불균형을 처리하기 위한 노드 및 엣지 분류용 Focal Loss 와 그래프 크기 예측용 smooth L1 손실의 조합입니다.
• 변분 손실: KL-발산 항이 잠재 공간을 표준 정규 사전 분포로 정규화합니다.
• 완화된 치환 불변성: 핵심적으로, 재구성 손실은 입력 노드 순서를 기반으로 적용됩니다. 저자들은 분자나 베이지안 네트워크와 같이 신뢰할 수 있는 표준 순서가 있는 도메인에서는 매칭을 통한 엄격한 치환 불변성 강제가 해롭다고 주장합니다. 일관된 순서에 의존함으로써 모델은 2 차 복잡도와 더 높은 재구성 충실도를 달성합니다.

주요 기여

1. 확장 가능한 아키텍처: GraViti 는 값비싼 치환 불변 매칭 절차를 피함으로써 대규모 그래프 데이터셋 (PubChem16, PubChem32) 에서 최첨단 재구성 정확도를 달성합니다.
2. 구조화된 잠재 공간: 학습된 잠재 공간은 잘 구조화되어 있어 부드러운 보간, 제어된 편집, 속성 기반 최적화를 지원합니다. 분자 벤치마크에서 모델은 잠재 공간으로부터 직접 확립된 화학 원리 (예: 알킬 사슬 길이와 소수성 간의 관계) 를 복원합니다.
3. 완화된 치환 불변성: 이 논문은 경험적으로 엄격한 치환 불변성이 항상 필요한 것은 아님을 보여줍니다. 일관된 표준 노드 순서가 사용 가능한 경우, 이 제약을 완화하는 것이 복잡한 매칭을 통해 불변성을 강제하는 방법들에 비해 재구성 품질과 확장성을 크게 향상시킵니다.

실험 결과

저자들은 GraViti 를 여러 벤치마크에서 평가했습니다.

• 재구성: 분자 데이터셋 (QM9, PubChem16, PubChem32) 에서 GraViti 는 재구성 정확도에서 베이스라인 (GRALE, PIGVAE) 보다 최대 25% 더 우수한 성능을 보였으며, 특히 매칭 기반 방법이 어려움을 겪는 더 큰 그래프에서 두드러졌습니다.
• Ablation 연구: 트랜스포머 인코더나 트랜스포머 기반 리드아웃을 제거하면 성능이 크게 저하되어 이러한 구성 요소의 중요성을 확인했습니다.
• 디노이징: 모델은 노이즈가 있는 입력으로부터 화학적으로 유효한 구조를 효과적으로 복원하며, AE 변형에 비해 VAE 변형이 더 뛰어난 견고성을 보입니다.
• 속성 최적화: 잠재 공간에서 CMA-ES 를 사용하여 모델은 logP (지용성) 를 성공적으로 최대화하여 점진적으로 더 긴 탄소 사슬이나 특정 소수성 부분을 가진 분자를 생성하며, 알려진 화학적 경향과 일치합니다.
• 보간: 잠재 공간에서의 선형 보간은 분자 구조 간의 부드러운 전환을 생성합니다.
• 하류 작업: PubChem 에서 학습된 잠재 벡터는 분자 특성 (쌍극자 모멘트, HOMO-LUMO 갭) 과 베이지안 네트워크의 베이지안 정보 기준 (BIC) 점수를 높은 정확도로 예측합니다.
• 생성: PubChem 에서는 생성 유효성이 QM9 에 비해 낮습니다 (방향족 표현에 따라 약 50-70%). 그러나 모델은 다양하고 새로운 분자를 생성합니다. 저자들은 더 큰 데이터셋에서의 낮은 유효성을 지역적 구조 실패가 아닌 잠재 공간의 글로벌 컴팩트성 때문이라고 설명합니다.

중요성 및 주장

이 논문은 GraViti 가 그래프 레벨 표현 학습을 위한 실용적인 기초를 제공한다고 주장합니다. 그 주요 중요성은 안정적인 노드 순서가 존재할 때 엄격한 치환 불변성이 효과적인 그래프 레벨 오토인코딩을 위한 전제 조건이 아님을 입증하는 데 있습니다. 이 제약을 완화함으로써 GraViti 는 그래프 매칭의 계산 병목 현상을 피하여 대규모 데이터셋에서의 확장 가능한 학습을 가능하게 하면서도 높은 재구성 품질을 유지합니다.

저자들은 GraViti 를 복잡한 생성 파이프라인에 대한 경량 대안으로 위치시키며, 생성 성능과 실용적인 샘플 품질 사이의 균형을 맞추는 단일 단계 디코딩 메커니즘을 제공합니다. 이 연구는 분자 그래프와 베이지안 네트워크와 같이 표준 순서가 내재되어 있거나 구축될 수 있는 도메인의 경우, 정렬 기반 손실의 오버헤드 없이 그래프 레벨 모델이 우수한 성능을 달성할 수 있음을 시사합니다. 향후 작업은 생성 유효성을 더욱 향상시키고 회로 및 단백질과 같은 다른 그래프 도메인으로의 전이를 탐구하기 위해 도메인 인식 제약을 통합하는 것을 제안합니다.