Teacher Feature Tokenizer Distillation

Minimal PyTorch implementation for testing feature distillation through a teacher-trained SimVQ tokenizer.

Methods

• ce: classification cross-entropy only.
• mlp_l2: CE plus a student-side channel MLP feature regression loss.
• tftd_code: CE plus code-distribution KL.
• tftd_rec: CE plus frozen-decoder reconstruction and commitment.
• tftd_full: CE plus code KL, reconstruction, and commitment.
• mlp_code: CE plus MLP-L2 feature regression and code-distribution KL.

The tokenizer uses a fixed random coefficient codebook C and a learnable linear basis W. Its hard quantization path uses:

zq_st = z + (zq_hard - z).detach()

The reconstruction loss updates the tokenizer encoder and decoder through the straight-through path. The codebook-side loss updates W; C is a buffer and never receives gradients.

Student feature maps are spatially aligned to the teacher feature map before feature regression, code KL, and reconstruction losses are computed. Matching spatial sizes are left unchanged; higher-resolution student features use adaptive average pooling, while lower-resolution student features use bilinear interpolation.

By default, loss_code is the raw KL term used by the V1-V8 experiments. For large temperature sweeps without re-tuning alpha_code, set distill.scale_code_loss_by_tau_s2: true to use alpha_code * tau_s^2 * loss_code in the total loss. Metrics record both the raw and scaled code loss.

Setup

Use a Python environment with a CUDA build of PyTorch for full CIFAR-100 training:

python -m pip install -e ".[test]"

The synthetic smoke test only requires PyTorch, PyYAML, NumPy, and pytest:

python -m tftd smoke
pytest

Training

Train the teacher:

python -m tftd train teacher \
  --config configs/cifar100_wrn.yaml \
  --data-dir /path/to/cifar100 \
  --download

Train the tokenizer with a frozen teacher:

python -m tftd train tokenizer \
  --config configs/cifar100_wrn.yaml \
  --teacher-ckpt outputs/teacher/best.pt

Train a student:

python -m tftd train student \
  --config configs/cifar100_wrn.yaml \
  --method tftd_full \
  --teacher-ckpt outputs/teacher/best.pt \
  --tokenizer-ckpt outputs/tokenizer/best.pt

For a two-batch, one-epoch dry run, add:

--epochs 1 --max-batches 2

Use --seed N to override the configured random seed for any training stage. Every stage writes last.pt, best.pt, and metrics.jsonl beneath its output directory. Use --resume <last.pt> to continue a run.

Validate a downloaded pretrained teacher without training it:

python -m tftd validate teacher \
  --config configs/cifar100_pairs/wrn_40_2__wrn_16_2.yaml \
  --teacher-ckpt outputs/pretrained_teachers/RepDistiller/save/models/wrn_40_2_vanilla/ckpt_epoch_240.pth

The stabilized tokenizer uses commitment weight 4.0, orthogonal basis weight 10.0, and metric temperature 0.02. V2 distillation calibrates temperatures around 0.01, 0.02, and 0.05 because the stabilized codebook has a much smaller distance scale than the original run.

Run the complete teacher, tokenizer, and student ablation sequence on GPU 0:

CUDA_VISIBLE_DEVICES=0 ./scripts/run_full_experiment.sh

The script automatically resumes each stage from its last.pt checkpoint.

Run the stabilized tokenizer and targeted V2 ablations:

CUDA_VISIBLE_DEVICES=0 \
DATA_DIR=/path/to/cifar100 \
BASELINE_DIR=/path/to/outputs/full_experiment \
./scripts/run_stabilized_v2.sh

Tokenizer runs save best_total.pt, best_rec.pt, last.pt, and a backward-compatible best.pt copy of best_rec.pt. best_total.pt is kept for tokenizer diagnostics; downstream distillation should use best_rec.pt.

Run the checkpoint, temperature, loss-weight, and three-seed V3 experiments with two concurrent jobs on GPU 0 and at most one job on GPU 1:

DATA_DIR=/path/to/cifar100 \
BASELINE_DIR=/path/to/outputs/full_experiment \
TOKENIZER_DIR=/path/to/outputs/stabilized_v2/tokenizer_stable \
./scripts/run_stabilized_v3.sh

The runner resumes each independent run from last.pt, stops after 9 hours 45 minutes by default, and writes final_summary.json and final_summary.md when all stages finish. Set V3_EPOCHS=1 and V3_MAX_BATCHES=2 for an end-to-end CUDA dry run.

Run the pretrained-teacher CIFAR-100 pair benchmark. This path downloads RepDistiller teacher checkpoints and does not train teacher models:

./scripts/fetch_repdistiller_teachers.sh
DATA_DIR=/path/to/cifar100 ./scripts/run_cifar100_pairs_pretrained.sh

The V8 runner trains a tokenizer, CE student, MLP-L2 student, and MLP-Code student for each configured pair. It runs four pairs concurrently by default: two on GPU 0 and two on GPU 1.

V8 showed that VGG13 -> VGG8 has tokenizer hard-code collapse with the default final 4x4 feature and 512-code tokenizer. Treat VGG as a separate feature layer/tokenizer repair case instead of first tuning student loss weights.

To run only selected pairs, pass a comma-separated V8_PAIRS list:

V8_PAIRS=wrn_40_2__wrn_16_2,resnet56__resnet20 \
DATA_DIR=/path/to/cifar100 \
./scripts/run_cifar100_pairs_pretrained.sh

Run the V9 diagnostic calibration protocol. This does not tune temperatures from student validation accuracy. It first computes unlabeled teacher/tokenizer assignment statistics, then trains one calibrated MLP-Code student per pair with fixed lambda_mlp=1.5 and alpha_code=0.5:

DATA_DIR=/path/to/cifar100 \
V8_DIR=/path/to/outputs/cifar100_pairs_pretrained_teachers \
TEACHER_ROOT=/path/to/RepDistiller/save/models \
./scripts/run_cifar100_pairs_calibrated_v9.sh

The V9 runner reuses V8 tokenizer checkpoints and V8 CE/MLP-L2/fixed MLP-Code results. It adds seven calibration.json files and seven mlp_code_calibrated student runs, then writes calibration_summary.md/json and final_summary.md/json. Use V9_EPOCHS=1 V9_MAX_BATCHES=2 for a dry run.

Run the V9.1 paired diagnostic after V9. It repeats only selected ResNet pairs with seeds 2/3/4, comparing default temperatures against the V9-calibrated temperatures under the same seed:

DATA_DIR=/path/to/cifar100 \
V8_DIR=/path/to/outputs/cifar100_pairs_pretrained_teachers \
V9_DIR=/path/to/outputs/cifar100_pairs_calibrated_v9 \
TEACHER_ROOT=/path/to/RepDistiller/save/models \
./scripts/run_cifar100_pairs_v9_1_diagnostic.sh

This produces 18 student runs and writes diagnostic_summary.md/json. Use V91_EPOCHS=1 V91_MAX_BATCHES=2 for a dry run.

Run the V10 layer/tokenizer repair diagnostic. It scans ResNet layer2 and VGG earlier blocks before running at most one seed-1 MLP-Code student per healthy candidate:

DATA_DIR=/path/to/cifar100 \
V8_DIR=/path/to/outputs/cifar100_pairs_pretrained_teachers \
TEACHER_ROOT=/path/to/RepDistiller/save/models \
./scripts/run_cifar100_pairs_v10_layer_repair.sh

This produces four tokenizer runs and at most three student runs, then writes layer_repair_summary.md/json. Use V10_EPOCHS=1 V10_MAX_BATCHES=2 for a dry run.

Run the V11 feature-standardization and linear-head diagnostic. It keeps tau/lambda/alpha fixed, trains standardized tokenizers, runs linear/std MLP-L2 and MLP-Code for all CIFAR-100 pairs, and adds two raw linear MLP-L2 probes:

DATA_DIR=/path/to/cifar100 \
V8_DIR=/path/to/outputs/cifar100_pairs_pretrained_teachers \
TEACHER_ROOT=/path/to/RepDistiller/save/models \
V11_GPU_SLOTS=0,0,1,1 \
./scripts/run_cifar100_pairs_v11_standardized_linear.sh

This writes standardized_linear_summary.md/json. Use V11_EPOCHS=1 V11_MAX_BATCHES=2 for a dry run.

Run the V12 VGG repair and WRN/ResNet degradation diagnostic. It reuses V11 standardized tokenizers, searches a small VGG structure grid, and runs two std/nonlinear controls for WRN and ResNet:

DATA_DIR=/path/to/cifar100 \
V8_DIR=/path/to/outputs/cifar100_pairs_pretrained_teachers \
V11_DIR=/path/to/outputs/cifar100_pairs_v11_standardized_linear \
TEACHER_ROOT=/path/to/RepDistiller/save/models \
V12_GPU_SLOTS=0,0,1,1 \
./scripts/run_cifar100_pairs_v12_vgg_repair.sh

This writes vgg_repair_summary.md/json. Use V12_EPOCHS=1 V12_MAX_BATCHES=2 for a dry run.

Run the V13a protocol cleanup. It keeps the common CIFAR-100 KD best-test protocol for comparability, but moves temperature selection into label-free assignment calibration artifacts:

python -m tftd calibrate assignment \
  --config configs/cifar100_pairs/resnet56__resnet20.yaml \
  --teacher-ckpt /path/to/resnet56_vanilla/ckpt_epoch_240.pth \
  --tokenizer-ckpt /path/to/tokenizer/best_rec.pt \
  --output-dir outputs/v13a_protocol/calibration/resnet56__resnet20/symmetric \
  --tau-t-grid 0.01,0.02,0.05,0.10,0.20 \
  --tau-s-grid 0.02,0.05,0.10,0.20,1.0 \
  --temperature-mode symmetric

Then run the V13a matrix:

DATA_DIR=/path/to/cifar100 \
V8_DIR=/path/to/outputs/cifar100_pairs_pretrained_teachers \
V11_DIR=/path/to/outputs/cifar100_pairs_v11_standardized_linear \
TEACHER_ROOT=/path/to/RepDistiller/save/models \
V13A_GPU_SLOTS=0,0,1,1 \
./scripts/run_v13a_protocol.sh

This writes final_summary.md/json. Use V13A_EPOCHS=1 V13A_MAX_BATCHES=2 for a dry run.