TabTreeFormer: Tabular Data Generation Using Hybrid Tree-Transformer

This repository contains the official implementation of the paper "TabTreeFormer: Tabular Data Generation Using Hybrid Tree-Transformer", published at The 29th International Conference on Artificial Intelligence and Statistics (AISTATS 2026). [Paper]

Pre-requisites

• Python>=3.9 installed.
• pip install -r requirements.txt.

Usage

To train,

python main.py train -d DATA_PATH -t TARGET_COLUMN -p TASK_TYPE -o OUT_DIR

After training, to generate,

python main.py sample -c OUT_DIR -n N_ROWS -o OUT_CSV_PATH

For instance, to train and sample iris dataset (one can get the dataset by sklearn.datasets.load_iris), one can run the following:

python main.py train -d iris.csv -t target -p mult -o out
python main.py sample -c out -n 150 -o synthetic-iris.csv

Evaluation

Our custom-implemented metrics are provided in evaluate.py. Dataset split and metadata specification should be prepared beforehand.

Machine Learning Efficacy

import pandas as pd
from evaluate import TabTransformer, ml_models
real_data: pd.DataFrame = ...
real_test_data: pd.DataFrame = ...
syn_data: pd.DataFrame = ...
model_name: str = ...  # "lin" / "rf" / "xgb"
ml_model = ml_models[model_name]
transformer = TabTransformer(target_column=..., is_clf=..., onehot=ml_model.need_onehot())
ml_model(transformer)
real_scores = ml_model.compute(real_data, real_test_data)
syn_scores = ml_model.compute(syn_data, real_test_data)

Distance to Closet Record

import pandas as pd
from evaluate import DistanceToClosestRecord
dcr_calculator = DistanceToClosestRecord(target_column=..., is_clf=...)
real_data: pd.DataFrame = ...
real_test_data: pd.DataFrame = ...
syn_test_data: pd.DataFrame = ...  # synthetic data generated from generator trained on real_data but size is same as test_data
dcr_calculator.fit(real_data, real_test_data)
score = dcr_calculator.compute(syn_test_data)

Fidelity Scores

We directly used sdv's quality report results for this. Code is omitted.

Citation

If you use this code in your research, please cite our paper:

@inproceedings{li2026tabtreeformer,
  title = {TabTreeFormer: Tabular Data Generation Using Hybrid Tree-Transformer},
  author = {Li, Jiayu and Zhao, Bingyin and Zhao, Zilong and Javaid, Uzair and Sikdar, Biplab},
  booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics},
  year = {2026},
  url = {https://openreview.net/forum?id=4Kj1Xoqi23}
}

License

This code is released for non-commercial research and academic purposes only. See LICENSE for details.