Energy-based data generation

Installation

pip install tabpfn-extensions

`TabEBM`

View source

TabEBM: Tabular Energy-Based Model for synthetic data generation. This class implements an energy-based model that uses TabPFN as the underlying classifier to define energy functions. It generates synthetic tabular data using Stochastic Gradient Langevin Dynamics (SGLD) sampling. The core idea is to treat each class as having its own energy landscape, where real data points have low energy and synthetic points are generated by following the energy gradient through SGLD sampling. Initialize TabEBM with optimized configuration.

TabEBM(
    max_data_size: int = 10000,
)

Parameters

Parameter	Type	Default	Description
`max_data_size`	`int`	`10000`	Maximum number of data points to use for training. Larger datasets will be subsampled to this size.

`TabEBM.add_surrogate_negative_samples`

View source

Create surrogate negative samples for TabEBM’s binary classification approach. This method creates artificial “negative” samples at specified distances from the origin to enable energy-based modeling through binary classification. The surrogate negatives help define the energy landscape by providing clear decision boundaries. For 2D data, negatives are placed at the four corners of a square centered at origin. For higher dimensions, random combinations of ±distance_negative_class are used.

TabEBM.add_surrogate_negative_samples(
    X: np.ndarray | torch.Tensor,
    distance_negative_class: float = 5,
) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]

Parameters

Parameter	Type	Default	Description
`X`	`np.ndarray \| torch.Tensor`	Required	Real data samples (expected to be approximately standardized) Shape: (num_samples, num_features)
`distance_negative_class`	`float`	`5`	Distance of surrogate negatives from origin Larger values create more distinct separation

Returns

Type	Description
`tuple[np.ndarray \| torch.Tensor, np.ndarray \| torch.Tensor]`	Tuple of (X_ebm, y_ebm) where: - X_ebm: Combined real and surrogate samples - y_ebm: Binary labels (0 for real data, 1 for surrogates)

`TabEBM.compute_energy`

View source

Compute TabEBM class-specific energy function. The energy function is defined as: E_c(x) = -log(exp(f^c(x)[0]) + exp(f^c(x)[1])) Where f^c(x) are the logits from the class-specific binary classifier. Lower energy corresponds to higher probability of belonging to the target class.

TabEBM.compute_energy(
    logits: torch.Tensor | np.ndarray,
    return_unnormalized_prob: bool = False,
) -> torch.Tensor | np.ndarray

Parameters

Parameter	Type	Default	Description
`logits`	`torch.Tensor \| np.ndarray`	Required	Model’s unnormalized logits for each class Shape: (num_samples, num_classes)
`return_unnormalized_prob`	`bool`	`False`	If `True`, return exp(-energy) instead of energy

Returns

Type	Description
`torch.Tensor \| np.ndarray`	Energy values (or unnormalized probabilities) for each sample

Raises ValueError If logits are not unnormalized or have wrong type

`TabEBM.generate`

View source

Generate synthetic samples using Stochastic Gradient Langevin Dynamics (SGLD). This method creates synthetic data by treating the TabPFN classifier as an energy function and using SGLD to sample from the learned energy landscape. For each class, it creates a binary classification problem (target class vs surrogate negatives) and samples new points by following energy gradients.

TabEBM.generate(
    X: np.ndarray | torch.Tensor | pd.DataFrame,
    y: np.ndarray | torch.Tensor | pd.Series,
    num_samples: int,
    starting_point_noise_std: float = 0.01,
    sgld_step_size: float = 0.1,
    sgld_noise_std: float = 0.01,
    sgld_steps: int = 200,
    distance_negative_class: float = 5,
    seed: int = 42,
    verbose: bool = False,
) -> dict[str, np.ndarray]

Parameters

Parameter	Type	Default	Description
`X`	`np.ndarray \| torch.Tensor \| pd.DataFrame`	Required	Input features of shape (n_samples, n_features)
`y`	`np.ndarray \| torch.Tensor \| pd.Series`	Required	Target labels of shape (n_samples,)
`num_samples`	`int`	Required	Number of synthetic samples to generate per class
`starting_point_noise_std`	`float`	`0.01`	Standard deviation of noise added to real data points when initializing SGLD chains
`sgld_step_size`	`float`	`0.1`	Step size for gradient updates in SGLD sampling
`sgld_noise_std`	`float`	`0.01`	Standard deviation of noise added at each SGLD step
`sgld_steps`	`int`	`200`	Number of SGLD steps to perform
`distance_negative_class`	`float`	`5`	Distance for placing surrogate negative samples from the origin (used to create binary classification)
`seed`	`int`	`42`	Random seed for reproducibility
`verbose`	`bool`	`False`	Whether to print verbose information during sampling

Returns

Type	Description
`dict[str, np.ndarray]`	Dictionary mapping class names to generated samples: { ‘class_0’: `np.ndarray` of shape (`num_samples`, n_features), ‘class_1’: `np.ndarray` of shape (`num_samples`, n_features), … }

`TabEBM.train_test_split_allow_full_train`

View source

Enhanced train-test split that supports full training mode. This method extends sklearn’s train_test_split to handle the case where test_size=0, which means we want to use all data for training (no validation). This is useful for TabEBM’s energy-based training approach.

TabEBM.train_test_split_allow_full_train(
    X: np.ndarray | torch.Tensor,
    y: np.ndarray | torch.Tensor,
    test_size: float | None = None,
    train_size: float | None = None,
    random_state: int | None = None,
    shuffle: bool = True,
    stratify: np.ndarray | torch.Tensor | None = None,
) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor, np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]

Parameters

Parameter	Type	Default	Description
`X`	`np.ndarray \| torch.Tensor`	Required	—
`y`	`np.ndarray \| torch.Tensor`	Required	—
`test_size`	`float \| None`	`None`	Fraction of data for testing (if 0, enables full train mode)
`train_size`	`float \| None`	`None`	Fraction of data for training
`random_state`	`int \| None`	`None`	Random seed for reproducibility
`shuffle`	`bool`	`True`	Whether to shuffle data before splitting
`stratify`	`np.ndarray \| torch.Tensor \| None`	`None`	Array-like for stratified splitting

Returns

Type	Description
`tuple[np.ndarray \| torch.Tensor, np.ndarray \| torch.Tensor, np.ndarray \| torch.Tensor, np.ndarray \| torch.Tensor]`	Tuple of (X_train, X_val, y_train, y_val) In full train mode, X_train=`X` and y_train=y

`seed_everything`

View source

Set random seeds for reproducibility across all libraries.

seed_everything(
    seed: int,
) -> None

Parameters

Parameter	Type	Default	Description
`seed`	`int`	Required	Random seed value

Returns

Type	Description
`None`	—

`to_numpy`

View source

Convert input data to numpy array format.

to_numpy(
    X: np.ndarray | torch.Tensor | pd.DataFrame | None,
) -> np.ndarray | None

Parameters

Parameter	Type	Default	Description
`X`	`np.ndarray \| torch.Tensor \| pd.DataFrame \| None`	Required	Input data in various formats (numpy array, torch tensor, pandas DataFrame, or `None`)

Returns

Type	Description
`np.ndarray \| None`	numpy array representation of the input data, or `None` if input is `None`

Raises ValueError If input type is not supported

Reference

tabpfn

tabpfn-client

tabpfn-extensions

REST API

Energy-based data generation

Installation

`TabEBM`

`TabEBM.add_surrogate_negative_samples`

`TabEBM.compute_energy`

`TabEBM.generate`

`TabEBM.train_test_split_allow_full_train`

`seed_everything`

`to_numpy`

​Installation

​TabEBM

​TabEBM.add_surrogate_negative_samples

​TabEBM.compute_energy

​TabEBM.generate

​TabEBM.train_test_split_allow_full_train

​seed_everything

​to_numpy

Installation

`TabEBM`

`TabEBM.add_surrogate_negative_samples`

`TabEBM.compute_energy`

`TabEBM.generate`

`TabEBM.train_test_split_allow_full_train`

`seed_everything`

`to_numpy`