Harpoon: Generalised Manifold Guidance for Conditional Tabular Diffusion

기술 요약: HARPOON – 조건부 표 형식 데이터 확산을 위한 일반화된 매니폴드 가이던스

문제 정의

특정 조건(제약 조건)을 준수하며 표 형식 데이터(tabular data)를 생성하는 것은 결측치 임퓨테이션(imputing missing values)이나 "what-if" 시나리오 시뮬레이션과 같은 응용 분야에서 매우 중요합니다. 확산 모델(diffusion models)이 생성 작업에서 최첨단(SOTA) 성능을 보여주고 있지만, 기존의 조건부 표 형식 확산 방법들은 다음과 같은 중대한 한계에 직면해 있습니다:

1. 일반화 능력 부족: 학습 시의 전략(예: 특정 제약 조건에 모델을 컨디셔닝하는 것)은 추론 시에 보지 못한 새로운 조건에 대해 일반화되지 못합니다.
2. 유연성 부족: 현재의 방법들은 비선형 제약 조건(예: $Age \ge 10$과 같은 특성 부등식)이나 복잡한 논리적 조합과 같은 다양한 추론 시 목적 함수를 처리하는 데 어려움을 겪습니다.
3. 이론적 공백: 데이터를 저차원 매니폴드 위에 존재하는 것으로 취급하는 최근의 기하학적 확산 접근법들은 연속적인 도메인과 평평한 기하 구조(flat geometries), 그리고 특정 추론 시 손실 함수(예: 제곱 오차)에 국한되어 있습니다. 이들은 표 형식 데이터의 혼합형 특성(이산형 및 연속형 특성)을 고려하지 못하거나 임의의 미분 가능한 목적 함수로 일반화되지 못합니다.

방법론: HARPOON

저자들은 무조건부 샘플을 매니폴드 기하 구조를 따라 유도하여, 재학습 없이도 다양한 조건을 충족하도록 하는 조건부 표 형식 확산 방법인 HARPOON을 소개합니다. 이 접근 방식은 매니폴드 이론을 표 형식 데이터 공간으로 확장하는 것을 기반으로 합니다.

이론적 토대

본 논문은 데이터 공간에서 작동하는 표 형식 확산에 관한 두 가지 핵심적인 이론적 결과를 입증합니다:

1. 직교 투영(Orthogonal Projection): 저자들은 표준 제곱 오차 손실로 학습된 확산 모델이 노이즈 수준이 0에 가까워질수록($\bar{\alpha}_t \to 1$) 기저의 데이터 매니폴드($M_0$)로의 직교 사상(orthogonal mapping)을 암묵적으로 학습함을 증명합니다. 이는 곡선형 기하 구조에서도 성립하며, 기존 연구의 평평한 매니폴드 가정을 완화합니다.
2. 접공간 그래디언트(Tangent-Space Gradients): 임의의 미분 가능한 추론 시 손실 함수($L_{inf}$)의 그래디언트가 "더러운" 추정치($\hat{x}_0$)에서의 데이터 매니폴드 접공간 내에 존재함을 증명합니다. 이는 제곱 오차 손실에 국한되었던 이전의 연구 결과들을 일반화한 것입니다.

이 이론은 디노이저(denoiser)가 매니폴드로의 직교 투영을 가리키는 나침반 역할을 하는 반면, 임의의 제약 조건으로부터 오는 그래디언트는 매니폴드 표면을 따라 접선 방향으로 업데이트를 조종하는 닻(anchor) 역할을 하여 샘플이 유효한 데이터 분포 껍질(shell) 내에 머물도록 한다는 것을 의미합니다.

알고리즘

HARPOON은 무조건부 디노이징 단계와 접선 방향의 그래디언트 수정을 교차하여 수행합니다:

1. 더러운 추정치(Dirty Estimate): 각 단계 $t$에서 모델은 순방향 노이징 과정을 역전시켜 "더러운" 추정치 $\hat{x}_0$를 계산합니다.
2. 접선 그래디언트(Tangential Gradient): 추론 시 손실 함수의 그래디언트($\nabla_{x_t} L_{inf}(\hat{x}_0, c)$)를 계산합니다. 이론적 증명에 따라, 이 그래디언트는 접공간 내로 제한됩니다.
3. 교차 업데이트(Interleaved Update): 샘플은 표준 무조건부 디노이징 단계를 거친 직후, 미리 계산된 그래디언트를 사용하여 접선 방향 업데이트를 즉시 수행합니다. 이 과정은 수렴할 때까지 반복됩니다.

결정적으로, 추론 시 손실 함수($L_{inf}$)는 학습 시의 손실과 다를 수 있습니다. 표 형식 데이터를 위해 저자들은 표준 MSE보다 이산적 특성에 더 적합한 희소성 유도 손실인 평균 절대 오차(MAE) 및 교차 엔트로피(CE)를 활용합니다.

주요 기여

1. 이론: 본 논문은 매니폴드와 표 형식 확산을 연결하는 첫 번째 이론적 결과를 제공하며, 확산 모델이 직교 투영체로 작동하고 임의의 미분 가능한 추론 시 목적 함수로부터의 그래디디언트가 곡선형 기하 구조와 혼합형 데이터를 포함하여 접공간에 존재함을 증명합니다.
2. 알고리즘: 재학습이나 파인튜닝 없이 추론 시 다양한 조건(임퓨테이션, 부등식, 논리적 제약 조건)에 적응할 수 있는 학습이 필요 없는(training-free) 가이던스 메커니즘인 HARPOON을 설계했습니다.
3. 실증적 검증: 다양한 데이터셋과 작업에서 매니폴드 인식 가이던스가 거절 샘플링(rejection sampling)이나 학습 시 컨디셔닝이 실패하는 부등식 제약 조건을 처리할 때 성능을 크게 향상시킨다는 것을 입증했습니다.

실험 결과

저자들은 여덟 개의 데이터셋(Adult, Default, Shoppers와 같은 UCI 벤치마크 포함)에 대해 두 가지 주요 작업으로 HARPOON을 평가했습니다.

1. 임퓨테이션 (Imputation)

• 성능: HARPOON은 다양한 결측 비율(25%, 50%, 75%)에 대해 SOTA 베이스라인(DiffPuter, Remasker, GReaT, MIRACLE, GAIN)과 비교하여 경쟁력 있거나 우수한 평균 제곱 오차(MSE) 및 정확도를 달 Vit했습니다.
• 강건성: 특정 학습 마스크에 의존하여 결측 비율이 높아질수록 성능이 저하되는 학습 시 방법들(예: Remasker)과 달리, HARPOON은 강건함을 유지했습니다.
• 손실 민감도: 절제 연구(Ablation studies) 결과, 추론 시 MAE(또는 MAE + CE)를 사용하는 것이 학습 시의 MSE 손실을 사용하는 것보다 우수함을 보여주었으며, 이는 매니폴드 기하 구조가 유연한 손실 선택을 가능하게 한다는 이론적 주장을 뒷받�습니다.

2. 부등식 및 논리적 제약 조건

• 제약 조건 충족: 범위 제약 조건(예: $Age \ge 50$) 및 범주형 제약 조건에 대해, 제약 조건을 임퓨테이션 문제로 취급하여 74~96%의 높은 위반율을 보인 DiffPuter 및 GReaT와 비교하여 HARPOON은 현저히 낮은 위반율(예: Adult/Default에서 <8%)을 달성했습니다.
• 복잡한 논리: HARPOON은 제약 조건의 결합(AND) 및 선택(OR)을 성공적으로 처리하였으며, 베이스라인들이 실패한 지점에서도 낮은 위반율과 높은 다운스트림 유틸리티(XGBoost 분류기 성능으로 측정)를 유지했습니다.
• 충실도(Fidelity): 제약 조건을 강제함에도 불구하고, HARPOON은 높은 샘플 충실도($\alpha$-score로 측정)와 유틸리티를 유지했습니다. 이는 매니폴드 가이드 업데이트가 샘플을 분포 밖으로 밀어내지 않았음을 나타냅니다.

효율성

HARPOON의 실행 시간은 추가적인 그래디언트 계산 단계로 인해 무조건부 확산 베이스라인(DiffPuter)보다 약 2배 정도 더 걸립니다. 그러나 LLM 기반(GReT) 또는 인과 그래프 기반(MIRACLE) 접근 방식보다는 훨씬 빠르며, 대규모 데이터셋에서도 5초 이내에 샘플링을 완료합니다.

의의 및 주장

본 논문은 HARPOON을 이론적 매니폴드 가이던스와 실용적인 표 형식 데이터 생성 사이의 가교로 포지셔닝합니다. 주요 의의는 **"학습은 단순하게 유지하고, 추론 시에 확장하라"**는 패러다임에 있습니다:

• 일반화: 모든 새로운 제약 조건마다 모델을 재학습할 필요를 없애, 단일 사전 학습된 모델이 보지 못한 조건에 적응할 수 있게 합니다.
• 이론적 엄밀성: 매니폴드 가이던스 이론을 연속적이고 평평한 도메인을 넘어 복잡한 혼합형 데이터의 기하 구조로 확장했습니다.
• 실용적 유용성: 의료, 금융 등 실제 도메인에서 흔히 발생하지만 표준 생성 모델이 분포 특성을 위반하지 않고 만족시키기 어려운 하드 제약 조건(부등식, 논리 조건)을 강제하는 강력한 메커니즘을 제공합니다.

저자들은 HARPOON이 일반적이고 이론적으로 근거가 확실하며, 적응력이 뛰어나고 효율적인 일반 가이던스 메커니즘을 제공함으로써 실제 세계의 표 형식 모델링에 대한 실질적인 이점을 제공한다고 결론짓습니다.