Geometry Adaptive Counterfactual Distribution Learning with Diffusion-Guided Smoothing

2026-05-25Machine Learning

Machine Learning
AI summary

The authors look at how to better estimate what could have happened (counterfactual outcomes) when dealing with very complex, high-dimensional data that often lies close to simpler, lower-dimensional shapes. They point out that usual smoothing methods don't work well because they treat all directions in the data the same, causing errors and instability. To fix this, they introduce new estimators that adapt the smoothing process based on the data's underlying geometry using diffusion techniques, which reduces bias and aligns smoothing to the data's true shape. Their theory shows that these new methods perform better, especially as data complexity grows, and they back this up with experiments using facial imagery data.

Counterfactual Distribution LearningHigh-Dimensional DataIsotropic SmoothingDiffusion-Guided EstimatorsSemiparametric DebiasingScore FunctionNuisance AdjustmentLocal GeometryAsymptotic AnalysisEffective Dimension
Authors
Kwangho Kim
Abstract
We study counterfactual distribution learning for high-dimensional outcomes whose counterfactual law may concentrate near lower-dimensional structure. Standard isotropic smoothing treats all ambient directions equally, leading to unfavorable scaling and unstable local inference. We propose two diffusion-guided estimators based on semiparametric debiasing: diffusion-informed smoothing for counterfactual densities and diffusion-informed score smoothing for counterfactual scores. The estimators combine causal nuisance adjustment with geometry-adaptive localization driven by diffusion score information, removing first-order nuisance bias while aligning smoothing with local outcome geometry. We establish asymptotic expansions, risk bounds, and inference procedures for smoothed density and score-based targets, with ambient density inference obtained under additional approximation conditions. Under structural geometry conditions, the leading stochastic error is governed by an effective dimension induced by the diffusion-guided kernel, rather than by the ambient dimension. Semi-synthetic experiments based on CelebA show steeper error decay for geometry-adaptive methods, supporting the proposed effective-dimension theory.