peal.generators.tabular_ddpm.TabularDDPM¶

class peal.generators.tabular_ddpm.TabularDDPM[source]¶

Bases: EditCapableGenerator

DDPM generator for tabular data with classifier-guided counterfactual edits.

On construction the beta/alpha/alpha_bar schedule is registered as buffers and train_model() is called with mode="train", i.e. the score network is (re)trained every time the generator is built.

Parameters:

config (TabularDDPMConfig) – Generator config (already a pydantic object, not a path).
dataset (Dataset, optional) – Tabular dataset with an attributes list; loaded from config.data via get_datasets when None. Columns named "Confounder" and "Target" are excluded from the diffused features (input_idx); "Target" columns form target_idx.
model_dir (str, optional) – Where the model and logs are written; defaults to config.base_path.
device (str) – Unused; tensors follow the classifier’s device in edit.

Variables:

alpha_bar (beta, alpha,) – [num_timesteps] schedule buffers.
model (BasicDiscreteTimeModel) – The trained score network (set by train_model()).
original_sample (counterfactuals,) – Last edit result and input, used by plot_counterfactuals().

Methods

`discard_counterfactuals`	Keep, per sample, the earliest candidate that reaches `target_confidence`.
`edit`	Produce counterfactuals for `x_in` (`EditCapableGenerator` protocol).
`forward_diffusion`	Sample `x_t = sqrt(alpha_bar_t) x_0 + sqrt(1 - alpha_bar_t) eps`.
`plot_counterfactuals`	Show a 2-D scatter of the dataset with arrows from each `original_sample` to its counterfactual (uses the first two `input_idx` columns).
`reverse_diffusion_ddpm`	One ancestral DDPM step `x_t -> x_{t-1}` with variance `beta_t`.
`sample_counterfactual_ddpm`	Classifier-guided reverse diffusion from a partially noised batch.
`sample_ddpm`
`sample_x`	Return `batch_size` fully denoised samples `[batch_size, input_dim]`.
`train_model`	Train the score network or load it from `model_dir`.

__init__(config, dataset=None, model_dir=None, device='cpu', **kwargs)[source]¶: Resolve the dataset and schedule, then train the score network.

forward_diffusion(clean_x, noise, timestep)[source]¶

Sample x_t = sqrt(alpha_bar_t) x_0 + sqrt(1 - alpha_bar_t) eps.

Parameters:

clean_x (torch.Tensor) – [B, D] clean features.
noise (torch.Tensor) – [B, D] standard normal noise eps.
timestep (int or torch.Tensor) – A single timestep applied to the whole batch, or a [B] tensor of per-sample timesteps.

Returns:

[B, D] noised features.

Return type:

torch.Tensor

reverse_diffusion_ddpm(noisy_x, model, timestep)[source]¶

One ancestral DDPM step x_t -> x_{t-1} with variance beta_t.

Parameters:

noisy_x (torch.Tensor) – [B, D] current sample x_t.
model (nn.Module) – Score network called as model(x=noisy_x, t=timestep).
timestep (int) – Scalar timestep shared by the batch; no noise is added at 0.

Returns:

[B, D] sample x_{t-1}.

Return type:

torch.Tensor

train_model(model_name='diffusion.pt', mode='train')[source]¶

Train the score network or load it from model_dir.

Parameters:

model_name (str) – State-dict filename inside model_dir.
mode (str) – "train" always trains for config.num_epochs epochs with Adam on the epsilon-MSE loss and saves the state dict plus config.yaml; any other value loads model_name if it exists and otherwise only logs a hint.

Notes

Training logs Loss/train and, every 10 epochs, a scatter plot of generated versus real points to <model_dir>/logs (TensorBoard). The trained/loaded network is stored as self.model. The nested run_epoch helper is defined but never called.

sample_ddpm(model, n_samples=256, label=None)¶

Iteratively denoise pure noise, keeping the sample after every step.

Parameters:

model (nn.Module) – Score network (switched to eval mode).
n_samples (int) – Number of samples drawn from N(0, I) in input_dim dims.
label (optional) – Unused.

Returns:

num_timesteps + 1 tensors of shape [n_samples, input_dim], from the initial noise to the final sample ([-1]).

Return type:

list of torch.Tensor

sample_x(batch_size=1)[source]¶: Return batch_size fully denoised samples [batch_size, input_dim].

sample_counterfactual_ddpm(clean_batch, model, classifier, num_noise_steps, target_classes, classifier_grad_weight)[source]¶

Classifier-guided reverse diffusion from a partially noised batch.

The clean batch is noised to num_noise_steps and denoised again; at every step the posterior mean is shifted by the (scaled) gradient of the cross-entropy towards target_classes, evaluated on the fully denoised estimate of the previous step. After each guided step the intermediate sample is denoised to t = 0 without guidance to obtain a candidate counterfactual.

Parameters:

clean_batch (torch.Tensor) – [B, D] inputs to edit.
model (nn.Module) – Score network.
classifier (nn.Module) – Differentiable predictor on the D features.
num_noise_steps (int) – Forward-diffusion depth t to start from.
target_classes (torch.Tensor) – [B] target labels.
classifier_grad_weight (float) – Scale of the classifier gradient.

Returns:

(counterfactuals, guided_grads, unguided_grads, total_series): counterfactuals [B, num_noise_steps, D] (the clean input first, then one candidate per guided step), the guided and unconditional update terms with the same shape, and a list of the per-step denoising trajectories. The first three are also stored as attributes counterfactuals_series, guided_grads and unguided_grads.

Return type:

tuple

discard_counterfactuals(counterfactuals, classifier, target_classes, target_confidence, minimal_counterfactuals, tolerance=0.1)[source]¶

Keep, per sample, the earliest candidate that reaches target_confidence.

Parameters:

counterfactuals (torch.Tensor) – [B, S, D] candidates per sample from sample_counterfactual_ddpm().
classifier (nn.Module) – Predictor whose softmax confidence is checked.
target_classes (torch.Tensor) – [B] target labels.
target_confidence (float) – Required target-class confidence.
minimal_counterfactuals (torch.Tensor) – [B, D] current best counterfactuals, updated in place.
tolerance (float) – Unused.

Returns:

minimal_counterfactuals: unchanged rows that already satisfy the confidence, otherwise the first satisfying candidate, or the last candidate when none satisfies it.

Return type:

torch.Tensor

edit(x_in, target_confidence_goal, source_classes, target_classes, classifier=None, **kwargs)[source]¶

Produce counterfactuals for x_in (EditCapableGenerator protocol).

Loops config.attack_iterations times over every noise_steps_for_counterfactuals x grad_scales combination, running sample_counterfactual_ddpm() from the current best counterfactuals and keeping the first candidate that reaches target_confidence_goal via discard_counterfactuals().

Parameters:

x_in (torch.Tensor) – [B, D] inputs (moved to the classifier’s device).
target_confidence_goal (float) – Required confidence for the target class.
source_classes (torch.Tensor) – [B] source labels (unused).
target_classes (torch.Tensor) – [B] target labels.
classifier (nn.Module, optional) – Predictor; taken from kwargs["predictor"] when None.
**kwargs – Other explainer arguments (predictor, explainer_config, …); only predictor is read.

Returns:

(counterfactual_list, diff_latent, y_target_end_confidence, x_list, history_list, cluster_list): per-sample counterfactuals, x_in - counterfactuals as a [B, D] tensor, per-sample target-class confidences, per-sample inputs, and two lists of empty/None placeholders. The counterfactuals and inputs are also stored as self.counterfactuals / self.original_sample.

Return type:

tuple

plot_counterfactuals()[source]¶: Show a 2-D scatter of the dataset with arrows from each original_sample to its counterfactual (uses the first two input_idx columns).