기술적 요약: 라플라시안 오토인코더를 통한 원칙적 그래프 잠재 확산

문제 정의

그래프 확산 모델은 분자 설계, 조합 최적화, 신경 아키텍처 검색 등 다양한 도메인의 생성 작업에서 최첨단 성능을 달성했습니다. 그러나 기존 접근법은 일반적으로 그래프 공간에서 직접 작동하여 노드 수에 대해 이차적인 계산 복잡도($O(n^2)$)를 초래합니다. 이러한 병목 현상은 이러한 모델이 노드와 엣지 표현을 모두 학습해야 하며, 종종 희소 그래프에서 엣지의 부재를 모델링하는 데 용량을 낭비하기 때문에 발생합니다. 또한, 확장 가능한 자기회귀 모델을 그래프에 적용하려는 시도는 종종 그래프 구조를 선형화함으로써 치환 등변성을 깨뜨립니다.

유망한 대안은 잠재 확산으로, 생성 과정을 적용하기 전에 데이터를 저차원 공간으로 압축합니다 (VAE 를 통한 이미지 생성에서 볼 수 있듯이). 그러나 이를 그래프에 적용하는 것은 독특한 과제를 제시합니다: 이미지에서는 작은 디코딩 오류가 시각적으로 미미한 것과 달리, 그래프 생성은 거의 손실 없는 재구성을 요구합니다. 인접 행렬을 재구성하는 단일 오류는 전체 샘플을 무효화할 수 있습니다 (예: 화학 결합 파괴 또는 구조적 제약 위반). 기존 잠재 그래프 확산 모델은 종종 이러한 충실도 요구 사항을 충족하지 못하거나, 여전히 이차적인 엣지 토큰에 의존하는 임시 아키텍처를 사용하여 효율성을 희생합니다.

방법론: LG-Flow

저자들은 높은 재구성 충실도를 유지하면서 이러한 장애물을 극복하도록 설계된 잠재 그래프 확산 프레임워크인 LG-Flow를 제안합니다. 이 프레임워크는 원칙적인 오토인코더와 잠재 흐름 매칭 생성 모델이라는 두 가지 주요 구성 요소로 구성됩니다.

1. 라플라시안 그래프 변분 오토인코더 (LG-VAE)

핵심 혁신은 노드 수 ($n \times d$) 에 대해 선형으로 확장되는 잠재 공간으로 그래프를 매핑하는 치환 등변 오토인코더로, 이차적 확장이 아닙니다.

• 이론적 기반: 아키텍처는 인접 식별 위치 인코딩에 기반합니다. 저자들은 노드 임베딩이 충분히 "인접 식별" 가능하면 인접 행렬을 이러한 임베딩에서 재구성할 수 있다는 속성을 활용합니다.
• 인코더: 인코더는 그래프 라플라시안 ($L = U\Lambda U^\top$) 의 고유분해에서 유도된 **라플라시안 위치 인코딩 (LPE)**을 활용합니다. 이는 노드 및 엣지 특징을 가우시안 사후 분포에 통합합니다. 방향성 비순환 그래프 (DAG) 를 처리하기 위해 프레임워크는 엣지 방향성을 보존하는 에르미트 행렬인 자기 라플라시안($L_q$) 으로 확장되어 **자기 LPE(mLPE)**를 생성합니다.
• 디코더: 디코더는 높은 충실도의 재구성을 위해 설계되었습니다.
  • 노드 레이블: 선형 레이어와 소프트맥스를 통해 직접 디코딩됩니다.
  • 인접 행렬: 잠재 임베딩에 적용된 이선형 레이어 ($Z W_Q (Z W_K)^\top$) 를 통해 재구성됩니다. 엣지 감지는 치환 등변성을 보장하고 가변적인 그래프 크기를 처리하기 위해 DeepSet아키텍처를 사용하여 이진 분류 문제로 처리됩니다.
  • 엣지 레이블: 멀티헤드 이선형 레이어를 사용하여 디코딩됩니다.
• 재구성 보장: 이 설계는 잠재 표현이 인접 행렬을 복구하기에 충분한 구조적 정보를 포함하도록 보장하여, 그래프 생성에서 거의 손실 없는 재구성에 대한 중요한 요구 사항을 해결합니다.

2. 잠재 흐름 매칭 (LG-Flow)

LG-VAE 가 훈련된 후, 그 매개변수는 고정됩니다. 그런 다음 잠재 공간 $Z \in \mathbb{R}^{n \times d}$에서 생성 모델이 훈련됩니다.

• 아키텍처: 저자들은 **흐름 매칭 (FM)**목표를 갖춘 **확산 트랜스포머 (DiT)**를 사용합니다.
• 프로세스: 이 모델은 노이즈를 잠재 분포로 운반하는 속도장을 예측하도록 학습합니다.
• 효율성: 잠재 공간에서 작동함으로써 DiT 는 전체 인접 행렬을 직접 예측하는 이차 복잡도를 피합니다. 유일한 이차 단계는 경량 디코더가 처리하는 최종 디코딩입니다. 이를 통해 훨씬 더 큰 생성 백본을 훈련할 수 있습니다.

주요 기여

1. LG-VAE: 무방향 그래프와 DAG 모두에 대해 거의 손실 없는 재구성을 달성하는 원칙적이고 치환 등변인 오토인코더입니다. 이는 인접 식별 속성을 통해 인접 행렬이 잠재 공간에서 복구 가능하다는 이론적 보장을 제공하는 최초의 모델입니다.
2. 선형 확장성: 이 프레임워크는 노드를 고정 차원 임베딩으로 매핑하여 잠재 표현 크기가 노드 수에 따라 선형적으로 확장되도록 보장함으로써, 그래프 공간 확산의 이차적 병목 현상을 효과적으로 제거합니다.
3. 통합 프레임워크: 이 접근법은 무방향 그래프와 방향성 그래프 생성을 단일 모듈식 프레임워크 내에서 통합하여 이미지 생성의 일반적인 탈노이즈 아키텍처 (DiT) 및 기법의 사용을 가능하게 합니다.
4. 성능: 이 방법은 경쟁력 있는 생성 성능을 달성하면서 최첨단 그래프 확산 모델에 비해 최대 1000 배까지의 상당한 속도 향상을 제공합니다.

실험 결과

저자들은 LG-Flow 를 세 개의 합성 (Extended Planar, Extended Tree, Ego), 두 개의 분자 (MOSES, GuacaMol), 두 개의 DAG 데이터세트 (TPU Tile, Price) 등 여섯 개의 데이터세트에서 평가했습니다.

• 재구성 충실도 (Q1): 구조적 벤치마크 (Extended Planar 및 Extended Tree) 에서 LG-VAE 는 거의 손실 없는 재구성을 달성하여 (각각 0.9961 및 0.9844 의 샘플 정확도), 샘플을 정확하게 재구성하지 못한 DGVAE 및 GLAD 와 같은 베이스라인을 크게 능가했습니다.
• 생성 성능 (Q2):
  • 합성 그래프: LG-Flow 는 분포 지표 (차수, 궤도, 군집화, 스펙트럼) 및 유효성, 고유성, 신규성 (V.U.N.) 에서 경쟁력 있는 결과를 달성했습니다.
  • 분자: GuacaMol 및 MOSES 에서 이 방법은 유효성 및 화학적 분포 일치 (FCD) 에서 최첨단 또는 경쟁력 있는 성능을 달성했으며, 종종 DiGress 및 DisCo 와 같은 이산 확산 모델을 능가했습니다.
  • DAG: TPU Tile 데이터세트에서 LG-Flow 는 높은 유효성을 유지하면서 분포 지표에서 자기회귀 방식인 LayerDAG 와 확산 기반인 Directo 를 능가했습니다.
• 효율성 (Q3):
  • LG-Flow 는 DeFoG 및 Cometh 와 같은 베이스라인에 비해 10 배에서 1000 배까지 빠른 추론을 보여주며 막대한 속도 향상을 입증했습니다.
  • 메모리 사용량은 배치 크기에 따라 선형적으로 확장되어 151 개의 샘플로 구성된 전체 테스트 배치를 단일 GPU 에 적합하게 만들었으며, 반면 DeFoG 는 32 개의 배치 크기에서도 메모리 부족이 발생했습니다.

중요성 및 주장

이 논문은 LG-Flow가 고희실도 노드 잠재 공간으로 그래프 확산을 이동시키는 것이 특히 중소형 그래프의 경우 추론 효율성을 크게 향상시킬 수 있음을 보여준다고 주장합니다. 이 프레임워크는 복잡한 그래프 전용 확산 모델 설계에서 강력하고 훈련 비용이 낮은 그래프 오토인코더 개발로 초점을 이동시킵니다.

무방향 및 방향성 그래프 생성을 통합하고 이미지 생성의 강력한 아키텍처인 확산 트랜스포머의 이전을 가능하게 함으로써, 저자들은 이전에 이차적 복잡성으로 인해 처리 불가능했던 훨씬 더 큰 그래프로 생성 모델을 확장할 수 있는 길을 제시합니다. 이 연구는 재구성 충실도를 보장하는 원칙적인 오토인코더를 통해 잠재 확산의 이점 (효율성, 모듈성, 용량) 을 그래프 도메인에서 성공적으로 실현할 수 있음을 강조합니다.