기술 요약: 확산 기반 상태 공간 모델 (DDSSM)

문제 정의

시계열 모델링을 수행하는 실무자들은 모델의 표현력(expressivity)과 잠재 시스템 역학(latent system dynamics)을 포착하는 능력 사이의 절충안에 직면하곤 합니다. 기존 방식들은 대개 한 가지 목표를 위해 다른 하나를 희생합니다:

• 딥 상태 공간 모델 (Deep State Space Models, DSSMs) (예: Deep Kalman Filter)은 원칙적인 추론과 엔드 투 엔드(end-to-end) 학습을 제공하지만, 일반적으로 가우시안 전이 밀도(Gaussian transition densities)를 가정합니다. 이러한 가정은 표현력을 제한하여 다봉형(multimodal) 역학을 포착하는 데 실패할 수 있으며, 가우시안 사전 분포가 잠재 표현을 과도하게 정규화하여 필요한 구조적 정보를 버리게 만들 수 있습니다.
• **확산 모델 (Diffusion Models)**은 매우 높은 표현력을 가지며 복잡한 다봉 분포를 모델링할 수 있습니다. 그러나 시계열에 적용될 때, 이들은 종종 과거 관측값에 조건화된 생성 샘플러로 사용되거나, 2단계 학습(먼저 오토인코더를 학습한 후 확산 모델을 학습)에 의존합니다. 이러한 2단계 접근 방식은 다운스트림 생성 역학에 필요한 정보를 버리는, 과도하게 정규화된 잠재 공간을 생성할 수 있습니다.

두 프레임워크의 강점을 결합하기 위해 시퀀스 데이터에 대해 오토인코더와 확산 모델을 원칙적인 방식으로 공동 학습시키는 방법에 대해서는 여전히 미해결 과제가 남아 있습니다.

방법론

저자들은 확산 모델을 상태 공간 모델(SSM)의 전이 밀도(transition density) 내에 직접 통합하는 **확산 기반 상태 공간 모델 (Diffusion-Driven State Space Model, DDSSM)**을 제안합니다.

생성 모델

모델은 관측값 $x_{1:T}$가 공변량 $u_{1:T}$에 조건화될 수 있는 마르코프 과정(Markov process)을 통해 진화하는 잠재 상태 $z_{1:T}$로부터 생성된다고 가정합니다. 결합 가능도(joint likelihood)는 다음과 같이 인수분해됩니다:
$$p(x_{1:T}, z_{1:T} | u_{1:T}) = p_{\eta,\theta}(z_{1:j}, x_{1:j} | u_{1:j}) \prod_{t=j+1}^T p^{(t)}_\psi(z_t | z_{t-j:t-1}, u_{t-j:t}) p^{(t)}_\theta(x_t | z_{t-j+1:t}, u_{t-j+1:t})$$

주요 혁신 사항은 다음과 같습니다:

1. 확산 전이 (Diffusion Transitions): 매개변수화된 가우시안 전이 $p^{(t)}_\psi$ 대신, 모델은 확산 과정을 사용합니다. 각 타임스텝 $t$에 대해, 전이 밀도는 중간 노이즈가 섞인 잠재 상태 $z^1_t, \dots, z^K_t$를 도입함으로써 모델링됩니다. 전이는 다음과 같이 정의됩니다:
   $$p^{(t)}_\psi(z^0_t, z^{1:K}_t | z_{t-j:t-1}, u_{t-j:t}) = p(z^K_t) \prod_{k=1}^K p^{(t)}_\psi(z^{k-1}_t | z^k_t, z_{t-j:t-1}, u_{t-j:t})$$
   여기서 $z^0_t = z_t$이고 $z^K_t \sim \mathcal{N}(0, I)$입니다. 이를 통해 모델은 임의의, 잠재적으로 다봉적인 전이 밀도를 학습할 수 있습니다.
2. 방출 (Emissions): 방출 밀도 $p^{(t)}_\theta$는 신경망에 의해 매개변수화된 표준 가우시안으로 유지되어, 잔차 불확실성(residual uncertainty)을 처리하는 한편, 확산 모델이 복잡한 비가우시안 구조를 포착하도록 합니다.
3. 초기화 (Initialization): 초기 상태 분포 $p_\eta(z_{1:j}|u_{1:j})$는 제한적인 가정을 피하기 위해 변이 계층적 사전 분포(Variational Hierarchical Prior, VHP)를 활용합니다.

추론 및 학습

진정한 사후 확률(true posterior)은 다루기 까다롭기 때문에, 저자들은 변이 추론(variational inference) 프레임워크를 채택합니다.

• 변이 사후 확률 (Variational Posterior): 사후 확률은 평활화된 전이 항(과거 잠재 상태 및 미래 관측값에 의존)과 확산 노이징 프로세스 항을 포함하도록 인수분해됩니다.
• 공동 학습 (Joint Training): 2단계 방식과 달리, DDSSM은 엔드 투 엔드로 학습됩니다. 목적 함수는 다음을 포함하는 ELBO를 최대화합니다:
  • 재구성 손실 (방출의 음의 로그 가능도).
  • 초기화를 위한 KL 발산(KL divergence) 항.
  • 학습된 전이 역학에 따라 예측 가능한 잠재 상태를 인코더가 생성하도록 장려하는 확산 체인 손실 (확산 모델의 역과정에서 유도됨).
• 프리컨디셔닝 (Preconditioning): 공동 학습 중 노이즈 수준에 따른 그래디언트 안정성을 보장하기 위해, 저자들은 확산 네트워크의 입력과 출력을 재조정하는 프리컨디셔닝(Karras et al., 2022)을 사용합니다.

주요 기여

1. DDSSM 프레임워크: 가우시안 가정을 배제하고 유연한 다봉 잠재 역학을 학습할 수 있도록 전이 밀도를 확산 모델로 대체한 상태 공간 모델을 도입했습니다.
2. 엔드 투 엔드 학습: 시퀀스 데이터에 대해 오토인코더와 확산 모델을 공동 학습시키는 문제에 대한 원칙적인 해결책을 제시했습니다. 이는 잠재 시계열 확산에서 흔히 나타나는 과도한 정규화 문제를 방지합니다.
3. 이론적 유도: 이 하이브리드 모델에 대한 ELBO의 엄격한 유도를 통해, 최적의 변이 사후 확률이 진정한 사후 확률의 구조와 일치하도록 인수분해됨을 입증함으로써 근사 격차(approximation gap)를 좁혔습니다.

실험 결과

저자들은 이봉 노이즈(bimodal noise, 이봉 전이를 가진 잠재 랜덤 워크)를 특징으로 하는 시뮬레이션 시계열 데이터셋을 통해 DDSSM을 평가했습니다.

• 설정: DDSSM을 VHP를 사용하는 Deep Kalman Filter(DKF) 및 LOCF(Last Observation Carried Forward) 베이스라인과 비교했습니다.
• 지표: 모델의 1-스텝 예측 분포와 실제 분포 사이의 젠슨-샤논 발산(Jensen-Shannon Divergence, JSD)을 사용하여 성능을 측정했습니다.
• 결과:
  • 가우시안 전이에 의해 제약된 DKF는 두 개의 실제 모드(mode) 사이에 확률 질량을 배치하여, 이봉 구조를 포착하는 데 실패했습니다.
  • DDSSM은 실제 분포의 두 모드를 모두 성공적으로 포착했습니다.
  • 정량적으로, DDSSM은 DKF보다 현저히 낮은 JSD를 달로 달성했습니다.
  • 정성적으로, DDSSM은 데이터를 대표하는 재구성과 예측을 생성한 반면, DKF의 예측은 일관된 패턴이 부족하고 관측 데이터를 잘 재구성하지 못했습니다.

의의 및 주장

본 논문은 잠재 시계열 확산을 상태 공간 모델로서 문제화함으로써, 잠재 표현을 학습하는 데 방해가 되지 않으면서도 유연한 재구성이 가능한 엔드 투 엔드 학습을 달성할 수 있다고 주장합니다.

저자들은 DDSSM을 확산 모델의 표현력과 SSM의 구조적 엄밀함 사이의 긴장을 해결하는 방법으로 포지셔닝합니다. 그들은 전통적인 Deep SSM의 가우시안 사전 분포 가정이 다봉 설정에서 잠재 표현을 학습하는 데 해롭다고 주장하며, 자신들의 접근 방식이 잠재 궤적에 대해 더 표현력 있는 사전 분포를 제공한다고 강조합니다. 이 연구는 2단계 학습 패러다임을 넘어 시계열을 위한 잠재 확산 모델의 문헌을 확장합니다.

저자들은 향-후 연구로서 DDSSM을 대규모 데이터셋에 적용하고, 확산 전이 분포의 능력을 완전히 활용하기 위해 딥 SSM 문헌의 추가적인 아이디어들을 통합할 것이라고 언급했습니다.