REVIEW 4 major objections 6 minor 24 references
ReMem: Mutual Information-Aware Fine-tuning of Pretrained Vision Transformers for Effective Knowledge Distillation
T0 review · 4 major / 6 minor · reviewed 2026-08-06 · deepseek-v4-flash
Pith's one-line read The paper claims that fine-tuning strong pretrained Vision Transformers with large-radius sharpness-aware minimization plus downweighted top MLP blocks raises the mutual information between inputs and teacher outputs, and that this makes…
desk verdict A plausible teacher-side fix for distilling strong ViTs, but the mutual-information story leans on an unvalidated proxy and the headline gains lack error bars. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing coordinates are the information plane $(R_T, I(X;F_T))$: teacher test error versus mutual information between inputs and distillation targets, with $I$ estimated by the reconstruction loss of a decoder trained on final-layer features. The heuristic engine is the residual reweighting $x_{l+1} = (2-\alpha)\tilde{x}_l + \alpha\,\mathrm{MLP}(\tilde{x}_l)$, which by simple arithmetic downweights top MLP blocks exponentially, with effective weight $\tilde{\alpha}_l = \alpha(2-\alpha)^{l_{\mathrm{tot}}-l}$. Supporting the diagnosis is the expertness measure $e(G)$, the maximal normalized cut of the bipartite neuron-activation graph solved by spectral co-clustering, and Proposition 5.4, which bounds $I(\mathrm{MLP}(g_\leftarrow(X)); X) \leq \log_2 M + \sum_{z\in[M]} |S_z| b$.
What would settle it
Run ReMem on CIFAR-100 with the paper's decoder proxy and simultaneously estimate $I(X;F_T)$ with a direct plug-in estimator such as MINE on the same final-layer features; if the two measures disagree in direction across reweighted and SAM fine-tuned checkpoints, the proposed mechanism is not established. As an observational check, apply the reweighting component alone to a teacher with low top-MLP expertness, such as ViT-B pretrained on ImageNet-1k rather than ImageNet-21k; the mutual-information story predicts a smaller student gain from reweighting, while a generic regularization story predicts a similar one.
Extended reading notes
Core claim
On the paper's own terms, the discovery is that distillation effectiveness from a strong pretrained ViT is governed by the mutual information $I(X;F_T)$ between the input and the teacher's output features, and that this quantity is depleted specifically by the top MLP blocks. These blocks spontaneously develop high 'expertness': for similar inputs only about 1% of neurons activate, so the outputs of a sub-population collapse toward a few experts and erase the information that distinguishes individual inputs. The paper proves a bound showing that in a mixture-of-experts style MLP the output's mutual information with the input is at most the log of the number of experts plus the sparse activations inside each expert. It then shows empirically that pruning or downweighting top MLP blocks raises the mutual information proxy, that fine-tuning with SAM at perturbation radius $\rho \gtrsim 0.05$ raises it further, and that the combination, ReMem, improves the student's accuracy after distillation even when the teacher's own accuracy drops.
Load-bearing premise
Everything the paper says about why ReMem works rests on treating the reconstruction error of a decoder trained on the teacher's final features as a faithful measure of how much information those features carry about the input; if that proxy does not track true mutual information, the diagnosis and the mechanism are not established.
Editorial extensions
If this is right
- Across the 16 datasets, ReMem improves the distilled student's accuracy on 15 of them; the single exception is Diabetic Retinopathy, where accuracy is essentially unchanged at 81.9% versus 81.8%, and the largest gain is Caltech-101, where student accuracy rises from 57.6% to 77.2%.
- The advantage grows with teacher scale: average student accuracy over the 16 datasets rises from 76.1/75.0/74.0/73.7 for ViT-Ti/ViT-S/ViT-B/ViT-L under vanilla fine-tuning to 77.9/78.4/78.3/78.5 under ReMem, so small students can exploit the largest teachers.
- ReMem transfers across student architectures (ResNet-18, MobileNetV2, EfficientNetV2), distillation losses (logit matching, DIST, patient distillation), and teacher fine-tuning methods including LoRA.
- The recipe often lowers the teacher's own test accuracy while raising the student's, implying that a teacher optimized for accuracy is not the same as a teacher optimized for distillation.
- Block reweighting alone helps on small or imbalanced datasets where SAM alone is less effective, and combining both gives the best student performance.
Reading between the lines
- If the mutual-information mechanism transfers, the same top-MLP-block bottleneck should appear in transformer language models, where late feed-forward layers also show sparse activation; downweighting them may improve distillation there, a testable extension the paper does not run.
- The exponential reweighting $\tilde{\alpha}_l = \alpha(2-\alpha)^{l_{\mathrm{tot}}-l}$ is a fixed heuristic; a natural refinement is to set each layer's weight from its measured expertness or its marginal effect on the mutual information proxy, which would handle cases where the topmost block is not the most expert, a pattern the paper itself observes.
- ReMem could double as a curriculum control: ramping $\alpha$ from small to 1 over distillation would expose the student to progressively harder targets, an idea the paper mentions only as future work.
- Because teacher accuracy can fall while student accuracy rises, deployments could tune $\alpha$ against the student's validation accuracy rather than the teacher's; the paper's hyperparameter grid suggests this is feasible.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes ReMem, a teacher-side fine-tuning strategy for pretrained Vision Transformers used in knowledge distillation. ReMem combines Sharpness-Aware Minimization with a large perturbation radius and a block-reweighting heuristic that downweights top MLP blocks, as defined in Eq. (2). The authors argue that strong pretrained ViTs have low mutual information I(X; F_T) between inputs and teacher outputs, that top MLP blocks are responsible for this loss because of their 'expertness', and that ReMem increases I(X; F_T) and thereby improves student accuracy after distillation. The paper reports experiments across 16 datasets, several teacher model sizes, student architectures, and distillation algorithms, and claims consistent student gains.
Significance. If the mechanism were established, the paper would make a useful practical contribution: a simple, low-cost way to convert large pretrained ViTs into better distillation teachers, with gains that appear to grow with teacher scale. The empirical sweep is broad, and the end-task student numbers are measured independently of the mutual-information proxy, so the student-improvement claim has some standalone value. However, the mutual-information explanation and the expertness story are not currently supported: the only MI estimator is an unvalidated reconstruction-loss proxy, and the evaluation protocol lacks error bars and selects the best result over a large hyperparameter grid that includes ReMem-specific parameters. The paper merits revision rather than acceptance in its present form.
major comments (4)
- [§3, footnote 1, Appendix A.5] All mutual-information estimates in Figures 1, 2, and 4 are obtained from the reconstruction loss of a decoder trained on final-layer features. This quantity is at best an upper bound on the conditional entropy H(X|F_T), and the paper provides no evidence that the bound is tight or that it tracks true I(X;F_T) across the compared models. The confound is particularly severe for ReMem itself: Eq. (2) changes the feature computation by attenuating top MLP blocks and amplifying residual connections, so the final features become less deeply processed and may be easier for a fixed decoder to reconstruct for geometric reasons unrelated to mutual information. The 'pareto-superior information plane' claims and the MI-based motivation are therefore not established. Please validate the proxy on data with known MI or against a second estimator (for example MINE or InfoNCE), or reframe the paper around the end-task gains without claiming the MI mechanism.
- [§5.2, Proposition 5.4] Proposition 5.4 provides an upper bound I(MLP(g←(X)); X) ≤ log2 M + Σ_z |S_z| b. A smaller |S_z| lowers this upper bound, but the bound does not imply that the actual mutual information is lower; the true MI can be far below the bound regardless of sparsity. In addition, the statement that 'if the sub-population Z happens to match the class label Y, then the only information that the output contains in addition to the class label is bounded by...' is an interpretive remark, not a consequence of the bound, and no experiment tests whether the expert partitions align with class labels. Consequently, the paper does not demonstrate that high expertness causes low MI in top MLP blocks.
- [§5.1, Definition 5.3] The expertness measure is defined as the maximum cut over a bipartition, but the implementation presets the number of experts to the number of classes and uses spectral co-clustering. With k fixed to the class count, the score measures how well activations cluster at that particular k, not 'how few experts would be sufficient,' which is what the MoE motivation requires. This weakens the link between the expertness values reported in Figure 5 and the sparsity argument in Proposition 5.4.
- [§6.1, Evaluation protocol; Tables 1–4, 8–10] The evaluation selects the best student over multiple teacher checkpoints, teacher and student learning rates, KD temperatures and loss weights, and ReMem-specific hyperparameters α and ρ (Table 6), and reports a single number without error bars or a validation-based selection rule. Selecting the best result on the test set can inflate ReMem's apparent advantage because ReMem adds hyperparameters to the sweep. The magnitude of some reported gains is difficult to interpret: Table 1 shows student gains as large as +19.6 on Caltech-101 and +9.4 on Flowers, while Table 8 shows a DIST run where the teacher drops from 75.1 to 57.2 on ImageNet-LT while the student improves from 40.5 to 44.3. Please report mean and standard deviation over seeds, use a fixed protocol with hyperparameters chosen on a validation split, or at least show sensitivity to the selection rule. Also, the claim of consistent improvement is overstated: Table 1 shows a tie on SVHN and a slight decrease on Retinopathy.
minor comments (6)
- [§1, Introduction] The phrase 'shear size of these models' should be 'sheer size of these models'.
- [§6.2, paragraph on efficient model architectures] 'As show in Table 3' should be 'As shown in Table 3'.
- [Table 2] The column header 'ViT-Ti' should be 'ViT-Tiny'.
- [Figure 5] The x-axis label '?-th MLP block' should be 'l-th MLP block' or 'block index'.
- [Table 6] The table lists 'Attention weight (α)' and 'Adapter reduction factor (r)', but the main text does not describe experiments using an attention-weight reweighting or adapters; please clarify or remove these entries.
- [Table 10] The table caption contains a stray ').' at the end; please remove it.
Circularity Check
No significant circularity: distilled-student accuracies are measured independently of the mutual-information proxy, and the MI-proxy concern is a validity caveat rather than a derivation that reduces to its inputs.
full rationale
The paper's central empirical claim is that fine-tuning a pretrained ViT with large-radius SAM plus downweighting of top MLP blocks (ReMem) makes the teacher more effective for knowledge distillation, as measured by downstream student accuracy. That outcome is evaluated by actually distilling students and testing them on held-out task data, which does not depend on the paper's mutual-information estimator. The MI motivation (Figures 1, 2, 4) uses decoder reconstruction loss as a proxy for I(X;F_T) (footnote 1, Appendix A.5), and this proxy is not validated; reconstruction loss is only a variational upper bound on conditional entropy and can be confounded by feature geometry, so the paper's MI claims are not as strong as stated. However, this is a measurement-validity limitation, not a circular derivation: no equation in the paper defines the claimed result in terms of the proxy, and the method is not fitted to the proxy by construction. The only self-citation (Dong et al., 2022, Soteacher) appears in related work and is not load-bearing for the paper's claims. The reported student gains are selected as the best over a hyperparameter grid that includes ReMem-specific choices and are reported without error bars, which weakens statistical strength, but best-of-grid selection is not equivalent to a fitted parameter renamed as a prediction. Because no specific reduction of a conclusion to its own inputs can be exhibited, the appropriate finding is no significant circularity.
Assumptions & free parameters
free parameters (3)
- MLP reweighting coefficient alpha =
0.8 or 0.9 (swept)
- SAM perturbation radius rho =
0.005, 0.05, or 0.5 (swept)
- Number of experts in expertness estimation =
set equal to number of classes
assumptions (4)
- domain assumption Distillation effectiveness is governed by mutual information I(X; F_T) and teacher error R_T (Equation 1).
- domain assumption Reconstruction loss of a decoder trained on final-layer features is a faithful proxy for I(X; F_T).
- domain assumption MLP blocks can be modeled as MoE with one active expert per input and sparse neuron support.
- domain assumption Spectral co-clustering with k equal to the number of classes correctly measures expertness.
Cite this review
Pith. "Pith review of ReMem: Mutual Information-Aware Fine-tuning of Pretrained Vision Transformers for Effective Knowledge Distillation." pith.science (2026). https://pith.science/paper/QJJVVSQP
@misc{pith2026250623041,
author = {Pith},
title = {Pith review of: ReMem: Mutual Information-Aware Fine-tuning of Pretrained Vision Transformers for Effective Knowledge Distillation},
year = {2026},
howpublished = {\url{https://pith.science/paper/QJJVVSQP}},
note = {Machine review of arXiv:2506.23041}
}
read the original abstract
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.
Figures
Figures from the paper (11 more)
Reference graph
Works this paper leans on
-
[3]
URL https://api.semanticscholar. org/CorpusID:235313572. Cheng, G., Han, J., and Lu, X. Remote sensing im- age scene classification: Benchmark and state of the art. Proceedings of the IEEE, 105(10):1865–1883, Oct
-
[8]
URL https://api.semanticscholar. org/CorpusID:231573431. 9 Mutual Information-Aware Fine-tuning of Pretrained Vision Transformers for Effective Knowledge Distillation Fei-Fei, L., Fergus, R., and Perona, P. Learning gen- erative visual models from few training examples: An incremental bayesian approach tested on 101 ob- ject categories. 2004 Conference on...
work page 2004
-
[10]
URL https://api.semanticscholar. org/CorpusID:222134093. Helber, P., Bischke, B., Dengel, A., and Borth, D. Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification, 2017. Hinton, G. E., Vinyals, O., and Dean, J. Distilling the knowledge in a neural network. ArXiv, abs/1503.02531, 2015. Hu, J. E., Shen, Y ., Wallis, ...
arXiv 2017
-
[11]
URL https://api.semanticscholar. org/CorpusID:235458009. Huang, G., Li, Y ., Pleiss, G., Liu, Z., Hopcroft, J. E., and Weinberger, K. Q. Snapshot ensembles: Train 1, get m for free. ArXiv, abs/1704.00109,
-
[12]
URL https://api.semanticscholar. org/CorpusID:6820006. Huang, T., You, S., Wang, F., Qian, C., and Xu, C. Knowledge distillation from a stronger teacher. ArXiv, abs/2205.10536, 2022. URL https: //api.semanticscholar.org/CorpusID: 248986690. Kaggle and EyePacs. Kaggle diabetic retinopathy detection, jul 2015. URL https://www.kaggle.com/c/ diabetic-retinopa...
arXiv 2022
-
[15]
Learning Student-Friendly Teacher Networks for Knowledge Distillation
URL https://api.semanticscholar. org/CorpusID:259138847. Liu, Z., Miao, Z., Zhan, X., Wang, J., Gong, B., and Yu, S. X. Large-scale long-tailed recognition in an open world. In IEEE Conference on Computer Vision and Pattern Recog- nition (CVPR), 2019. URL https://github.com/ zhmiao/OpenLongTailRecognition-OLTR. Loshchilov, I. and Hutter, F. Decoupled weig...
work page Pith review arXiv 2019
-
[17]
URL https://api.semanticscholar. org/CorpusID:235485156. Tan, M. and Le, Q. V . Efficientnetv2: Smaller mod- els and faster training. In International Confer- ence on Machine Learning , 2021. URL https: //api.semanticscholar.org/CorpusID: 232478903. Tishby, N. and Zaslavsky, N. Deep learning and the informa- tion bottleneck principle. 2015 IEEE Informatio...
arXiv 2021
-
[19]
URL https://api.semanticscholar. org/CorpusID:47021742. Wang, C., Yang, Q., Huang, R., Song, S., and Huang, G. Efficient knowledge distillation from model check- points. ArXiv, abs/2210.06458, 2022. URL https: //api.semanticscholar.org/CorpusID: 252846591. Wang, Y ., Ni, Z., Song, S., Yang, L., and Huang, G. Revisiting locally supervised learning: an alte...
arXiv 2022
Show all 24 references
-
[20]
org/CorpusID:231709464
URL https://api.semanticscholar. org/CorpusID:231709464. Xiao, J., Hays, J., Ehinger, K. A., Oliva, A., and Tor- ralba, A. Sun database: Large-scale scene recogni- tion from abbey to zoo. 2010 IEEE Computer Soci- ety Conference on Computer Vision and Pattern Recog- nition, pp....
2010
-
[21]
org/CorpusID:214317405
URL https://api.semanticscholar. org/CorpusID:214317405. Zhang, C., Bengio, S., and Singer, Y . Are all layers created equal? ArXiv, abs/1902.01996,
1902 arXiv
-
[22]
org/CorpusID:59606288
URL https://api.semanticscholar. org/CorpusID:59606288. Zhang, H., Ciss ´e, M., Dauphin, Y ., and Lopez-Paz, D. mixup: Beyond empirical risk minimization. ArXiv, abs/1710.09412, 2017. URL https://api. semanticscholar.org/CorpusID:3162051. Zhang, Z., Lin, Y ., Liu, Z., Li, P., ...
2017 arXiv
-
[23]
org/CorpusID:247958465
URL https://api.semanticscholar. org/CorpusID:247958465. Zhu, Y ., Liu, N., Xu, Z., Liu, X., Meng, W., Wang, L., Ou, Z., and Tang, J. Teach less, learn more: On the undistillable classes in knowledge distilla- tion. In Neural Information Processing Systems ,
-
[24]
org/CorpusID:258509000
URL https://api.semanticscholar. org/CorpusID:258509000. 11 Mutual Information-Aware Fine-tuning of Pretrained Vision Transformers for Effective Knowledge Distillation A. Additional experiment setup details A.1. Dataset specifications In table 5, we list the specifications of ...
2018
-
[2000]
org/CorpusID:8936496
URL https://api.semanticscholar. org/CorpusID:8936496. Van Horn, G., Mac Aodha, O., Song, Y ., Cui, Y ., Sun, C., Shepard, A., Adam, H., Perona, P., and Belongie, S. The inaturalist species classification and detection dataset. In The IEEE Conference on Computer Vision and Pat...
2018
-
[2004]
org/CorpusID:2156851
URL https://api.semanticscholar. org/CorpusID:2156851. Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. Sharpness-aware minimization for efficiently improving generalization. ArXiv, abs/2010.01412,
2010 arXiv
-
[2013]
org/CorpusID:16632981
URL https://api.semanticscholar. org/CorpusID:16632981. Krizhevsky, A. Learning multiple layers of features from tiny images. 2009. URL https://api. semanticscholar.org/CorpusID:18268744. Li, Z., You, C., Bhojanapalli, S., Li, D., Rawat, A. S., Reddi, S. J., Ye, K. Q., Chern, ...
2009
-
[2014]
org/CorpusID:12726540
URL https://api.semanticscholar. org/CorpusID:12726540. Chen, X., Hsieh, C.-J., and Gong, B. When vision transformers outperform resnets without pretraining or strong data augmentations. ArXiv, abs/2106.01548,
-
[2017]
doi: 10.1109/jproc.2017
ISSN 1558-2256. doi: 10.1109/jproc.2017. 2675998. URL http://dx.doi.org/10.1109/ JPROC.2017.2675998. Cho, J. H. and Hariharan, B. On the efficacy of knowledge distillation. ICCV, pp. 4793–4801, 2019. Cimpoi, M., Maji, S., Kokkinos, I., Mohamed, S., and Vedaldi, A. Describing t...
2017
-
[2018]
org/CorpusID:4555207
URL https://api.semanticscholar. org/CorpusID:4555207. Steiner, A., Kolesnikov, A., Zhai, X., Wightman, R., Uszkoreit, J., and Beyer, L. How to train your vit? data, augmentation, and regularization in vi- sion transformers. Trans. Mach. Learn. Res. , 2022,
2022
-
[2019]
org/CorpusID:214728308
URL https://api.semanticscholar. org/CorpusID:214728308. Kornblith, S., Shlens, J., and Le, Q. V . Do better ima- genet models transfer better? 2019 IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), pp. 2656–2666, 2018. URL https://api. semanticscholar.o...
2019
-
[2020]
org/CorpusID:225039882
URL https://api.semanticscholar. org/CorpusID:225039882. Fedus, W., Zoph, B., and Shazeer, N. M. Switch transform- ers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23:120:1–120:39,
-
[2021]
org/CorpusID:235376877
URL https://api.semanticscholar. org/CorpusID:235376877. Bossard, L., Guillaumin, M., and Gool, L. V . Food- 101 - mining discriminative components with random forests. In European Conference on Computer Vision ,
-
[2022]
org/CorpusID:249642077
URL https://api.semanticscholar. org/CorpusID:249642077. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for i...
2010 arXiv
-
[2023]
org/CorpusID:256808367
URL https://api.semanticscholar. org/CorpusID:256808367. Dhillon, I. S. Co-clustering documents and words using bipartite spectral graph partitioning. In Knowledge Dis- covery and Data Mining, 2001. URL https://api. semanticscholar.org/CorpusID:11847258. Dong, C., Liu, L., and...
2001 arXiv
Reviewed August 6, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.