peal.adaptors.counterfactual_knowledge_distillation.CFKD

class peal.adaptors.counterfactual_knowledge_distillation.CFKD

Bases: Adaptor

This class implements the counterfactual knowledge distillation approach.

The constructor wires up everything one run needs (student, dataloaders, generator, explainer, teacher, sparse dictionary); run() executes the finetune iterations and returns the corrected student. See the module docstring for the algorithm and the per-iteration artifacts.

Parameters:

• student (torch.nn.Module, optional) – The classifier to repair. If None it is built from adaptor_config.student with get_predictor.

• datasource (list or tuple, optional) – Datasets/dataloaders handed to create_dataloaders_from_datasource; None means the data is built from adaptor_config.data.

• generator (InvertibleGenerator, Path or str, optional) – Generator instance or config path; falls back to adaptor_config.generator.

• base_dir (str or Path, optional) – Run directory; falls back to adaptor_config.base_dir.

• teacher (str or TeacherInterface, optional) – Teacher spec (e.g. "human@8000", "SegmentationMask", a dict {type: llm}) or instance; falls back to adaptor_config.teacher.

• adaptor_config (dict, str, Path or CFKDConfig) – The run configuration (yaml path, dict or object).

• overwrite (bool, optional) – Whether cached artifacts in base_dir are discarded; falls back to adaptor_config.overwrite. The stored config always gets overwrite=False so a resumed run reuses its cache.

Variables:

• original_student (torch.nn.Module) – The uncorrected student, kept for the progress visualisation.

• student (torch.nn.Module) – The student being finetuned (a deep copy of the input).

• test_dataloader (train_dataloader, val_dataloader,) – The original data; val_batch_size is forced to the CFKD batch_size because validation batches go straight into the explainer.

• dataloader_mixer (DataloaderMixer) – Original data plus every counterfactual dataset added so far.

• datastack (DataStack) – Per-class sample buffer the batch builder pops factuals from.

• tracked_keys (list of str) – Keys of the explainer output that are collected and cached (extended with cluster_list, z_difference_list/collage_path_list at tracking_level >= 4, hint_list and idx_list as needed).

• validation_data_config (data_config,) – Copies of the config used to load the serialised counterfactual datasets for training and validation respectively.

Methods

add_dataset_to_dataloader_mixer	Load a serialised counterfactual dataset and mix it with the old data.
calculate_feedback_stats	Explainer metrics from the teacher verdicts of one validation pass.
calculate_unbiasedness	Unbiasedness score of the explainer from the distilled feedback accuracy.
create_dataset	Serialise judged counterfactuals as a dataset for finetuning.
ensure_group_accuracies	The latest per-group test accuracies, computing them once if needed.
finetune_student	Finetune the student on the mixed data of one iteration.
generate_x_counterfactual_list	Generate training counterfactuals until min_train_samples are found.
get_batch	Assemble one batch of factuals with (source, target) class pairs.
initialize_run	Prepare base_dir and compute or reload the iteration-0 state.
retrieve_counterfactual_list	Training counterfactuals of an iteration, generated or loaded from cache.
retrieve_feedback	Teacher verdicts for a set of counterfactuals, cached per iteration.
retrieve_validation_prestats	Validation counterfactuals and pre-feedback statistics of an iteration.
retrieve_validation_stats	Complete the validation statistics with teacher feedback.
run	Run the counterfactual knowledge distillation
visualize_progress	Render the before/after comparison figure for a binary task.

__init__(student=None, datasource=None, generator=None, base_dir=None, teacher=None, adaptor_config='<PEAL_BASE>/configs/adaptors/symbolic_cfkd.yaml', overwrite=None)

Build all components of a CFKD run; see the class docstring for the parameters.

Besides constructing the objects this also mutates adaptor_config: test_data defaults to data, in_memory/tracking_level/ transition_restrictions/clustering_strategy are pushed into the explainer config, training.val_batch_size is set to batch_size and training.steps_per_epoch to one pass over the training set. A collage-based teacher (human, web, llm) requires tracking_level >= 4 and fails an assertion otherwise.

Parameters:

• student (torch.nn.Module)

• datasource (list | tuple)

• generator (InvertibleGenerator | Path | str)

• base_dir (str | Path)

• teacher (str | TeacherInterface)

• adaptor_config (dict | str | Path | AdaptorConfig)

• overwrite (bool)

initialize_run()

Prepare base_dir and compute or reload the iteration-0 state.

On a fresh run (overwrite or no logs dir yet) the existing base_dir is moved to <base_dir>_old_<timestamp>, the initial validation/test (group) accuracies, sample batches and a generator sample (with FID) are written to TensorBoard under base_dir/logs, config.yaml and platform.txt are written and the validation counterfactuals of iteration 0 are generated. On a resumed run the cached validation values are loaded instead, every counterfactual dataset up to current_iteration is re-added to the dataloader mixer and the last model.cpl becomes the student. Optionally a decision-boundary plot (2-D latent datasets, tracking_level >= 4) and the progress figure (binary tasks, tracking_level >= 6) are rendered.

Returns:

(validation_stats, validation_tracked_values, writer): the statistics dict of the last completed validation pass, the tracked validation counterfactuals and the SummaryWriter.

Return type:

tuple

get_batch(error_matrix=None, cm_idx_in=0)

Assemble one batch of factuals with (source, target) class pairs.

Pairs are walked through the flattened output_size x output_size confusion-matrix index, skipping the diagonal, or sampled from error_matrix when use_confusion_matrix is set. Each factual is popped from the datastack of its source class; in "1sided" mode it is kept only when the student predicts it correctly. Hints and dataset indices are split off the label when enabled.

Parameters:

• error_matrix (torch.Tensor, optional) – Flattened error distribution over class pairs from the validation stats; only used with use_confusion_matrix.

• cm_idx_in (int) – Starting offset (in rows) of the confusion-matrix walk, so that alternate batches start from different pairs.

Returns:

x_list (stacked tensor [B, ...]), y_list, y_target_list (tensor [B]), y_source_list, y_target_start_confidence_list (softmax of the tempered logits at the target class), hint_list (zeros unless a SegmentationMaskTeacher is used) and idx_list (zeros unless the explainer is PerfectFalseCounterfactualConfig).

Return type:

dict

generate_x_counterfactual_list(error_matrix, confidence_score_stats, finetune_iteration, tracked_keys)

Generate training counterfactuals until min_train_samples are found.

Repeatedly builds batches with get_batch(), runs explainer.explain_batch on them and accumulates the requested keys. Collages go to base_dir/<iteration>/collages (an existing directory is moved aside with an _old_<timestamp> suffix); the explainer’s state directory is base_dir/<iteration - 1>.

Parameters:

• error_matrix (torch.Tensor) – Passed to get_batch().

• confidence_score_stats (object) – Currently unused.

• finetune_iteration (int) – Iteration whose directory receives the collages.

• tracked_keys (list of str) – Keys of the explainer output to collect.

Returns:

{key: list} for every key in tracked_keys, with at least adaptor_config.min_train_samples entries each (batches are not truncated, so usually a few more).

Return type:

dict

retrieve_counterfactual_list(validation_stats, finetune_iteration)

Training counterfactuals of an iteration, generated or loaded from cache.

Generates them with generate_x_counterfactual_list() and saves base_dir/<iteration>/tracked_values.npz (tracking_level >= 3), or loads that file (expanding old per-factual layouts with flatten_explanations). Then collage_path_list is rebuilt from the .png files in the collages directory and, if the explainer clusters, cluster_explanations adds the clusters<k> entries.

Parameters:

• validation_stats (dict) – Must contain error_matrix and confidence_score_stats.

• finetune_iteration (int) – Iteration index (1-based).

Returns:

The tracked values, one list per key.

Return type:

dict

retrieve_feedback(tracked_values, finetune_iteration, mode)

Teacher verdicts for a set of counterfactuals, cached per iteration.

As a side effect the distilled predictor used by the explainer metrics is prepared: with calculate_explainer_stats the student is distilled (distill_predictor with a leaky-softplus activation) into base_dir/<current_iteration>/distilled_predictor or reloaded from there; otherwise the student itself is used. The verdicts are written to base_dir/<iteration>/<mode>_feedback.txt and read back from it on later runs.

Parameters:

• tracked_values (dict) – Tracked counterfactual values, unpacked as keyword arguments of teacher.get_feedback.

• finetune_iteration (int) – Iteration whose directory holds the feedback file.

• mode (str) – "train" or "validation"; also names the teacher’s working directory <mode>_teacher.

Returns:

One verdict per counterfactual ("true", "false", "ood" or a teacher-specific marker string).

Return type:

list of str

ensure_group_accuracies()

The latest per-group test accuracies, computing them once if needed.

Returns:

adaptor_config.group_accuracies[-1] if present; otherwise the groups’ accuracies freshly measured on the test loader (also appended to the config and saved to config.yaml), or None when calculate_group_accuracies is off.

Return type:

list or None

calculate_unbiasedness(feedback_accuracy_distilled)

Unbiasedness score of the explainer from the distilled feedback accuracy.

r_dominant = 1 - mean(two lowest group accuracies) estimates how often the student uses the shortcut; r_actual = 1 - feedback_accuracy_distilled how often the explainer reveals it. The score is exp(-|log(r_actual / r_dominant)|), i.e. 1 when both rates agree.

Parameters:

feedback_accuracy_distilled (float) – Fraction of “true” verdicts among counterfactuals that flip both the student and its distilled copy.

Returns:

The score, 0.0 when a rate is zero, None when group accuracies are unavailable or fewer than two groups exist.

Return type:

float or None

calculate_feedback_stats(tracked_values, feedback, finetune_iteration)

Explainer metrics from the teacher verdicts of one validation pass.

Always computes flip_rate (target confidence >= 0.51, relative to max_validation_samples at iteration 0), ood_rate, feedback_accuracy (fraction of “true” among true+false verdicts on correctly classified factuals) and counterfactuals_per_second. With several attempts per factual only the highest-confidence attempt counts; the attempt layout (interleaved vs sequential) is inferred from y_list.

With calculate_explainer_stats at iteration 0 the distilled predictor is additionally run on every counterfactual and cluster member, cluster collages are written to base_dir/<iteration>/latent_cluster_collages<k>/ and the following are added: flip_rate_distilled (NAFR: flips shared by student and distilled copy), non_adversarial_rate, feedback_accuracy_distilled, counterfactual_quality (dataset track_generator_performance on the non-adversarial flips, with quality collages in latent_quality_collages/), unbiasedness (when group accuracies exist) and the latent sparsity/diversity stats of explainer.calculate_latent_difference_stats.

Parameters:

• tracked_values (dict) – Tracked validation counterfactuals; extended in place with y_target_end_confidence_distilled_list and cluster_confidence_distilled<k>.

• feedback (list of str) – One verdict per counterfactual.

• finetune_iteration (int) – Iteration the values belong to.

Returns:

Metric name to float.

Return type:

dict

create_dataset(x_counterfactual_list, feedback, y_source_list, y_target_list, finetune_iteration, hint_list=None, mode='', **kwargs)

Serialise judged counterfactuals as a dataset for finetuning.

“false” counterfactuals are stored with their source class as label (the change did not alter the class, so the student must not flip); “true” ones with the target class when use_true_counterfactuals is set. Sample names encode verdict, source, target and index. The dataset is written with train_dataloader.dataset.serialize_dataset to base_dir/<iteration>/<mode>_dataset and skipped if that directory already exists.

Parameters:

• x_counterfactual_list (list of torch.Tensor) – Counterfactual inputs.

• feedback (list of str) – One verdict per counterfactual; must match in length.

• y_source_list (list) – Source and target class per counterfactual.

• y_target_list (list) – Source and target class per counterfactual.

• finetune_iteration (int) – Iteration directory to write into.

• hint_list (list of torch.Tensor, optional) – Segmentation hints stored alongside the samples.

• mode (str) – "train" or "validation".

• **kwargs – Other tracked values, ignored.

Returns:

Path of the dataset directory.

Return type:

str

Raises:

Exception – If x_counterfactual_list and feedback differ in length.

add_dataset_to_dataloader_mixer(dataloader_old, dataset_path, mixing_ratio, writer, finetune_iteration)

Load a serialised counterfactual dataset and mix it with the old data.

Parameters:

• dataloader_old (DataloaderMixer) – The data used so far.

• dataset_path (str) – Directory written by create_dataset().

• mixing_ratio (float) – Weight of the new counterfactual data; the old mixer is appended with weight 1 - mixing_ratio.

• writer (SummaryWriter) – Sample images of the new data are logged as train_<iteration>.

• finetune_iteration (int) – Used for the log tag.

Returns:

New mixer with return_src_internal=True and hints enabled when the run uses them.

Return type:

DataloaderMixer

finetune_student(finetune_iteration, dataset_path, writer)

Finetune the student on the mixed data of one iteration.

Loads the validation counterfactual dataset of the iteration (written by retrieve_validation_stats()) and appends it to the joint validation loaders, adds dataset_path to the dataloader mixer, rebuilds the datastack and trains with ModelTrainer into base_dir/<iteration>/finetuned_model (an existing directory is moved aside). continuous_learning selects "finetune", "retrain" (fresh resnet18) or "deep_feature_reweighting" (last layer only). Hints and dataset indices are disabled during training. Afterwards model.cpl is loaded as the new student and handed to the explainer. If the validation split is empty or training stops early, error_iteration_<n>.txt is written and the method returns without changing the student.

Parameters:

• finetune_iteration (int) – Current iteration.

• dataset_path (str) – Training counterfactual dataset from create_dataset().

• writer (SummaryWriter) – Receives sample images of the new loaders.

visualize_progress(paths)

Render the before/after comparison figure for a binary task.

Uses create_comparison to show test samples with counterfactuals of the uncorrected and of the CFKD-corrected student side by side, with criteria for class, confounder (if the dataset exposes a Confounder attribute) and both predictions. Ranged explainer parameters (two-element lists) are replaced by their midpoint. Two images are saved per path: <path> and <path>_success.png (the latter filtered by a checkbox pattern of expected predictions).

Parameters:

paths (list of str) – .png output paths.

Returns:

The unfiltered comparison image.

Return type:

PIL.Image.Image

retrieve_validation_prestats(finetune_iteration)

Validation counterfactuals and pre-feedback statistics of an iteration.

Runs calculate_validation_statistics validation_runs times (interpolating ranged explainer parameters between runs) on the joint validation loaders and averages the statistics (error matrix, confidence score stats, …). Results are cached as validation_tracked_values.npz and validation_prestats.npz in base_dir/<iteration>/ at tracking_level >= 3 and loaded from there otherwise, with collage_path_list rebuilt from validation_collages0. If the explainer clusters, the explanations are clustered and stored as validation_tracked_cluster_values.npz; 2-D latent datasets additionally get val_counterfactuals_global.png.

Parameters:

finetune_iteration (int) – Iteration directory to use.

Returns:

(validation_tracked_values, validation_stats).

Return type:

tuple

retrieve_validation_stats(finetune_iteration, validation_tracked_values, validation_prestats)

Complete the validation statistics with teacher feedback.

If base_dir/<iteration>/validation_stats.npz exists and overwrite is off it is loaded (back-filling unbiasedness for old files when possible). Otherwise the teacher judges the validation counterfactuals, calculate_feedback_stats() adds the explainer metrics to validation_prestats, the judged counterfactuals are serialised as <iteration + 1>/validation_dataset for the next finetuning and the merged dict is saved.

Parameters:

• finetune_iteration (int) – Iteration the validation values belong to.

• validation_tracked_values (dict) – From retrieve_validation_prestats().

• validation_prestats (dict) – From retrieve_validation_prestats(); extended in place.

Returns:

Statistics including the feedback metrics.

Return type:

dict

run()

Run the counterfactual knowledge distillation

Starting at current_iteration + 1 and up to finetune_iterations, every iteration retrieves training counterfactuals and their feedback, finalises the previous iteration’s validation statistics (logged as validation_* scalars), builds the counterfactual dataset, finetunes the student, logs val_accuracy/test_accuracy (plus group accuracies and gain, the relative improvement of the mean group accuracy, when enabled), computes fresh validation prestats for the next iteration, saves model.cpl and advances current_iteration in config.yaml. The post-loop validation pass is skipped unless PEAL_FULL_VALIDATION_TAIL=1.

Returns:

The finetuned student.

Return type:

torch.nn.Module