peal.adaptors.didae.DiDAE

class peal.adaptors.didae.DiDAE

Bases: Adaptor

DiDAE adaptor: discover shortcut directions in a diffusion autoencoder’s latent with a sparse dictionary, have a teacher judge them, and repair the student with CFKD.

The constructor loads the student, the train/validation/test dataloaders, the generator (whose sparse_dictionary is replaced by adaptor_config.sparse_dictionary when given) and the teacher, and writes config.yaml into base_dir. It refuses a dictionary whose fitted_on_encoder tag does not match the generator’s feature space, and refuses decode budgets outside service_mode. The workflow itself runs in run().

Parameters:

• adaptor_config (dict, str, Path or DiDAEConfig) – Loaded with load_yaml_config into a DiDAEConfig.

• **kwargs – Ignored.

Variables:

• student (torch.nn.Module) – The classifier being explained; replaced by the finetuned one after step 9.

• generator (InvertibleGenerator) – Provides encode, sweep_all_directions and compute_bipartite_gt_sae_matching.

• teacher (TeacherInterface) – Judges directions in step 8.

• generator_to_classifier (classifier_to_generator,) – Normalisation converters between the two datasets’ conventions.

• sweep_found_nothing (bool) – Set by run() when the sweep produced no direction.

Methods

run	Execute the full 9-step DiDAE workflow and return the (repaired) student.

__init__(adaptor_config=None, **kwargs)

Load student, data, generator, dictionary and teacher from the config.

Parameters:

adaptor_config (dict | str | Path | AdaptorConfig)

run()

Execute the full 9-step DiDAE workflow and return the (repaired) student.

Steps, with the artefacts written under base_dir:

• Step 1: distil the student into a linear probe on z_sem (distilled_probe/w.pt, bias.pt) and compute the empirical component bounds c_min_and_maxes.txt of the dictionary.

• Step 1.5: match ground-truth attributes to dictionary concepts on up to 5000 samples of the generator’s (or the adaptor’s) data via generator.compute_bipartite_gt_sae_matching.

• Steps 2-6: draw a CFKD-style class-balanced pool of min_train_samples + max_validation_samples images and call generator.sweep_all_directions; the ranked metadata goes to sweep_results.pt. With latent_only the run stops here.

• Step 7: render collages of the top-k directions into direction_collages/.

• Step 8: collect true/false verdicts from the teacher (a Model2ModelTeacher judges every direction with a verified flip) into direction_feedback.txt; false directions are ranked by false flips and pruned by cfkd_false_direction_min_share and cfkd_max_false_directions.

• Step 9: unless disabled, run CFKD on the false directions (cfkd_finetuning/, standalone config cfkd_on_false_directions.yaml) and evaluate the result on the test set.

Scalars are logged to logs/ (TensorBoard) and the final student is saved by _save_final_student. An empty sweep sets self.sweep_found_nothing and returns early.

Returns:

The student, finetuned when step 9 ran.

Return type:

torch.nn.Module