Conditioning Tree-Based Diffusions and Flows for Probabilistic Tabular Regression

기술 요약: 확률적 표 형식 회귀를 위한 조건부 트리 기반 확산 및 흐름 모델링

문제 정의
표 형식 데이터(tabular data)에 대한 확률적 회귀는 그래디언트 부스팅 트리(gradient-boosted trees)의 예측 강점과 보정(calibration), 의사결정 및 불확실성 정량화에 적합한 완전한 예측 분포를 결합하는 모델을 필요로 한다. 트리 기반 확산 모델(예: Treeffuser)이 노이즈가 섞인 조건부 밀도의 스코어(score)를 추정하도록 LightGBM 회귀 모델을 학습시킴으로써 해결책으로 등장했으나, 이들은 현재 신경망 확산 설정에서 유래된 설계 기본값(design defaults)을 그대로 물려받고 있다. 노이징 경로(noising path), 파라미터화, 학습 분포, 피처(features), 샘플러를 포함한 이러한 기본값들은 히스토그램 기반 트리 학습기에게 최적이 아닐 수 있다. 본 논문은 이러한 기본값들이 구속 요인으로 작용한다고 주장한다. 즉, 그래디언트 부스팅 앙상블이 해결하는 특정 지도 학습 회귀 문제의 성격은 이러한 조건부 선택들이 어떻게 구성되느냐에 의해 결정된다.

방법론: DiffGBM
저자는 트리 기반 확산의 조건부 선택 사항들을 고정된 기본값이 아닌 명시적이고 튜닝 가능한 축(axis)으로 취급하는 Treeffuser의 확장형인 DiffGBM을 소개한다. 방법론은 두 가지 주요 축을 따라 개발되었다:

1. 스코어 측면의 수정 ( "Score-Flex" 공간):
   저자는 공유된 LightGBM 표면 위에서 스코어 기반 확산을 수행하기 위한 공동 튜닝 가능한 레시피를 제안한다. 주요 수정 사항은 다음과 같다:

   • 조건부 평균 잔차화(Conditional-Mean Residualization): 먼저 교차 검증된 LightGBM 평균 예측기 $\hat{\mu}(x)$를 적합시킨다. 확산은 $y$가 아닌 잔차 $y - \hat{\mu}(x)$에 적용되며, 이를 통해 조건부 위치(conditional location)를 스코어 회귀 타겟에서 분리한다.
   • EDM 프리컨디셔닝(EDM Preconditioning): EDM 원칙에 따라, 모델은 단위 분산을 갖도록 스케일링된 디노이징 타겟 $D_\theta(y_t, t)$를 예측하며, 이를 통해 스코어를 대수적으로 재구성한다. 이는 트리를 위한 회귀 타겟을 안정화하여, 노이즈 수준에 따라 노이즈 예측 크기를 외삽(extrapolate)할 필요를 방지한다.
   • Log-Sigma 시간 샘플링 및 피처: 균등한 시간 샘플링 대신, 학습 분포는 $\log \sigma(t) \sim \mathcal{N}(-1.2, 1.2^2)$를 따르도록 한다. 결정적으로, 입력 벡터에 명시적인 $\log \sigma(t)$ 피처를 추가하여 트리 분할(split)이 노이즈 체제(noise regime)에 직접적으로 조건화될 수 있도록 한다.
   • 손실 가중치(Loss Weighting) 및 히스토그램 해상도: 이 요소들은 모델의 용량(capacity)을 어디에 집중할지 제어하기 위해 튜닝 가능한 차원으로 노출된다.

2. 가우시안 경로 흐름 매칭 (Gaussian-Path Flow Matching, FM):
   저자는 $p(y|x)$를 위한 가우시안 경로 흐름 매칭 트레이너를 구현한다.

   • 속도장 학습(Velocity Field Learning): 스코어를 직접 학습하는 대신, 모델은 그래디언트 부스팅 트리를 통해 속도장 $v_\theta$를 학습한다.
   • 대수적 스코어 복구: 임의의 가우시안 경로에 대해, 조건부 스코어는 속도로부터 대수적으로 복구될 수 있다(Lemma 1). 이를 통해 모델이 원한다면 확률적 인터폴런트(stochastic-interpolant) 샘플러를 사용할 수 있게 한다.
   • 샘플러 효율성: 이 접근 방식은 결정론적이고 단계가 적은 ODE 샘플러(특히 5단계 Heun ODE)의 사용을 가능하게 하여, 확률적 SDE에 비해 추론 지연 시간을 크게 줄인다.
   • 경로 선택: 저자는 선형(linear), 삼각 함수(trigonometric), 분산 보존(variance-preserving, VP) 경로를 비교한다. VP 경로는 입력을 단위 스케일 근처에 유지하지만, 최적의 경로는 데이터셋과 잔차화(residualization) 사이의 상호작용에 따라 달라진다고 언급한다.

주요 기여

1. 조건부 트리 기반 예측을 위한 흐름 매칭: 그래디언트 부스팅 트리를 사용하여 속도장을 직접 학습하는 가우시안 경로 FM 트레이너를 제공하며, 이는 가장 저렴한 샘플러와 가장 잘 보정된 동작 지점을 제공한다.
2. 공동 튜닝 가능한 스코어 레시피: 스코어 측면의 수정 사항(잔차화, EDM 프리컨디셔닝, log-sigma 샘플링, 명시적 노이즈 피처, 손실 가중치, 히스토그램 해상도)을 튜able 축으로 노출했다. 이 "score-flex" 공간은 발표된 레시피를 하나의 특수한 사례로 포함하면서도 데이터셋별 최적화를 허용한다.
3. 회귀 조건화가 공유 병목 구간임: 논문은 두 축 모두에서의 성능 향상이 히스토그램 트리의 유한한 분할 해상도에 맞춰 유도된 지도 학습 문제를 더 잘 매칭하는 데서 기인한다고 상정한다. 제약 조건은 목적 함수의 모집단 수준 특성이 아니라, 모든 노이즈 수준을 동시에 처리하기 위해 요구되는 유한한 빈(bin) 할당이다.

실험 결과
방법론은 11개의 표준 표 형식 회귀 벤치마크(UCI 9개, scikit-learn 2개)에 대해 중첩 튜닝 프로토콜(튜닝을 위한 fold 0, 평가를 위한 folds 1–5)과 매칭된 40회 실험 예산(trial budget)을 사용하여 평가되었다.

• 정확도 (Score-Flex): 공동 튜닝된 score-flex 레시피는 가장 우수한 집계 CRPS skill score(기존 베이스라인 0.699 대비 0.725)와 정규화된 CRPS(1.248 대비 1.106)를 달성했다. 모든 데이터셋에서 기존 베이스라인을 압도했다 (paired Wilcoxon $p < 10^{-3}$).
• 보정 및 속도 (Flow Matching): DiffGBM-FM 변형은 가장 잘 보정된 DiffGBM 행(50/90% 수준에서 가장 낮은 커버리지 에러와 가장 높은 PIT 통과율)이었으며 가장 빠른 샘플러였다. 경쟁력 있는 CRPS(0.707)를 유지하면서도 기존 베이스라인보다 5.2배 빨랐다 (총 196.3초 대비 61.3초).
• 절제 연구 (Ablation Studies):
  • 다른 변경 없이 스코어 모델에 Euler SDE 대신 Heun ODE를 사용했을 때 정확도가 감소했는데, 이는 샘플러가 자유로운 속도 향상 수단이 아님을 나타낸다.
  • EDM 프리컨디셔닝이 스코어 측면 성능에서 가장 큰 단일 도약을 제공했다.
  • 잔차화(Residualization)는 Flow Matching에 필수적이었다. 잔차화가 없으면 FM 성능이 크게 저하되었다.
  • 확률적 샘플러($\epsilon > 0$)는 집계 CRPS나 지연 시간 측면에서 결정론적 샘플러($\epsilon = 0$)를 파레토 우위(Pareto-dominate)로 압도하지 못했으나, 높은 불확실성 영역에서의 꼬리 커버리지(tail coverage)를 개선할 수는 있었다.

의의 및 주장
논문은 신경망 확산 설정에서 물려받은 설계 기본값이 트리 기반 확률적 회귀의 주요 병목 현상이라고 주장한다. 노이징 경로, 파라미터화, 학습 분포를 튜닝 가능한 조건부 선택으로 취급함으로써, DiffGBM은 정확도와 보정 모두에서 상당한 개선을 이끌어낸다.

저자는 "score-flex"와 "flow-matching" 접근 방식이 상호 보완적임을 강조한다:

• Score-Flex는 정확도를 위한 선택이다. 데이터가 증가함에 따라 우위를 점하며 가장 큰 데이터셋에서 승리하지만, 확률적 샘플러로 인해 지연 시간이 길고 중앙 커버리지가 느슨해질 수 있다.
• Flow Matching은 빠르고 잘 보정된 선택이다. 잔차화된 소단계(few-step) ODE가 가장 강력한 소규모~중규모 데이터셋에서 승리한다.

본 논문은 단일한 범용 레시피는 존재하지 않으며, 최적의 구성은 데이터셋의 특성(예: ct_slices와 같이 결정론적인 데이터셋은 더 미세한 히스토그램 빈과 노이즈 예측을 선호함)에 달려 있다고 결론짓는다. 이 연구는 히스토그램 분할의 유한한 해상도에 유도된 지도 학습 문제를 조건화하는 것이 트리 기반 확산 및 흐름 모델의 잠재력을 끌어내는 핵심임을 입증한다.