🤖 AI Summary
Existing continuous generative models struggle to effectively model mixed continuous-categorical tabular data, particularly when categorical variables exhibit severe class imbalance. This work proposes a unified framework based on logit coordinates, encoding categorical variables as smooth natural parameters and jointly modeling them with transformed numerical variables to construct Logit Flow Matching (Logit FM) and Logit diffusion models. The approach innovatively introduces a hybrid distribution discrepancy measure that disentangles categorical marginal error from conditional continuous Wasserstein error, and establishes imbalance-aware nonparametric convergence rates and stability bounds. Experiments on four real-world datasets demonstrate that Logit FM significantly improves distribution fidelity on three datasets and outperforms one-hot encoding under extreme class imbalance, while Logit diffusion consistently matches or surpasses existing one-hot diffusion methods.
📝 Abstract
Mixed continuous--categorical data pose a representation problem for continuous generative models. Flow Matching and Gaussian diffusion operate in Euclidean spaces, whereas categorical laws lie on probability simplices and may be highly imbalanced. We study a logit-coordinate framework that encodes categorical variables as smoothed natural parameters and combines them with transformed numerical variables. This yields common formulations of Logit Flow Matching and Logit Diffusion. We introduce a mixed-distribution discrepancy separating categorical marginal error from conditional continuous Wasserstein error, and derive stability bounds and imbalance-aware nonparametric rates linking vector-field or drift error to decoded mixed-distribution error. Controlled simulations show that scaled-logit coordinates improve or match one-hot coordinates, especially under severe rare-cell imbalance. Across four real-data benchmarks and ten splits per dataset, Logit FM improves the primary distributional metrics on three datasets and is comparable on Churn2; Block-Conditional Logit FM consistently improves the flat model; and Logit Diffusion generally improves over or matches One-Hot Diffusion.