기술 요약: 스티키 점프 확산 (Sticky Jump Diffusions, SJDs)

문제 정의

확산 기반 생성 모델링은 연속적인 데이터에서 이산적인 시퀀스(예: 텍text)로 성공적으로 확장되었으나, 기존 방식들은 이산적 역학과 연속적 역학을 통합하는 데 있어 근본적인 한계에 직면해 있습니다. 두 가지 지배적인 계열이 존재합니다:

1. 마스크 확산 (Masked Diffusion): 이산 토큰 공간에서 작동하며, 구조적으로 유효한 토큰을 출력합니다. 그러나 이는 각 위치를 단일 흡수 토큰(absorbing token)에 매핑되는 것으로 취급하여, 오염(corruption)이 후보 토큰과의 근접성에 대해 제공하는 단계적인 정보를 삭제합니다.
2. 연속 확산 (Continuous Diffusion): 표준 스코어 기반 학습을 사용하여 토큰 임베딩 상에서 작동합니다. 이들의 종단 상태(terminal state)는 이산적 지지 집합(discrete support)을 벗어나 있으며, 이산 토큰을 복구하기 위해 외부 투영 단계(예: 디코더 또는 양자화)를 필요로 합니다. 이는 학습된 역학의 일부가 아닙니다.
3. 하이브리드 확산 (Hybrid Diffusion): 마스크 체인과 연속 잠재 변수를 결합합니다. 효과적이기는 하지만, 이러한 모델들은 약속률(commitment rates), 목적지(destinations), 드리프트(drifts)를 하나의 통합된 순방향 과정으로부터 도출하기보다는 별개의 설계 선택지로 취급하여 역방향 샘플러를 수동으로 조립하는 경우가 많습니다.

핵리 과제는 이 두 영역을 자연스럽게 연결하는 연속 시간 마르코프 과정을 정의하는 것입니다. 여기서 역방향 역학(이산 토큰으로의 약속 및 연속 확산)은 휴리스틱하게 설계되는 것이 아니라, 플럭스 균형(flux balance)을 통해 순방향 과정으로부터 수학적으로 도출되어야 합니다.

방법론: 스티키 점프 확산 (SJDs)

저자들은 $\mathbb{R}^d$ 상의 유한 또는 가산 개의 앵커(anchors) $A$(토큰 임베딩과 동일시됨)를 가진 연속 시간 마르코프 과정 클래스인 **스티키 점프 확산 (SJDs)**을 소개합니다.

1. 순방향 역학 (Unsticking)

순방향 시간 동안, 데이터는 앵커 $a \in A$에서 시작합니다. 프로세스는 다음과 같이 진화합니다:

• 스티킹 단계 (Sticking Phase): 입자는 해저드 비율(hazard rate) $\vec{\lambda}_t$에 의해 제어되는 무작위 지속 시간 동안 앵커 $a$에 머뭅니다.
• 언스킹 (Unsticking): 해저드 이벤트 발생 시, 입자는 커널 $R_t(\cdot | a)$에 의해 생성되며, 이 커널은 연속 영역 $X_A = \mathbb{R}^d \setminus A$ 상에 지지됩니다.
• 확산 (Diffusion): $X_A$에 들어서면, 궤적은 표준 확률 미분 방정식(SDE) $d\vec{X}_t = \vec{b}_t(\vec{X}_t)dt + g_t d\vec{W}_t$에 따라 확산됩니다.
• 질량 분할 (Mass Split): 주변 분포 $p_t$는 $X_A$ 상의 절대 연속 밀도와 앵커 $A$ 상의 원자 질량(atomic masses)으로 분리됩니다.

핵심 혁신 (교차 위치 블렌딩 - Cross-Position Blending): 일반적인 하이브리드 모델과 달리, 언스킹 커널이 로컬 앵커에만 의존하는 것과 달리, SJD는 블렌딩 행렬 $W$를 통해 전체 시퀀스에 의존할 수 있도록 허용합니다. 위치 $i$에서의 언스킹 분포의 평균은 이웃 임베딩들의 가중 블렌드입니다: $\mu_i(\vec{X}_0) = (W E(\vec{X}_0))_i$. 이는 구조적 귀납 편향(예: 공간적 국소성 또는 제약 그래프)을 오염 과정에 직접 인코딩합니다.

2. 역방향 역학 (Sticking)

SJD의 시간 역전은 두 가지 결합된 메커니즘을 가진 생성 과정을 산출합니다:

• 스코어 구동 SDE: 연속 영역 $X_A$ 상에서 프로세스는 스코어 함수 $\nabla \log p_t(x)$에 의해 구동되는 역방향 SDE를 따릅니다.
• 스티키 점프 커널 (Sticky Jump Kernel): 프로세스는 연속 영역에서 앵커로 점프합니다. 역방향 해저드 비율(언제 점프할 것인가)과 점프 목적지(어디로 점 jump 할 것인가)는 플럭스 균형 방정식(순방향 흐름과 역방향 흐름 사이의 상세 균형)에 의해 유일하게 결정됩니다.
  • 결정적으로, 약속 비율과 목적지는 수동으로 설계된 스케줄이 아니라, 순방향 법칙에 의해 고정된 값입니다.

3. 학습: 디노이징 해저드 매칭 (Denoising Hazard Matching, DHM)

역방향 과정을 학습 가능하게 만들기 위해, 저자들은 디노이징 스코어 매칭의 해저드 버전인 **디노이징 해저드 매칭 (DHM)**을 제안합니다.

• 단일 분류기: 단일 디노이징 분류기 $P_\theta(a | y, t)$가 SJD로 오염된 쌍 $(\vec{X}_0, \vec{X}_t)$에 대해 표준 교차 엔트로피를 통해 학습됩니다.
• 해석적 재가중치 부여 (Analytic Reweighting): 역방향 해저드와 할당 확률은 두 번째 네트워크 헤드 없이 DHM을 통해 분석적으로 복구됩니다.
  • 역방향 해저드 $\vec{\lambda}^\star_t(y)$는 SJD 오염 하에서 순방향으로 알려진 타겟(DHM 타겟)의 조건부 기대값으로 추정됩니다.
  • 할당 확률 $\pi^\star_t(a|y)$는 분류기의 로짓(logits)을 DHM 타겟으로 재가중하여 도출됩니다.
• 시뮬레이션 프리 학습 (Simulation-Free Training): 표준 조건(폐쇄형 생존 함수, 가우시안 합성의 폐쇄성) 하에서, 학습은 SJD 경로의 시뮬레이션을 필요로 하지 않습니다. 타겟에 필요한 혼합 밀도는 폐쇄형 쿼드러처(quadrature)를 통해 평가되므로, 학습 파이프라인은 시뮬레이션이 필요 없고 효율적입니다.

주요 기여

1. SJD 프로세스 클래스: 유한 또는 가산 개의 앵커를 가진, 순방향 주변 분포가 연속 밀도와 이산 원자로 분리되는 점프 확산에 대한 통합 프레임워크를 제공합니다. 이는 고차원 앵커 집합으로 클래식한 스티키 경계 확산을 일반화합니다.
2. 단일 분류기로부터의 역방향 역학: 저자들은 단일 디노이징 분류기로부터 스코어와 각 앵커별 역방향 강도를 모두 DHM을 통해 추정할 수 있음을 입증했습니다. 이는 별도의 해저드 헤드를 제거하고, 분류기 출력에 대한 Rao-Blackwellized 재가중치 부여에 의존합니다.
3. 확산 패밀리에 대한 통합적 관점: SJDs는 기존 모델들을 특정 극한으로서 복구합니다:
   • 마스크 확산 (Masked Diffusion): 언스킹 커널이 단일 흡수점으로 붕괴될 때 발생합니다. 마스크 확산의 "정보 공백"은 구조적으로 설명됩니다: 언스킹 커널이 앵커 독립적이므로, 오염된 상태는 소스 토큰에 대한 증거를 전혀 담고 있지 않습니다.
   • 연속 확산 (Continuous Diffusion): 점프 비율이 0일 때 발생합니다. 터미널 투영이 필요한 이유는 순방향 주변 분포에 원자가 부재하여 역방향 점프가 발생하지 않기 때문입니다.
   • 하이브리드 확산 (Hybrid Diffusion): 언스킹 커널이 순방향 섭동 커널과 일치할 때 발생합니다. 약속 비율, 목적지, 드리프트는 별도의 설계가 아닌 플럭스 균형으로부터 자연스럽게 뒤따름이 입증되었습니다.
4. 설계 공간으로서의 점프 커널: 블렌딩 행렬 $W$를 변화시킴으로써, 저자들은 새로운 설계 공간을 도입합니다. 교차 위치 블렌딩은 각 위치를 이웃들의 깨끗한 값들의 블렌드로 오염시키며, 이를 통해 의존 구조(공간적 국소성 또는 제약 그래프)를 오염 자체의 귀납적 편향으로 전환합니다.

실험 결과

저자들은 구조적 의존성이 알려진 세 가지 데이터셋에서 교차 위치 블렌딩 축을 경험적으로 검증했습니다:

• CIFAR-10 (이미지 생성): 공간적 국소성을 인코딩하기 위해 값 공간 블러 커널(value-space blur kernel)을 사용하였으며, 비로컬 커널($\sigma_B = 2.5$)을 사용한 SJD는 256 NFE에서 아이덴티티 커널 하이브리드(CADD, 15.88) 및 다른 베이스라인(MDLM, Bit Diffusion)보다 낮은 FID(14.57)를 달달성했습니다.
• Text8 (텍스트 생성): 국소적 순차 상관관계를 인코딩하기 위해 1D 가우시안 컨볼루션을 사용하였으며, $\sigma_W \ge 1.0$인 SJD는 특히 긴 시퀀스 길이($\ge 6$)와 다양한 NFE 예산에 걸쳐 CANDI 및 MDLM보다 높은 유효 단어 생성 수를 기록했습니다.
• Sudoku (제약 충족): 행/열/박스 의존성을 인코딩하기 위해 제약 그래프 커널을 사용하였으며, SJD는 아이덴티티 커널 하이브리드(CADD)를 크게 압도했습니다. CADD는 성능이 양극화(종종 무작위 수준으로 붕괴)되는 반면, SJD는 지속적으로 퍼즐을 해결하며 95.65%의 정확도를 달성하고 훈련 단계에서 4배 빠르게 수렴했습니다.

의의 및 주장

본 논문은 스티키 점프 확산이 마스크, 연속, 하이브리드 확산을 통합하는 프로세스 수준의 설명을 제공한다고 주장합니다.

• 이론적 통합: 이산적 좌표와 연속적 좌표가 수동으로 결합해야 하는 별개의 엔티티가 아니라, 약속 비율과 목적지가 플럭스 균형에 의해 고정되는 단일한 역방향 상태임을 입증합니다.
• 설계 공간의 확장: 언스킹 커널을 명시적인 설계 공간으로 다룸으로써, 연구자들이 역방향 샘플러의 근본적인 아키텍처를 변경하지 않고도 구조적 사전 지식(공간적 국소성 또는 논리적 제약 등)을 순방향 오염 역학에 직접 인코딩할 수 있게 합니다.
• 실용적 효율성: SJD 특유의 구조는 오직 분석적 재가중치 부여와 폐쇄형 쿼드러처를 통해서만 반영되므로, 훈련 및 샘플링 비용은 기존 하이브리드 모델과 동일합니다.

저자들은 이러한 통합적 관점이 확산 모델을 확장하는 것을 더 용이하게 만든다고 결론짓습니다. 새로운 오염, 스케줄, 약속 규칙은 순방향 역학으로 지정될 수 있으며, 시간 역전 정리를 통해 무료로 역방향 모델을 얻을 수 있습니다.