Over-parameterized Student Model via Tensor Decomposition Boosted Knowledge Distillation

Summary

In this work, the authors focus on the knowledge distillation (KD) task, using tensor decomposition to enhance the performance of the student model. Leveraging the principle of overparameterization, the authors employ the Matrix Product Operator (MPO), also known as tensor train matrix, to reformulate the original weight matrix. Additionally, they propose a “distillation loss” to measure the distance between the student’s weights and the teacher’s weights. In the experiments, the proposed method is integrated with existing KD methods, and the results demonstrate a clear improvement in KD performance.

Strengths

The paper introduces an innovative application of tensor decomposition. Typically, tensor decomposition is used for dimension reduction, adhering to the low-rank principle. However, this paper utilizes tensor decomposition in a novel manner: employing MPO to construct overparameterized learning models. While I am not entirely convinced why the overparameterized MPO is superior to traditional matrix decomposition or other tensor networks, the numerical results indicate that this approach could be a promising new avenue for using tensor networks to solve more machine learning problems.

Weaknesses

The clarity of the paper needs improvement. For instance, Figure 1 is not fully comprehensible as the meaning of the arrows is not clearly explained. Additionally, the experiment settings do not clearly describe how the weights are reshaped into the higher-order tensor format.

Questions

1. Since the entire paper is based on the principle of overparameterization of the student model, it would be beneficial to explain in the preliminary section why the overparameterization principle is relevant to the KD problem. This addition would help non-experts follow the main idea of the paper more smoothly. 2. In lines 176-177, the phrase “losing its ability to think independently” is highlighted. I am confused by this statement. Could you provide more interpretation of this claim, supported with sufficient evidence? 3. Please offer more interpretation of the central and auxiliary tensors mentioned in lines 184-185, using formulas or figures, for example. I cannot clearly understand their differences from the current non-rigorous descriptions. 4. Why does using MPO provide better performance than SVD? Is it because the MPO forms more “linear layers” than SVD? Have you considered other tensor network formats such as tree tensor networks? 5. How does the selection of TT-ranks (e.g., d_k in Eq. 2) affect the performance?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper lacks a discussion on the limitations of the proposed method. I suggest that the authors address this aspect to provide a more balanced and comprehensive evaluation of their work.

Summary

This paper introduces the Over-Parameterization Distillation Framework (OPDF), which addresses performance degradation in limited-parameter student networks after knowledge distillation (KD). OPDF proposes an overparameterized student model that utilizes the tensor-decomposition technique known as matrix product operator (MPO), allowing for a significant increase in parameters during student training time without imposing additional inference time burden. Experimental validation is performed across multiple KD tasks to assess the effectiveness of the proposed technique.

Strengths

+ The proposed technique appears to be quite novel. The use of MPO to expand parameters in the student model during the KD process, along with the tensor alignment loss function to improve student model performance, introduces innovative approaches that could offer significant advantages, particularly on low-computational devices. + The proposed methodology is easy to understand, even though it includes some abstract concepts. The authors have effectively structured their methodology by first providing a high-level overview of their technique with an illustrative figure, followed by detailed explanations of their important components. + Extensive experiments are conducted across both NLP and CV tasks. Multiple knowledge distillation (KD) techniques are employed to demonstrate how their contributions is orthogonal to existing methods. Also, parameters introduced in training as well as inference process are clearly shown. + The study extensively examines the impact of overparameterization scale, learning rate, and other components of ODPF through ablation experiments. This study helps to justify the effectiveness of their technique.

Weaknesses

- The paper does not provide information on the time required for the student network's overparameterization using MPO and the contraction of decomposed matrices into the original matrix. Having this information would be important for understanding the practical implications of the technique. - The experiments are conducted on a relatively smaller model. I am curious about the feasibility of applying this technique to LLM/VLMs with billions of parameters. A significant concern is whether their approach, which involves decomposing the teacher network, can scale effectively to such large models. A deeper discussion on this topic would provide valuable insights.

Questions

Please refer to Weaknesses section

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

NA

Summary

The authors propose to start with an over-parameterised student model. This is realised using high-order tensors that can reconstruct the original parameter matrices. The idea is that this over-parameterised model will benefit more from knowledge distillation.

Strengths

The ideas is quite interesting/novel for over-parameterising the student during training but preserving the inference parameters (since the higher-order tensors can just be contracted to reconstruct the original weights).

Weaknesses

TinyViT-5M⚗ achieves 80.7% top-1 with 5.4M parameters using a just a distillation token and a logit distillation loss. Comparing this with TinyViT-5M+OPDF 80.0% top-1 and 9.9M/5.4M inference parameters, it is hard to see the benefit? if anything the additional SVD operations and reconstruction losses, make OPDF much more difficult to adopt. What is worse it that this difference is even more significant when going to the larger TinyViT models. It is understandable that transformers have been the main focus of this work, however it would be good to see some experiments with CNNs. Matrix decomposition for the transformer linear layers is easy (with SVD), but when going to higher-order tensors it will be NP-hard. All the theory presented in this paper is for generalising to arbitrary dimensions, yet all the experiments are done on 2-dimensional weights. The paper would be a lot easier to follow without the pre-mature generalisation. i.e. showing equation (4) in 2-dimensions with SVD and I am not sure if the introduction of the MPO framework is needed. L177: I am not sure what you mean by "think independantly?". Besides the point that models don't think, if the student can match the teacher fully then that is perfectly fine. The only goal is preserve the performance of the teacher for a much smaller model. L595: 160 days to pre-train TinyViT is very significant! I understand this is because ImageNet-21K is very large, but why was this chosen benchmark chosen over, for example, a more standard KD benchmark, such as that in DeIT [1] or CRD [2]. Is 160 days an expected scale of time for the number of GPUs used here? [1] Training data-efficient image transformers & distillation through attention ICML 2021 [2] Contrastive Representation Distillation ICLR 2020 Small points/spelling: Fig 1a "Batchs -> Batches"

Questions

What is this normalization step in figure 1 and Alg S.1 L9? and its importance? If it is just to preserve the scale and stop exploding values after reconstruction, this could be explained in the text a bit.

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have adequately addressed these limitations in the checklist.

Summary

This paper proposes a novel over-parameterization framework designed to enhance the effectiveness of knowledge distillation. This framework employs MPO as a tensor decomposition technique to expand small models into larger ones to give the student model more capacity. Moreover, to enhance the effectiveness of knowledge distillation, a tensor constraint loss is introduced to align the teacher and student model. Extensive experiments verify the superiority of the method.

Strengths

1. Enhancing the capacity of the student model through tensor decomposition is novel and it does not incur additional inference overhead. 2. The loss constraint for aligning the auxiliary tensors between the student and teacher models is also quite different from the conventional logit or feature matching, representing a new approach of matching in KD. 3. The experiments are comprehensive; the authors test the method on many benchmarks in both CV and NLP, proving the effectiveness of the method.

Weaknesses

1. I would like to know how much additional distillation cost (memory and time cost) will be incurred by introducing such tensor decomposition technique and loss constraint for aligning the auxiliary tensors between the student and teacher models. 2. Especially in the era of LLM, the high cost of distillation often limits its application. Methods like LoRA have been proposed to reduce the trainable model parameters. I am concerned that if this approach goes against the current mainstream research directions? 3. Can the authors provide some distillation results on larger models to further validate the applicability of the approach on large models?

Questions

Please refer to weaknesses.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Additional vision KD results (part 2)

The following response reports the additional KD results for more recent vision models. **Q2.The Performance of OPDF on some architecture pairs provided in CRD and several recent CNN KD methods.** **Reply:** Thanks for your great question. We have applied OPDF to several architecture pairs in CRD and recent KD methods to show the capability of our method in enhancing the performance of CNN distillation models. The new results are reported in **Table B**. It can be seen that our proposed OPDF can enhance student model performance in several architecture pairs provided in the CRD benchmark using more recent KD methods. We will include these results in the revised paper. **Table B.** Distillation results of various CNN architectures on CIFAR-100. | Teacher to Student | ResNet110 to ResNet20 | ResNet110 to ResNet32 | ResNet32x4 to ResNet8x4 | | --------- | :--------: | :--------:| :--------:| | Teacher | 74.31 | 74.31 | 79.42 | | CRD | 71.46 | 73.48 | 75.51 | | CRD+OPDF | **72.35** | **74.57** | **75.83** | | NST | 69.53 | 71.96 | 73.30 | | NST+OPDF | **70.98** | **72.34** | **74.00** | | ITRD [1] | 71.99 | 74.26 | 76.19 | | ITRD+OPDF | **72.48** | **75.92** | **77.73** | | KD+LSKD [2] | 71.99 | 74.26 | 76.19 | | KD+LSKD+OPDF | **72.33** | **74.93** |**76.89** | Since the discussion period is soon ending, this is the best result we could supply at this moment after nightless working on the new experiments. Thank you very much for your time and efforts. Your consideration of further improving the rating of our paper will be much appreciated! **Reference:** [1] Miles, et al. Information theoretic representation distillation. In British Machine Vision Conference (BMVC), 2022. [2] Sun, et al. Logit standardization in knowledge distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15731–15740, 2024.

Thanks for the Rebuttal

Dear Authors, Thank you for your effort toward addressing my concerns. My concerns have been addressed and therefore, I would like to maintain my original score.