# ReMem: Mutual Information-Aware Fine-tuning of Pretrained Vision Transformers for Effective Knowledge Distillation

arXiv: https://arxiv.org/abs/2506.23041 · HTML: https://mlsanity.com/paper/2506.23041/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in General knowledge distillation methods, one of 12 directions in Knowledge distillation.
- Builds on 33 of its 51 references that are in the map.
- First posted on arXiv 29 Jun 2025.

## Placed in

- Knowledge distillation / [General knowledge distillation methods](https://mlsanity.com/distill/general-knowledge-distillation-methods/)

## Builds on (most cited first, 30 of 33)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Describing Textures in the Wild](https://mlsanity.com/paper/1311.3618/) (2013)
- [EuroSAT: A Novel Dataset and Deep Learning Benchmark for Land Use and Land Cover Classification](https://mlsanity.com/paper/1709.00029/) (2017)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [EfficientNetV2: Smaller Models and Faster Training](https://mlsanity.com/paper/2104.00298/) (2021)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [Remote Sensing Image Scene Classification: Benchmark and State of the Art](https://mlsanity.com/paper/1703.00121/) (2017)
- [Big Transfer (BiT): General Visual Representation Learning](https://mlsanity.com/paper/1912.11370/) (2019)
- [Large-Scale Long-Tailed Recognition in an Open World](https://mlsanity.com/paper/1904.05160/) (2019)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [MINE: Mutual Information Neural Estimation](https://mlsanity.com/paper/1801.04062/) (2018)
- [Improved Knowledge Distillation via Teacher Assistant](https://mlsanity.com/paper/1902.03393/) (2019)
- [How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers](https://mlsanity.com/paper/2106.10270/) (2021)
- [Do Better ImageNet Models Transfer Better?](https://mlsanity.com/paper/1805.08974/) (2018)
- [A Large-scale Study of Representation Learning with the Visual Task Adaptation Benchmark](https://mlsanity.com/paper/1910.04867/) (2019)
- [Rotation Equivariant CNNs for Digital Pathology](https://mlsanity.com/paper/1806.03962/) (2018)
- [On the Efficacy of Knowledge Distillation](https://mlsanity.com/paper/1910.01348/) (2019)
- [Snapshot Ensembles: Train 1, get M for free](https://mlsanity.com/paper/1704.00109/) (2017)
- [When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations](https://mlsanity.com/paper/2106.01548/) (2021)
- [Knowledge distillation: A good teacher is patient and consistent](https://mlsanity.com/paper/2106.05237/) (2021)
- [Knowledge Distillation from A Stronger Teacher](https://mlsanity.com/paper/2205.10536/) (2022)
- [MoEfication: Transformer Feed-forward Layers are Mixtures of Experts](https://mlsanity.com/paper/2110.01786/) (2021)
- [The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers](https://mlsanity.com/paper/2210.06313/) (2022)
- [Revisiting Locally Supervised Learning: an Alternative to End-to-end Training](https://mlsanity.com/paper/2101.10832/) (2021)
- [Can Neural Network Memorization Be Localized?](https://mlsanity.com/paper/2307.09542/) (2023)
- [Learning Student-Friendly Teacher Networks for Knowledge Distillation](https://mlsanity.com/paper/2102.07650/) (2021)

## Abstract (arXiv)

Knowledge distillation from pretrained visual representation models offers an effective approach to improve small, task-specific production models. However, the effectiveness of such knowledge transfer drops significantly when distilling from strong models that are pretrained in a large scale. In this paper, we address this challenge for pretrained Vision Transformers (ViTs) by exploring methods to fine-tune them for more effective knowledge transfer. Motivated by the connection between mutual information and distillation effectiveness, we propose to employ mutual information-aware optimization during finetuning. For small or highly-imbalanced downstream datasets where such optimization becomes less effective, we introduce a simple yet effective heuristic of reweighting MLP blocks. This approach is inspired by our observation that top MLP blocks are primarily responsible for mutual information loss. Our method enables small student models to benefit from those pretrained models among the strongest.