How to distill task-agnostic representations from many teachers?

We thank Reviewer yAYt for their detailed review and the effort they put into evaluating our work.

Questions

1. We extended the discussion on the metric space distortion (see Sec 6.2: Embedding space structure), and moved results located in the appendix to the main body of the paper (Fig.5). The metric space distortion stems from the intrinsic properties of the mutual information. The mutual information is a metric that is invariant through composition of invertible mappings (for two invertible maps $f_1$ et $f_2$, $I(X;Y) = I(f_1(X); f_2(Y))$). As a result, this objective cannot preserve structural properties of the teachers (such as the pairwise cosine similarity). This invariance is thus a limitation of our method: it only enforces the preservation of information (retrievable by arbitrarily complex schemes or extraction mechanisms). This is because of this “extraction mechanism” requirement that our results for classification tasks hold (we train a classifier to extract the relevant information on top of the embedding) but not for tasks such as clustering or sentence similarity that do not rely on an additional complex extraction mechanism. Regardless, we evaluated our models in both these settings in text. In section C.2.2, we provided the results and analysis of clustering and STS.

2. Our experiments were conducted on larger models and datasets in NLP. We evaluate models from 22 to 109M parameters (we updated the paper to present the results for all three sizes in the body of the paper). We utilized standard architectures commonly employed in practical scenarios, specifically the Snowflake family, we trained on a common large-scale embedding-oriented datasets from hugging faces comprising around 1B tokens. We evaluated the resulting embedders on the standardized MTEB benchmark: comprising 12 classifications datasets, 7 clustering datasets, and 10 STS datasets, all of which are widely accepted and commonly used for embedding evaluation.

We hope these clarifications alleviate your concerns, please let us know if you have any additional questions.

We would like to thank Reviewer qU1H for their detailed account of our work and the efforts they put into their review.

Weaknesses

1. We added details to the Section E (“Computational cost and Complexity”) to show the effect of the number of teachers on the runtime of the training, showing that an additional teacher increases the runtime by less than 1%. Our method is scalable, in text, we specifically trained models up to 100m parameters with limited resources and got good performance. (See Table 3, Fig 6). We trained on 5M text samples in text (on a SOTA/modern dataset: GIST) and on 2M samples for molecular embedding. Overall, we trained on SOTA realistic settings both in terms of model sizes and dataset sizes in text and molecular data releasing SOTA models. We edited the main table of the NLP section to reflect the results on the three model sizes that were previously relegated to the appendices.

2. In the multi-teacher setting, to the best of our knowledge, no method covers the task agnostic distillation framework. Most of the single teacher approaches cannot be directly adapted to this setting, we therefore compared our method to two baselines that could easily be adapted to the multi-teacher setting: distillations based on the cosine similarity and the L2 distance.

3. We evaluated a total of 33 datasets in text. Encompassing classification tasks, clustering tasks and STS tasks. Our theoretical results hold only on classification tasks thus these are the main results reported in the body of the paper but we provide results in clustering and STS in appendices (See Fig 6). Furthermore, we also evaluated our method on 32 datasets in molecular modeling, covering both regressions and classification tasks.

Questions

1. Our method is task-agnostic in the sense that it does not require models to be trained on specific tasks nor labels. While training the student, only views of the same datapoint through different teachers is required, building a single representation that encompasses as much information about all other representations.

2. We added a discussion and a visualization to discuss the robustness of our method on tasks containing varying amounts of samples, and for imbalanced classification datasets (Sec.B.5). We observed that while remaining competitive, as the datasets become more imbalanced, or as the amount of available sample decreases, the performance of our students and the top performing teachers become more similar. Finally, there are tasks where our student’s performance are not as competitive, for instance in molecular modeling, on ClinTox, our model achieves an accuracy of 66% in our experiments while in this same experimental setup, ThreeDInfomax achieves an accuracy of 84%. Here, it can be explained by the fact that ClinTox contains drugs that have gone through clinical tests, while our training set focuses on lead molecules, most of which have not passed the pre-clinical phase.

3. We extended the discussion on computational cost and included additional empirical results in Sec.E. Notably we provide a visualization of the computational overhead induced by the addition of multiple teachers to our distillation process.

4. Compared to other distillation techniques, our approach performs similarly in terms of time complexity. Indeed, the estimation of Gaussian Mixtures does not require significantly more parameters than a simple dense cross encoder used in previous SOTA baselines using an L2 loss or a Cosine Similarity, and these cross-encoders are significantly smaller than the backbone of the student. As a result for all of these methods, the principal bottleneck in terms of runtime is the backbone of the student (see Fig.14 in Sec.E).

5. We rely on the theoretical results on mutual information and conditional entropy estimation methods proposed in KNIFE [1]. Those results would not hold for an arbitrary class of function [2]. Mixtures of Gaussian distributions are dense in the set of probability distributions, with respect to the weak topology. By "weak topology" we mean the probabilists' weak topology, also called the topology of convergence in distribution, the vague topology, and the weak-* topology. In general, under smoothness conditions, people perform "density estimation" with Gaussian "kernel functions". I.e. convolve the empirical distribution (the finite discrete distribution of the points in your sample) with a Gaussian density. We try to approximate the unknown population density with a weighted average of Gaussians.

[1] Georg Pichler, Pierre Colombo, Malik Boudiaf, G¨unther Koliander, and Pablo Piantanida. A differential entropy estimator for training neural networks, 2022

[2] Donoho, David L. “One-Sided Inference about Functionals of a Density.” The Annals of Statistics 16, no. 4 (1988): 1390–1420.

We would like to thank Reviewer CDoW for their detailed account of our work and the efforts they put into their review.

Weakness: Related work

Thank you for pointing out these additional references. In this related work section, we refer to task-agnostic distillation, that is to say the distillation process was performed without labels, or task-specific loss. We revised the formulation of this section, to include the proposed references. Please note that in the various references proposed, 2 are task specific (the student is trained using a task specific loss). While the other citations do address task-agnostic knowledge distillation, none of them treats the multi-teacher framework. We included those references to our related work section.

While we believe they do not invalidate our claim, as they all rely on a single teacher, we still revised our claim to focus it on the multi-teacher setting: To the best of our knowledge, there are few methods that address task-agnostic representation distillation in the context of multi-teacher approaches.

For a more complete discussion on the proposed references please refer to the following general comment.

Weakness: Experimental Validation

1,2 ) We have included in the revised version of the article larger models such as DINOv2, Swin, BEiT, ViT, and PVTv2 as teachers, evaluated on CUB, FGVC-Aircraft, and DTD, in addition to CIFAR10, SVHN and STL10.

To present the expanded evaluations, we have added a visualization of the Pareto frontier between model size and performance to the main paper (Figure 2). This highlights the trade-offs achieved by our method. Additionally, the full set of results, including the new datasets and teacher models, is provided in Table 22 of the appendix for detailed comparison.

3 ) To perform downstream tasks, we chose to freeze the model’s backbone and train classifiers using its embeddings. This experimental setup (Sec 3.3) was designed specifically to assess the quality and informativeness of the embeddings we produce. This explains the lower values of all our metrics, compared to experimental setups where the full backbone is fine-tuned. We added a remark in our analysis of the results in vision to make sure this information appears clearly in the article.

Finally, we evaluated resnet18 with the suggested augmentation. We observed improvements over certain benchmarks (notably CIFAR10), even though none of them outperform the student without any augmentation. To see the results, please refer to the following table (some experiments are still running, notably the student with data augmentation, we will update the revised article in the following days).

Embedder accuracy on various benchmarks using data-augmentations:

Minor Remark.

Thank you for the remark, we modified this sentence in the revised document.

We extend our thanks to Reviewer CDoW for their thorough review of our work, and we hope our answers and these additional discussions addressed all their concerns so they can reassess our work’s quality.

We thank reviewer LqTb for evaluating our work.

Weakness 1

CelebA dataset. Our method is designed to use a variety of unlabeled data in order to distill the representations, CelebA’s training images are only part of our diverse training set to ensure generalizability. Following the reviewers' suggestions, we are now running additional experiments on CelebA to provide a more detailed evaluation, and we will include these results in the revised submission during this discussion period.

Experimental performances. The performance discrepancies you pointed out come from the evaluation setting we chose. For each embedder, we use its embeddings to train classifiers on each task (hence freezing the backbone), rather than finetuning everything end-to-end. We chose this unified setting for all experiments for clarity’s sake. It specifically highlights the informativeness of the embeddings and fits our theoretical setting that evaluates how much information can be extracted from a frozen representation. While these details were in the “Baselines and Evaluation” paragraph (Sec.3.3), we updated the discussion section of the computer vision section to clarify our setting and how it differs from standards.

Standard KD methods Traditional KD methods, such as those proposed by Hinton et al. (2015) and FitNets (Romero et al., 2015), are inherently task-specific. This means that the distillation process must be repeated for each downstream task. In contrast, our framework is task-agnostic: it requires only a single distillation. Only one distillation phase is performed to obtain representations enabling competitive performances across multiple datasets from a variety of domains.

Weakness 2

We are pleased that you appreciated our methodology and presentation of the problem setting. . While mutual information maximization has been used previously in knowledge distillation and self-supervised learning, our approach leverages it in a novel context: task-agnostic multi-teacher feature distillation. Our method and contribution significantly differ in several aspects:

• existing methods are limited to a single teacher, or use task labels to distill representations, while ours enables the distillation of varying architectures in an unsupervised manner.
• We leveraged SOTA mutual information estimator with convergence guarantees (KNIFE [1]).
• We link our training objective to the “majority-vote” of the teachers on any task (See Corollary 1).

Our experimental results demonstrate consistent performance gains compared to individual teacher models, highlighting the distinctive advantage of our proposed approach. We hope this explanation clarifies the novel aspects of our work and its contributions beyond existing methods.

We hope that our answers alleviate your concerns and that you would reconsider our work in the light of this discussion.

[1] Georg Pichler, Pierre Colombo, Malik Boudiaf, G¨unther Koliander, and Pablo Piantanida. A differential entropy estimator for training neural networks, 2022