PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion models for missing value imputation in tabular data

Shuhan Zheng, Nontawat Charoenphakdee

arXiv:2210.17128v2cs.LGcs.AI

TL;DR

Missing-value imputation with diffusion models had received limited attention for tabular data, which may mix categorical and numerical features. The paper proposes TabCSDI, a CSDI-based approach with three ways to handle categorical variables. It reports competitive performance against established imputation methods, particularly for numerical imputation, and finds that categorical embedding choice can affect results.

  • Problem

    Although diffusion models had been used in other domains, the authors knew of no diffusion model proposed for tabular imputation, where categorical and numerical features must be handled together.

    Source: §1 · p. 2

  • Method

    TabCSDI adapts CSDI [18] for tabular data and evaluates one-hot encoding, analog bits encoding, and feature tokenization for categorical variables.

    Source: §1 · p. 2, §3.2 · p. 3

  • Results

    TabCSDI achieves competitive performance against well-known imputation methods and works particularly well for numerical imputation.

    Source: §5 · p. 4

  • Takeaways & Limitations

    Categorical embedding choice can affect performance, with feature tokenization performing evidently better than one-hot and analog bits in the Census dataset.

    Source: §1 · p. 2, §5 · p. 4

Abstract

from arXiv · show

Missing value imputation in machine learning is the task of estimating the missing values in the dataset accurately using available information. In this task, several deep generative modeling methods have been proposed and demonstrated their usefulness, e.g., generative adversarial imputation networks. Recently, diffusion models have gained popularity because of their effectiveness in the generative modeling task in images, texts, audio, etc. To our knowledge, less attention has been paid to the investigation of the effectiveness of diffusion models for missing value imputation in tabular data. Based on recent development of diffusion models for time-series data imputation, we propose a diffusion model approach called "Conditional Score-based Diffusion Models for Tabular data" (TabCSDI). To effectively handle categorical variables and numerical variables simultaneously, we investigate three techniques: one-hot encoding, analog bits encoding, and feature tokenization. Experimental results on benchmark datasets demonstrated the effectiveness of TabCSDI compared with well-known existing methods, and also emphasized the importance of the categorical embedding techniques.

1 Introduction

Missing values are common in real-world datasets and can arise from privacy concerns, collection difficulty, or differing examinations. The paper identifies an unaddressed opportunity: applying diffusion models to tabular imputation, and proposes TabCSDI for that setting.

1 Introduction

  • Missing values can arise from human error, privacy concerns, or difficult data collection; census privacy and differing healthcare examinations are examples.The paper distinguishes missing completely at random, missing at random, and missing not at random, referencing Rubin and Van Buuren.Source: §1 · p. 1
  • Imputation methods are grouped into iterative conditional-feature estimation and deep generative approaches, with MICE as a well-known iterative method.The iterative procedure repeats feature-wise conditional estimation until a convergence criterion is met.Source: §1 · p. 1
  • Diffusion models had not, to the authors’ knowledge, been proposed for tabular missing-value imputation, despite their effectiveness in several other domains.The cited domains include computer vision, time-series data, chemistry, and natural language processing.Source: §1 · p. 2
  • TabCSDI adapts CSDI [18] to support both categorical and numerical features, achieving competitive performance against existing iterative and generative methods.The authors also report that categorical embedding choice can affect performance.Source: §1 · p. 2

2 Problem formulation

The paper formulates imputation as learning a function that maps incomplete d-dimensional inputs to d real-valued outputs. Evaluation compares imputed missing entries with their ground truths using metrics suited to numerical or categorical features.

2 Problem formulation

  • In the inductive setting, the imputation function maps an input with missing, numerical, or categorical features to d real-valued outputs.The input space is (R ∪ {∅})^d, where ∅ denotes a missing value.Source: §2 · p. 2
  • A desirable imputation function replaces missing values with reasonable values, and evaluation uses test inputs and corresponding ground truths.The test-input passage introduces the test data and ground truths used to assess the function.Source: §2 · p. 2
  • Numerical features are scored with RMSE and categorical features with error rate over missing entries.The error-rate formula uses an indicator that returns 1 when its condition holds and 0 otherwise.Source: §2 · p. 2

3 TabCSDI: Conditional Score-based Diffusion Models for Tabular data

TabCSDI adapts conditional score-based diffusion for tabular imputation, where inputs may combine observed and missing values, numerical features, and categorical features.

3 TabCSDI: Conditional Score-based Diffusion Models for Tabular data

  • TabCSDI applies a diffusion-model approach to missing-value imputation in tabular data by adapting CSDI [18].The method section describes the original model and its modification for tabular inputs.Source: §3 · p. 2

3.1 Conditional Score-based Diffusion Model (CSDI)

  • CSDI conditions on observed values and models the unobserved target values, while its reverse denoising process predicts only that target part.The forward process adds noise iteratively; the reverse process denoises, and the tabular architecture removes the temporal transformer in favor of a residual connection of transformer encoder and multilayer perceptron.Source: §3.1 · pp. 2–3

3.2 Handling categorical variables

  • TabCSDI handles categorical features with one-hot encoding, analog bits, or feature tokenization; tokenization embeds numerical and categorical variables at a shared length.One-hot decoding selects the largest output; analog bits threshold outputs, while tokenization recovers numerical values from embeddings and assigns categories by nearest embedding.Source: §3.2 · p. 3, Figure 1 · p. 3

4 Experimental results

Across the evaluated datasets, TabCSDI performs well on numerical imputation, while categorical handling results vary by dataset and encoding. The experiments compare it with four established imputation baselines on mixed-variable and numerical-only data.

Experimental setup

  • The experiments used seven datasets and compared TabCSDI with Mean / Mode, two MICE variants, MissForest, and GAIN.The datasets included three mixed-variable and four pure numerical datasets.Source: §4 · pp. 3–4

Results

  • TabCSDI has the best performance on 3 of 4 pure numerical datasets and the best RMSE on 5 of 7 datasets overall.The paper also reports that TabCSDI achieved the lowest RMSE on Diabetes and Census among the mixed-variable datasets.Source: Table 2 · p. 4, §4 · p. 4
  • On Diabetes and Census, TabCSDI reached the lowest RMSE, while MissForest reached the lowest error rate.In Census, FT had the lowest error rate among TabCSDI’s categorical-handling schemes, and analog bits outperformed one-hot encoding.Source: §4 · p. 4, Table 1 · p. 4

Discussions

  • FT’s advantage for categorical handling was evident in Census, the only mixed-variable dataset with multiple categories, possibly by avoiding column imbalance.The authors note that one-hot and analog-bits encodings can use different numbers of columns depending on category count.Source: §4 · p. 4

5 Conclusions and future work

The paper concludes that TabCSDI achieves competitive imputation performance, particularly for numerical variables, and that feature tokenization performs better than the other tested categorical encodings on Census. It identifies inference time, architecture, and loss-function analysis as future work.

Conclusions

  • TabCSDI achieved competitive performance with established imputation methods and worked particularly well for numerical-variable imputation.FT performed better than one-hot encoding and analog bits in the Census dataset.Source: §5 · p. 4

Future work

  • Future work includes investigating inference time, improving the model architecture, and theoretically analyzing the loss function.Source: §5 · p. 4

A Dataset information

Table 3 is identified as providing dataset information.

Dataset information

B Hyperparameters

The experiments used mostly default hyperparameter settings, with dataset-dependent choices for batch size, epochs, layer count, and reverse steps.

Hyperparameters

  • TabCSDI used 2 layers for COVID-19 and 4 for the other datasets, with 64 channels and 4 transformer encoder heads.Source: B Hyperparameters · p. 7
  • The initial learning rate was 0.0005, and Adam used MultiStepLR with 0.1 decay at 25%, 50%, 75%, and 90% of total epochs.Source: B Hyperparameters · p. 7
  • The reverse process used 150 steps for Diabetes, COVID-19, and numerical datasets, and 100 steps for Census.The authors used default settings for most hyperparameters and did not heavily tune the models.Source: B Hyperparameters · p. 7
Loading 2210.17128v2…