Data Reconstruction: Identifiability and Optimization with Sample Splitting

Code for the ICML 2026 paper Data Reconstruction: Identifiability and Optimization with Sample Splitting.

This repository implements sample splitting algorithm on top of three existing reconstruction methods to demonstrate its effectiveness: Haim et al. for binary KKT reconstruction, Buzaglo et al. for multiclass KKT reconstruction, and Loo et al. for binary NTK-based reconstruction.

Setup

Create a local settings file from the default template and edit the dataset, model, and result paths:

cp settings.default.py settings.py

Create and activate the conda environment:

conda env create -f environment.yaml
conda activate rec

The code expects settings.py to define:

datasets_dir = "/path/to/datasets"
results_base_dir = "/path/to/runs"
models_dir = "/path/to/models"

Running the Code

All training and reconstruction experiments are launched through Main.py. Example command templates for training, sample-splitting reconstruction, and no-splitting baselines are stored in command_line_args/. The reconstruction templates use the default values in GetParams.py for --extraction_init_scale, --extraction_lr, --extraction_min_lambda, and --extraction_model_relu_alpha; these parameters can be tuned with a sweep when reproducing or extending the experiments.

The following three commands illustrate one consistent CIFAR-10 vehicles-vs-animals setting. First train a 3-layer MLP with 250 samples per class:

python Main.py \
  --run_mode=train \
  --problem=cifar10_vehicles_animals \
  --proj_name=cifar10_vehicles_animals \
  --data_per_class_train=250 \
  --model_hidden_list=[1000,1000] \
  --model_init_list=[0.0001,0.0001] \
  --train_epochs=1000000 \
  --train_lr=0.01 \
  --train_evaluate_rate=1000

Then reconstruct from the trained model with Haim-style KKT reconstruction and sample splitting.

python Main.py \
  --run_mode=reconstruct \
  --extraction_method=Haim \
  --problem=cifar10_vehicles_animals \
  --data_per_class_train=250 \
  --extraction_data_amount_per_class=500 \
  --extraction_epochs=50000 \
  --extraction_evaluate_rate=1000 \
  --model_hidden_list=[1000,1000] \
  --model_init_list=[0.001,0.001] \
  --pretrained_model_path=weights-cifar10_vehicles_animals_d250_cifar10_vehicles_animals.pth \
  --wandb_active=False

The Loo-style NTK reconstruction uses the same trained model and additionally needs the corresponding initialization checkpoint saved before training.

python Main.py \
  --run_mode=reconstruct \
  --extraction_method=Loo \
  --problem=cifar10_vehicles_animals \
  --data_per_class_train=250 \
  --extraction_data_amount_per_class=500 \
  --extraction_epochs=50000 \
  --extraction_evaluate_rate=1000 \
  --model_hidden_list=[1000,1000] \
  --model_init_list=[0.001,0.001] \
  --pretrained_model_path=weights-cifar10_vehicles_animals_d250_cifar10_vehicles_animals.pth \
  --initial_model_path=weights-cifar10_vehicles_animals_d250_cifar10_vehicles_animals_initial.pth \
  --wandb_active=False

Visualization

Use the triplet visualization scripts to compare reconstructions without and with sample splitting:

python reconstruction_mnist_triplet.py
python reconstruction_cifar10_triplet.py

Acknowledgements

This codebase builds on nivha/dataset_reconstruction.