Skip to content

Add TabFMDataGenerator for synthetic tabular data generation - #83

Open
LckyLke wants to merge 2 commits into
google-research:mainfrom
LckyLke:synthetic-data-generation
Open

Add TabFMDataGenerator for synthetic tabular data generation#83
LckyLke wants to merge 2 commits into
google-research:mainfrom
LckyLke:synthetic-data-generation

Conversation

@LckyLke

@LckyLke LckyLke commented Jul 24, 2026

Copy link
Copy Markdown

Adds TabFMDataGenerator, a scikit-learn style generator that samples entirely new rows mimicking a reference dataset, using the existing pretrained classifier as the conditional-distribution engine. Closes #82.

Summary

The joint distribution over columns is factorized with the chain rule, p(x_1,...,x_d) = prod_j p(x_j | x_{<j}). Columns are visited in a seeded random (or user-supplied) order; each column is sampled for all rows at once from a TabFMClassifier fit with that column as the target and the already-sampled columns as features. This is a pure inference-time composition of the existing sklearn API — no model, architecture, or weight changes, and only the classification checkpoint is used.

Numerical columns reach fine resolution despite max_classes = 10 via hierarchical quantile refinement: n_bins ** n_levels equal-mass bins (100 by default, 1000 with n_levels=3), whose index is sampled digit by digit in base n_bins with already-sampled digits appended to the conditioning features. Every refinement level trains on all reference rows, so resolution grows exponentially while cost grows linearly. Values are drawn uniformly within the sampled bin.

Design notes

  • The regression head decodes to a single scalar (a point estimate), so it is never used; the binned classifier distribution plays the role of TabPFN's bar distribution.
  • Per sample() call: one classifier fit + one predict_proba per categorical column, n_levels of each per numeric column, all rows batched. With default settings (n_estimators=4), fit() makes no model forward passes (no NNLS, no calibration), so the cost is exactly the predict_proba calls.
  • Categorical columns beyond max_classes are modeled top-k with the tail merged into one class that is re-sampled from its empirical frequencies.
  • Constant columns never join the conditioning set (they would be dropped by UniqueFeatureFilter).
  • Sampling temperature t sharpens (< 1) or flattens (> 1) each conditional; dtypes (including integer columns and pandas string/categorical dtypes) round-trip.

Validation with the released v1.0.0 weights

  • Mixed-type dataset with known structure (age/job/income): corr(age, income) 0.817 real vs 0.833 synthetic; manager salary premium preserved with correct group ordering; per-column KS tests vs real p = 0.78–0.98.
  • iris: correlation matrix recovered with 12/12 sign agreement, mean abs error 0.059; per-species petal lengths within ~1%.
  • Temperature controls conditional (residual) spread monotonically: 6.1k / 5.1k / 4.5k at t = 1.0 / 0.5 / 0.1, vs 5.2k in the real data.
  • Novelty: median normalized nearest-real-neighbor distance 0.057; generated values are fresh draws, not copies (no formal privacy guarantee, as documented).

Testing

  • tabfm/src/generation_test.py: 25 backend-free tests (absltest) plus a JAX end-to-end test guarded like classifier_and_regressor_test.py.
  • tabfm/src/generation_pytorch_test.py: 5 end-to-end tests with a tiny random PyTorch model (no weight downloads); conftest-guarded like its sibling when torch is absent.
  • Registered in tabfm/src/BUILD (the torch-only test file is deliberately not registered, matching classifier_and_regressor_pytorch_test.py, since torch is not in the Bazel pip lock).
  • Full suite green under pytest -vv -n auto; lint profile matches the existing modules under the repo .pylintrc.
  • Runnable example: examples/synthetic_data_example.py; README section and CHANGELOG entry under [Unreleased] (no version bump).

Known v1 limits

No extrapolation beyond the observed min/max (TabPFN adds half-normal tails — a candidate follow-up), no causal/DAG column ordering, no missingness simulation.

@google-cla

google-cla Bot commented Jul 24, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support for synthetic data generation

1 participant