From Per-Image Low-Rank to Encoding Mismatch

This repository releases the Lift and WideLast checkpoints from the ICML 2026 paper:

From Per-Image Low-Rank to Encoding Mismatch: Rethinking Feature Distillation in Vision Transformers
Huiyuan Tian, Bonan Xu, and Shijian Li

The paper studies why feature distillation between heterogeneous Vision Transformers can remain difficult even when individual feature maps are strongly low-rank. Its analysis distinguishes:

• Per-image low-rank structure, measured by sample-wise singular value decomposition (SVD).
• Dataset-level representational structure, measured by PCA over features from many images.
• Spectral Energy Patterns (SEP), which characterize how feature energy is distributed across channels.

The central finding is an encoding mismatch: compact students and wider teachers can encode dataset-level information using substantially different channel-space geometries. Two simple architecture modifications reduce this mismatch:

• Lift retains a linear endpoint projector that lifts the student's final representation from 192 to 384 channels.
• WideLast widens only the last transformer block from 192 to 384 channels.

Links

• Paper on arXiv
• Paper on Hugging Face
• Paper on OpenReview
• Official code, analysis scripts, raw numeric data, and training recipes
• Analysis dataset and Dataset Viewer
• ICML 2026 poster

Released Checkpoints

The reported baseline distilled DeiT-Tiny accuracy is 74.86% under the paper's experimental protocol.

Download Individual Checkpoints

Lift and WideLast use the custom architecture definitions and checkpoint-loading code in the GitHub repository. They cannot be loaded through timm's generic Hugging Face Hub entry point. Download each checkpoint directly as follows, then use the evaluation commands below with the matching architecture implementation:

from huggingface_hub import hf_hub_download

lift_checkpoint = hf_hub_download(
    repo_id="Huiyuancs/Encoding_Mismatch",
    filename="models/Lift/pytorch_model.pth",
)

widelast_checkpoint = hf_hub_download(
    repo_id="Huiyuancs/Encoding_Mismatch",
    filename="models/WideLast/pytorch_model.pth",
)

print("Lift:", lift_checkpoint)
print("WideLast:", widelast_checkpoint)

Model Details

Shared setting

• Task: ImageNet-1K image classification
• Input resolution: 224 × 224
• Student family: DeiT-Tiny-derived Vision Transformer
• Teacher: CaiT-S24
• Framework: PyTorch and timm
• Training duration: 300 epochs
• Primary metric: ImageNet-1K validation top-1 accuracy

These models use custom architectures and checkpoint-loading code from the accompanying GitHub repository. They are not packaged as transformers.AutoModel checkpoints and are not expected to run directly through the Hugging Face hosted inference widget.

Lift

Lift preserves the standard 192-dimensional DeiT-Tiny transformer blocks and adds a learned linear projection at the student endpoint:

$$ \mathbb{R}^{192} \rightarrow \mathbb{R}^{384}. $$

Unlike a training-only projector, this layer is retained at inference. The released Lift checkpoint corresponds to the SoftKD + SpectralKD result reported in Table 2 of the paper.

WideLast

WideLast keeps the first eleven transformer blocks at 192 channels and widens only the final block to 384 channels. The final block uses six attention heads, while the preceding 192-dimensional blocks use three heads. The released WideLast checkpoint corresponds to the SoftKD + MSE result reported in Table 2.

Evaluation Results

Main ImageNet-1K results

Inference cost

Results are taken from the paper. Reproduction can vary with software versions, distributed-training configuration, data access, and numerical nondeterminism.

Installation

Clone the implementation repository and install its dependencies:

git clone https://github.com/thy960112/From-Per-Image-Low-Rank-to-Encoding-Mismatch.git
cd From-Per-Image-Low-Rank-to-Encoding-Mismatch

conda create -n encoding-mismatch python=3.10 -y
conda activate encoding-mismatch
pip install -r requirements.txt
pip install huggingface_hub

Download this Hugging Face repository:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="Huiyuancs/Encoding_Mismatch",
    local_dir="./Encoding_Mismatch",
)

Expected checkpoint paths:

Encoding_Mismatch/
└── models/
    ├── Lift/
    │   └── pytorch_model.pth
    └── WideLast/
        └── pytorch_model.pth

Data Preparation

The evaluation and training scripts expect ImageNet-1K in the standard ImageFolder layout:

/path/to/ILSVRC/
├── train/
└── val/

ImageNet is not redistributed with this repository. Users must obtain it under the dataset's applicable terms.

Evaluate the Released Checkpoints

Run these commands from the cloned GitHub repository.

Lift

cd training/lift

python main.py \
  --eval \
  --resume /absolute/path/to/Encoding_Mismatch/models/Lift/pytorch_model.pth \
  --data-path /path/to/ILSVRC \
  --model deit_tiny_patch16_224 \
  --batch-size 256 \
  --num_workers 8 \
  --use-modified-student \
  --expansion-start-layer 11 \
  --expansion-type step \
  --expansion-use-ln \
  --expansion-target-dim 384

WideLast

cd training/widelast

python main.py \
  --eval \
  --resume /absolute/path/to/Encoding_Mismatch/models/WideLast/pytorch_model.pth \
  --data-path /path/to/ILSVRC \
  --model deit_tiny_patch16_224 \
  --batch-size 256 \
  --num_workers 8 \
  --custom-arch \
  --arch-schedule heads_step

The repository's evaluation code expects a PyTorch checkpoint dictionary containing a model state dictionary.

Reproduce Training

Set the dataset and distributed-training configuration:

export DATA_PATH=/path/to/ILSVRC
export GPUS=0,1,2,3
export NPROC=4

Train the released Lift recipe:

cd training/lift
bash scripts/run_cait_softkd_spectralkd.sh

Train the released WideLast recipe:

cd training/widelast
bash scripts/run_cait_softkd_mse.sh

Additional ablation recipes are available in the training/lift/scripts/ and training/widelast/scripts/ directories.

Analysis Artifacts and Raw Data

The official GitHub repository includes:

• sample-wise SVD analysis;
• dataset-level PCA analysis;
• Spectral Energy Pattern analysis;
• SEP permutation-robustness analysis;
• scripts for regenerating figures and tables;
• prepared .npz and .csv numeric results in Raw data.

See the ICML 2026 poster for a compact visual overview.

Intended Use

The checkpoints are intended for:

• research on feature and knowledge distillation in Vision Transformers;
• analysis of teacher–student representation mismatch;
• controlled ImageNet-1K evaluation of Lift and WideLast;
• studies of low-rank structure, dataset-level subspaces, and spectral channel statistics;
• reproduction and extension of the ICML 2026 paper.

Citation

@inproceedings{tian2026encodingmismatch,
  title     = {From Per-Image Low-Rank to Encoding Mismatch:
               Rethinking Feature Distillation in Vision Transformers},
  author    = {Tian, Huiyuan and Xu, Bonan and Li, Shijian},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  year      = {2026}
}

License

This Hugging Face repository is marked as Apache-2.0. Users are responsible for complying with the licenses and terms of the code dependencies, ImageNet-1K, and any upstream assets.