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
Released Checkpoints
| Checkpoint | File | Architecture | Distillation recipe | ImageNet-1K top-1 |
|---|
| Lift | models/Lift/pytorch_model.pth | DeiT-Tiny-derived student with a retained 192→384 endpoint projector | SoftKD + SpectralKD, CaiT-S24 teacher | 77.53% |
| WideLast | models/WideLast/pytorch_model.pth | DeiT-Tiny-derived student with blocks 1–11 at 192 dimensions and block 12 at 384 dimensions | SoftKD + MSE, CaiT-S24 teacher | 78.23% |
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:
1from huggingface_hub import hf_hub_download
2
3lift_checkpoint = hf_hub_download(
4 repo_id="Huiyuancs/Encoding_Mismatch",
5 filename="models/Lift/pytorch_model.pth",
6)
7
8widelast_checkpoint = hf_hub_download(
9 repo_id="Huiyuancs/Encoding_Mismatch",
10 filename="models/WideLast/pytorch_model.pth",
11)
12
13print("Lift:", lift_checkpoint)
14print("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
| Student | Distillation objective | Top-1 accuracy |
|---|
| Distilled DeiT-Tiny baseline | Original baseline | 74.86% |
| Lift | SoftKD | 77.23% |
| Lift | SoftKD + MSE | 77.50% |
| Lift | SoftKD + SpectralKD | 77.53% |
| WideLast | SoftKD | 77.88% |
| WideLast | SoftKD + SpectralKD | 78.16% |
| WideLast | SoftKD + MSE | 78.23% |
Inference cost
| Model | Parameters | Change vs. baseline | FLOPs | Change vs. baseline |
|---|
| Distilled DeiT-Tiny baseline | 5,717,416 | — | 2.507G | — |
| Lift | 5,983,528 | +4.65% | 2.536G | +1.17% |
| WideLast | 7,239,016 | +26.61% | 3.089G | +23.20% |
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:
1git clone https://github.com/thy960112/From-Per-Image-Low-Rank-to-Encoding-Mismatch.git
2cd From-Per-Image-Low-Rank-to-Encoding-Mismatch
3
4conda create -n encoding-mismatch python=3.10 -y
5conda activate encoding-mismatch
6pip install -r requirements.txt
7pip install huggingface_hub
Download this Hugging Face repository:
1from huggingface_hub import snapshot_download
2
3snapshot_download(
4 repo_id="Huiyuancs/Encoding_Mismatch",
5 local_dir="./Encoding_Mismatch",
6)
Expected checkpoint paths:
1Encoding_Mismatch/
2└── models/
3 ├── Lift/
4 │ └── pytorch_model.pth
5 └── WideLast/
6 └── pytorch_model.pth
Data Preparation
The evaluation and training scripts expect ImageNet-1K in the standard ImageFolder layout:
1/path/to/ILSVRC/
2├── train/
3└── 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
1cd training/lift
2
3python main.py \
4 --eval \
5 --resume /absolute/path/to/Encoding_Mismatch/models/Lift/pytorch_model.pth \
6 --data-path /path/to/ILSVRC \
7 --model deit_tiny_patch16_224 \
8 --batch-size 256 \
9 --num_workers 8 \
10 --use-modified-student \
11 --expansion-start-layer 11 \
12 --expansion-type step \
13 --expansion-use-ln \
14 --expansion-target-dim 384
WideLast
1cd training/widelast
2
3python main.py \
4 --eval \
5 --resume /absolute/path/to/Encoding_Mismatch/models/WideLast/pytorch_model.pth \
6 --data-path /path/to/ILSVRC \
7 --model deit_tiny_patch16_224 \
8 --batch-size 256 \
9 --num_workers 8 \
10 --custom-arch \
11 --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:
1export DATA_PATH=/path/to/ILSVRC
2export GPUS=0,1,2,3
3export NPROC=4
Train the released Lift recipe:
1cd training/lift
2bash scripts/run_cait_softkd_spectralkd.sh
Train the released WideLast recipe:
1cd training/widelast
2bash 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
1@inproceedings{tian2026encodingmismatch,
2 title = {From Per-Image Low-Rank to Encoding Mismatch:
3 Rethinking Feature Distillation in Vision Transformers},
4 author = {Tian, Huiyuan and Xu, Bonan and Li, Shijian},
5 booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
6 year = {2026}
7}
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.