동적 그래프를 위한 연속 상태공간 모델 DyG‑Mamba
초록
**
DyG‑Mamba는 불규칙한 타임스탬프를 제어 신호로 활용해 기억 소멸을 동적으로 조절하는 연속‑시간 상태공간 모델이다. 시간 간격을 학습 가능한 step‑size 함수에 매핑하고, 핵심 매개변수 A, B, C를 입력‑의존적으로 재정의함으로써 장기 의존성 포착, 메모리 효율, 그리고 노이즈에 대한 강인성을 동시에 달성한다. 12개의 벤치마크에서 링크 예측·노드 분류 모두 최첨단 성능과 선형 시간·메모리 복잡도를 보였다.
**
상세 분석
**
DyG‑Mamba는 기존 동적 그래프 모델이 안고 있던 두 가지 근본적인 한계를 극복한다. 첫 번째는 장기 의존성을 효율적으로 학습하지 못한다는 점이다. RNN 기반 방법은 기울기 소실·폭발 문제로 긴 시퀀스에 취약하고, Transformer 기반 방법은 O(N²) 복잡도로 메모리·시간 비용이 급증한다. DyG‑Mamba는 연속‑시간 상태공간 모델(SSM)을 기반으로 하면서, Mamba에서 제안된 parallel‑scan 최적화를 그대로 적용해 O(N) 복잡도를 유지한다. 특히, “step‑size” Δt를 고정값이 아니라 시간 간격에 비례하는 학습 가능한 함수(Δtₖ = w₁ ⊙ (1‑exp(‑w₂·(tₖ₊₁‑tₖ)/(τ‑t₁)))) 로 정의함으로써, 오래된 이벤트일수록 더 큰 Δt가 적용되어 Aₖ = exp(Δtₖ·Aₖ) 가 급격히 감소한다. 이는 Ebbinghaus의 망각 곡선과 일치하는 “fast‑then‑slow” 기억 소멸 메커니즘을 구현한다. 결과적으로 모델은 불규칙한 타임스탬프를 자연스럽게 반영하면서도, 장기 정보를 과도하게 보존하거나 급격히 소멸시키는 문제를 피한다.
두 번째 한계는 노이즈에 취약하다는 점이다. 기존 SSM은 B, C 매개변수를 데이터‑의존적으로 초기화하지만, 입력 노이즈가 직접 영향을 미쳐 선택적 복사(selective copy) 능력이 저하된다. DyG‑Mamba는 B와 C를 입력‑조건부 선형 변환(Bₖ = Linear′(input), Cₖ = Linear(input)) 으로 재정의하고, 각각에 스펙트럴 노름 제약을 부여해 Lipschitz 연속성을 보장한다. 이렇게 하면 노이즈가 들어와도 B, C가 과도히 증폭되지 않아, 중요한 과거 정보를 효과적으로 “리뷰”하고 잡음은 억제한다. 또한, Ebbinghaus의 복습 주기 이론을 차용해, 일정 간격마다 과거 상태를 재활성화하는 메커니즘을 모델 내부에 내재시켰다.
구조적으로는 (1) 노드·엣지·시간·공동이웃 인코딩을 각각 선형 변환 후 concat하여 4d 차원의 입력 Z를 만든다. (2) Z에 1‑D Conv + SiLU를 적용해 차원을 8d로 확장하고, 이를 SSM에 입력한다. (3) 재정의된 Δt, A, B, C를 사용해 연속‑시간 상태 전이와 출력 투영을 수행한다. (4) 최종 출력은 C·h와 원본 Z의 SiLU‑활성화된 선형 변환을 element‑wise 곱한 뒤, 또다시 선형 변환해 4d 차원의 노드 임베딩을 얻는다. 이 파이프라인은 완전한 미분 가능성을 유지하면서도, GPU‑friendly parallel‑scan을 활용해 대규모 시퀀스(길이 10⁴ 이상)도 메모리·시간 효율적으로 처리한다.
실험에서는 12개의 동적 그래프 데이터셋(링크 예측·노드 분류 모두)에서 DyG‑Mamba가 기존 최첨단 모델(JODIE, TGN, DyGFormer, GraphMixer 등)을 대부분 능가했다. 특히, 시간·메모리 효율 측면에서 동일 GPU 메모리 제한 하에 전체 시퀀스를 풀‑프롤링(pooled 없이) 처리했으며, 노이즈 주입 실험(50% 가짜 엣지)에서도 성능 저하가 10% 이하로 제한돼 강인성을 입증했다. 코드와 데이터는 공개 저장소에서 재현 가능하도록 제공된다.
**
댓글 및 학술 토론
Loading comments...
의견 남기기