# mixgb: Multiple Imputation Through XGBoost

#### Yongshi Deng

#### 2025-04-07

## Introduction

Mixgb offers a scalable solution for imputing large datasets using XGBoost, subsampling and predictive mean matching. Our method utilizes the capabilities of XGBoost, a highly efficient implementation of gradient boosted trees, to capture interactions and non-linear relations automatically. Moreover, we have integrated subsampling and predictive mean matching to minimize bias and reflect appropriate imputation variability. Our package supports various types of variables and offers flexible settings for subsampling and predictive mean matching. We also include diagnostic tools for evaluating the quality of the imputed values.

## Impute missing values with `mixgb`

We first load the `mixgb` package and the `nhanes3_newborn` dataset, which contains 16 variables of various types (integer/numeric/factor/ordinal factor). There are 9 variables with missing values.

    library(mixgb)
    str(nhanes3_newborn)
    #> tibble [2,107 × 16] (S3: tbl_df/tbl/data.frame)
    #>  $ HSHSIZER: int [1:2107] 4 3 5 4 4 3 5 3 3 3 ...
    #>  $ HSAGEIR : int [1:2107] 2 5 10 10 8 3 10 7 2 7 ...
    #>  $ HSSEX   : Factor w/ 2 levels "1","2": 2 1 2 2 1 1 2 2 2 1 ...
    #>  $ DMARACER: Factor w/ 3 levels "1","2","3": 1 1 2 1 1 1 2 1 2 2 ...
    #>  $ DMAETHNR: Factor w/ 3 levels "1","2","3": 3 1 3 3 3 3 3 3 3 3 ...
    #>  $ DMARETHN: Factor w/ 4 levels "1","2","3","4": 1 3 2 1 1 1 2 1 2 2 ...
    #>  $ BMPHEAD : num [1:2107] 39.3 45.4 43.9 45.8 44.9 42.2 45.8 NA 40.2 44.5 ...
    #>   ..- attr(*, "label")= chr "Head circumference (cm)"
    #>  $ BMPRECUM: num [1:2107] 59.5 69.2 69.8 73.8 69 61.7 74.8 NA 64.5 70.2 ...
    #>   ..- attr(*, "label")= chr "Recumbent length (cm)"
    #>  $ BMPSB1  : num [1:2107] 8.2 13 6 8 8.2 9.4 5.2 NA 7 5.9 ...
    #>   ..- attr(*, "label")= chr "First subscapular skinfold (mm)"
    #>  $ BMPSB2  : num [1:2107] 8 13 5.6 10 7.8 8.4 5.2 NA 7 5.4 ...
    #>   ..- attr(*, "label")= chr "Second subscapular skinfold (mm)"
    #>  $ BMPTR1  : num [1:2107] 9 15.6 7 16.4 9.8 9.6 5.8 NA 11 6.8 ...
    #>   ..- attr(*, "label")= chr "First triceps skinfold (mm)"
    #>  $ BMPTR2  : num [1:2107] 9.4 14 8.2 12 8.8 8.2 6.6 NA 10.9 7.6 ...
    #>   ..- attr(*, "label")= chr "Second triceps skinfold (mm)"
    #>  $ BMPWT   : num [1:2107] 6.35 9.45 7.15 10.7 9.35 7.15 8.35 NA 7.35 8.65 ...
    #>   ..- attr(*, "label")= chr "Weight (kg)"
    #>  $ DMPPIR  : num [1:2107] 3.186 1.269 0.416 2.063 1.464 ...
    #>   ..- attr(*, "label")= chr "Poverty income ratio"
    #>  $ HFF1    : Factor w/ 2 levels "1","2": 2 2 1 1 1 2 2 1 2 1 ...
    #>  $ HYD1    : Ord.factor w/ 5 levels "1"<"2"<"3"<"4"<..: 1 3 1 1 1 1 1 1 2 1 ...
    colSums(is.na(nhanes3_newborn))
    #> HSHSIZER  HSAGEIR    HSSEX DMARACER DMAETHNR DMARETHN  BMPHEAD BMPRECUM
    #>        0        0        0        0        0        0      124      114
    #>   BMPSB1   BMPSB2   BMPTR1   BMPTR2    BMPWT   DMPPIR     HFF1     HYD1
    #>      161      169      124      167      117      192        7        0

To impute this dataset, we can use the default settings. The default number of imputed datasets is `m = 5`. Note that we do not need to convert our data into dgCMatrix or one-hot coding format. Our package will automatically convert it for you. Variables should be of the following types: numeric, integer, factor or ordinal factor.

    # use mixgb with default settings
    imputed.data <- mixgb(data = nhanes3_newborn, m = 5)

### Customize imputation settings

We can also customize imputation settings:

- The number of imputed datasets `m`

- The number of imputation iterations `maxit`

- XGBoost hyperparameters and verbose settings. `xgb.params`, `nrounds`, `early_stopping_rounds`, `print_every_n` and `verbose`.

- Subsampling ratio. By default, `subsample = 0.7`. Users can change this value under the `xgb.params` argument.

- Predictive mean matching settings `pmm.type`, `pmm.k` and `pmm.link`.

- Whether ordinal factors should be converted to integer (imputation process may be faster) `ordinalAsInteger`

- Whether or not to use bootstrapping `bootstrap`

- Initial imputation methods for different types of variables `initial.num`, `initial.int` and `initial.fac`.

- Whether to save models for imputing newdata `save.models` and `save.vars`.

    # Use mixgb with chosen settings
    params <- list(
      max_depth = 5,
      subsample = 0.9,
      nthread = 2,
      tree_method = "hist"
    )

    imputed.data <- mixgb(
      data = nhanes3_newborn, m = 10, maxit = 2,
      ordinalAsInteger = FALSE, bootstrap = FALSE,
      pmm.type = "auto", pmm.k = 5, pmm.link = "prob",
      initial.num = "normal", initial.int = "mode", initial.fac = "mode",
      save.models = FALSE, save.vars = NULL,
      xgb.params = params, nrounds = 200, early_stopping_rounds = 10, print_every_n = 10L, verbose = 0
    )

### Tune hyperparameters

Imputation performance can be affected by the hyperparameter settings. Although tuning a large set of hyperparameters may appear intimidating, it is often possible to narrowing down the search space because many hyperparameters are correlated. In our package, the function `mixgb_cv()` can be used to tune the number of boosting rounds - `nrounds`. There is no default `nrounds` value in `XGBoost,` so users are required to specify this value themselves. The default `nrounds` in `mixgb()` is 100. However, we recommend using `mixgb_cv()` to find the optimal `nrounds` first.

    params <- list(max_depth = 3, subsample = 0.7, nthread = 2)
    cv.results <- mixgb_cv(data = nhanes3_newborn, nrounds = 100, xgb.params = params, verbose = FALSE)
    cv.results$evaluation.log
    #>      iter train_rmse_mean train_rmse_std test_rmse_mean test_rmse_std
    #>     <num>           <num>          <num>          <num>         <num>
    #>  1:     1       47.953069     0.02887792      47.943522     0.2022316
    #>  2:     2       33.696521     0.02157801      33.690149     0.1994458
    #>  3:     3       23.738196     0.01598159      23.749785     0.1977378
    #>  4:     4       16.769628     0.02555824      16.782115     0.2085698
    #>  5:     5       11.919335     0.02170435      11.938467     0.2127997
    #>  6:     6        8.581559     0.03198676       8.616413     0.2644566
    #>  7:     7        6.295901     0.03219033       6.336111     0.3095945
    #>  8:     8        4.786384     0.05469584       4.851058     0.3579966
    #>  9:     9        3.800631     0.06778147       3.900530     0.3995183
    #> 10:    10        3.194553     0.08204794       3.336379     0.4238159
    #> 11:    11        2.828824     0.09147260       3.012366     0.4541305
    #> 12:    12        2.622492     0.09779780       2.835412     0.4676225
    #> 13:    13        2.502706     0.09922714       2.748699     0.4574018
    #> 14:    14        2.420211     0.09825536       2.717810     0.4556033
    #> 15:    15        2.375800     0.09800962       2.689850     0.4593614
    #> 16:    16        2.334678     0.10517289       2.676588     0.4522710
    #> 17:    17        2.315079     0.10752478       2.665412     0.4504448
    #> 18:    18        2.289319     0.10471521       2.673480     0.4397764
    #> 19:    19        2.270184     0.10218612       2.670527     0.4378170
    #> 20:    20        2.253590     0.10354572       2.665598     0.4403487
    #> 21:    21        2.231716     0.10002020       2.670749     0.4395124
    #> 22:    22        2.220662     0.10168917       2.673183     0.4429654
    #> 23:    23        2.200946     0.10075459       2.706757     0.4538523
    #> 24:    24        2.186310     0.09924333       2.706875     0.4540316
    #> 25:    25        2.167716     0.09681875       2.713598     0.4604745
    #> 26:    26        2.156608     0.10162004       2.726646     0.4567566
    #> 27:    27        2.139597     0.10030344       2.732775     0.4519637
    #>      iter train_rmse_mean train_rmse_std test_rmse_mean test_rmse_std
    cv.results$response
    #> [1] "BMPRECUM"
    cv.results$best.nrounds
    #> [1] 17

By default, `mixgb_cv()` will randomly choose an incomplete variable as the response and build an XGBoost model with other variables as explanatory variables using the complete cases of the dataset. Therefore, each run of `mixgb_cv()` will likely return different results. Users can also specify the response and covariates in the argument `response` and `select_features` respectively.

    cv.results <- mixgb_cv(
      data = nhanes3_newborn, nfold = 10, nrounds = 100, early_stopping_rounds = 1,
      response = "BMPHEAD", select_features = c("HSAGEIR", "HSSEX", "DMARETHN", "BMPRECUM", "BMPSB1", "BMPSB2", "BMPTR1", "BMPTR2", "BMPWT"), xgb.params = params, verbose = FALSE
    )

    cv.results$best.nrounds
    #> [1] 19

Let us just try setting `nrounds = cv.results$best.nrounds` in `mixgb()` to obtain 5 imputed datasets.

    imputed.data <- mixgb(data = nhanes3_newborn, m = 5, nrounds = cv.results$best.nrounds)

## Inspect multiply imputed values

The `mixgb` package used to provide a few visual diagnostics functions. However, we have moved these functions to the `vismi` package, which provides a wide range of visualisation tools for multiple imputation.

For more details, please check the `vismi` package on GitHub Visualisation Tools for Multiple Imputation.