ComboStoc: 확산 생성 모델을 위한 조합적 확률성

기존 확산 모델은 고차원 데이터의 모든 속성을 동시에 학습하려다 보니 특정 조합의 데이터를 충분히 학습하지 못하는 한계가 있었다. 이 논문은 데이터의 각 차원이나 속성마다 서로 다른 시간 단계를 적용하는 비동기식 학습법을 통해 학습 효율을 높이고, 생성 시점에 정교한 부분별 제어를 가능하게 한다.

기존의 Diffusion Model은 이미지 전체를 하나의 시간 단계(t)에 맞춰 동시에 노이즈를 섞거나 제거한다. 이는 데이터가 이동하는 경로를 단순한 직선으로 제한하여, 고차원 공간에서 발생할 수 있는 다양한 데이터 속성 간의 조합을 충분히 탐색하지 못하게 만든다. 결과적으로 학습 데이터가 부족한 영역에서는 생성 품질이 급격히 저하되는 문제가 발생한다.

ComboStoc은 이 시간 단계 t를 스칼라 값이 아닌 데이터와 동일한 크기의 텐서(Tensor)로 확장한다. 즉, 이미지의 각 패치나 3D 모델의 각 부품이 서로 다른 t 값을 가지도록 설정하여, 모델이 데이터의 모든 차원 조합을 골고루 학습하게 유도한다. 이는 모델이 특정 속성이 고정된 상태에서 다른 속성이 변화하는 관계를 더 깊이 이해하게 만든다.

이러한 비동기적 접근은 학습 효율을 높일 뿐만 아니라 추론 시에도 강력한 이점을 제공한다. 사용자가 특정 부분의 t 값을 조절함으로써, 이미지의 특정 영역은 유지하면서 나머지 영역만 새로 생성하거나 3D 모델의 특정 부품만 교체하는 식의 '단계별 제어(Graded Control)'가 가능해진다.

전체적인 접근 방식은 기존의 선형 보간(Linear Interpolant) 모델을 확장하는 것이다. 기존 식 x_t = (1-t)z + tx_1에서 스칼라 t를 데이터 x와 동일한 형상을 가진 텐서 t로 대체하여 x_t = (1-t)⊙z + t⊙x_1로 정의한다. 여기서 ⊙는 요소별 곱(Elementwise Product)을 의미한다.

핵심 메커니즘은 학습 시 각 차원의 t 값을 [0, 1] 범위에서 독립적으로 균일하게 샘플링하는 것이다. [입력값: 노이즈 z, 원본 데이터 x_1, 텐서 t] → [연산: 요소별 보간 수행] → [결과값: 비동기적으로 오염된 데이터 x_t] → [의미: 모델이 데이터의 모든 부분적 조합을 학습할 수 있는 상태가 됨].

추론 시 발생하는 오프-대각선 드리프트(Off-diagonal Drift) 문제를 해결하기 위해 드리프트 최소화 기법을 도입했다. 이는 모델이 예측한 속도 벡터에 보정 성분을 더해 궤적을 다시 목표 데이터 방향으로 유도하는 방식이다. [입력값: 현재 상태 x_t, 예측 속도 v] → [연산: 목표 지점 x_1과의 오차를 계산하여 음의 그래디언트 방향으로 보정] → [결과값: 보정된 속도 벡터] → [의미: 비동기 추론 시에도 정확한 데이터 생성을 보장함].

ImageNet 256x256 벤치마크에서 ComboStoc-XL 모델은 800K 학습 단계에서 FID 11.41을 기록하여, 동일 조건의 SiT-XL(12.6)과 DiT-XL(14.3)보다 우수한 성능을 보였다. 특히 400K 단계에서도 FID 15.69를 달성해 기존 모델들의 800K 시점 성능에 근접하는 빠른 수렴 속도를 증명했다.

구조화된 3D 형상 생성 실험(PartNet 데이터셋)에서는 기존 모델들이 유의미한 형태를 생성하지 못하는 상황에서도 ComboStoc은 정확한 부품 조합과 형태를 생성해냈다. FPD(Frechet Point Distance) 지표에서 baseline인 'none' 설정이 7.99를 기록한 반면, 모든 속성을 비동기화한 'all' 설정은 4.04로 성능이 약 2배 향상됐다.

Ablation Study를 통해 시간 단계의 비동기화 수준이 높을수록(패치 단위 < 벡터 단위 < 전체 요소 단위) 생성 품질이 일관되게 향상됨을 확인했다. 또한 데이터가 부족한 상황(1,000장의 이미지)에서도 ComboStoc은 SiT보다 훨씬 빠르고 안정적으로 고품질 이미지를 생성하는 강건함을 보였다.

ComboStoc의 아키텍처는 SiT(Scalable Interpolant Transformer)를 기반으로 하되, 타임스텝 임베딩 모듈을 텐서 입력을 처리할 수 있도록 수정했다. 기존의 스칼라 t를 주입하던 방식에서 탈피하여, 각 차원별 t 값을 주파수 인코딩한 후 패치 임베딩과 동일한 방식으로 처리하여 트랜스포머 블록에 전달한다.

수학적으로 ComboStoc은 조건부 벡터 필드 u(x|x0, x1)를 직사각형 서브스페이스(Rectangular Subspace) 전체로 확장한다. 이는 기존 Flow Matching이 x0와 x1을 잇는 대각선 경로(Diagonal)만 다루던 것을 보완하여, 주변 영역에 대한 모델의 근사 능력을 강화한다. 연속 방정식(Continuity Equation) 분석을 통해 이러한 비동기 샘플링이 여전히 유효한 확률 경로를 생성함을 이론적으로 입증했다.

학습 효율성 측면에서 텐서 기반 타임스텝 임베딩으로 인해 GFLOPs는 약 48% 증가(237.34 -> 352.46)하지만, 실제 추론 속도(Inference Speed)는 병렬 연산 덕분에 50ms에서 48ms로 거의 차이가 없거나 오히려 소폭 개선되는 효율적인 구조를 가진다.

데이터 차원 간의 독립성이 강한 벡터 공간에서는 조합적 확률성의 효과가 제한적일 수 있다. 또한 현재 3D 형상 생성 모델에서는 부품의 회전(Rotation) 속성은 고정된 것으로 가정하고 있어, 향후 회전 속성까지 포함한 완전한 자유도의 생성이 과제로 남아있다.