Learning Task-Agnostic Representations through Multi-Teacher Distillation

TL;DR

The paper addresses the challenge of creating broadly useful representations without task-specific distillation. It introduces a task-agnostic multi-teacher distillation framework based on a majority-vote objective and an information-theoretic bound, implemented via a tractable Gaussian-kernel estimator of conditional entropy to maximize mutual information between the student and multiple teachers. The approach yields high-density embeddings that perform well across NLP, molecular, and vision downstream tasks, often outperforming task-specific or single-teacher baselines while offering favorable Pareto efficiency. The work demonstrates practical utility by releasing high-quality embedding models and showing robust generalization across modalities, with potential for cross-modal extensions in future work.

Abstract

Casting complex inputs into tractable representations is a critical step across various fields. Diverse embedding models emerge from differences in architectures, loss functions, input modalities and datasets, each capturing unique aspects of the input. Multi-teacher distillation leverages this diversity to enrich representations but often remains tailored to specific tasks. In this paper, we introduce a task-agnostic framework based on a ``majority vote" objective function. We demonstrate that this function is bounded by the mutual information between student and teachers' embeddings, leading to a task-agnostic distillation loss that eliminates dependence on task-specific labels or prior knowledge. Our evaluations across text, vision models, and molecular modeling show that our method effectively leverages teacher diversity, resulting in representations enabling better performance for a wide range of downstream tasks such as classification, clustering, or regression. Additionally, we train and release state-of-the-art embedding models, enhancing downstream performance in various modalities.

Figures (12)

• Figure 1: Unsupervised training of our student through task-agnostic distillation. The student embedder $S$ is trained to minimize the negative log-likelihood of multiple teachers' outputs conditioned on the student’s predictions. During this multi-teacher distillation procedure, both the student's weights and those of the teacher-specific Gaussian kernels $\{f_{\theta_k}\}_{k\leqslant K}$ are updated in an end-to-end fashion. Post-training, we discard the Gaussian kernels and evaluate the student embedders by freezing their weights and training a feed-forward network on their embeddings for an unseen dataset.
• Figure 2: (a) Pareto frontier size/performance in NLP. Our method (in blue) yields Pareto optimal model. (b) Global ranking of embedders on clustering and classification tasks for our xs model ($23$M). The NLL-distilled model rank $1$ in most tasks and in average, outperforming all other baselines of its weight category and closing the gap with models 10 times bigger.
• Figure 3: Ranking on the TDC ADMET tasks. Our student consistently achieves competitive performances across the evaluated tasks compared to its teachers (denoted by ${}^{(t)}$) and the other baselines, achieving the best average rank for both regression and classification tasks.
• Figure 4: Pareto frontier of vision models. The figure compares the performance of student model distilled using our method (named ViT-Student shown with color blue) with baselines (shown in yellow) across various datasets. The distilled student consistently lies on the Pareto frontier.
• Figure 5: Training loss of different students using different GNN architectures on the ZINC-250k dataset.

Theorems & Definitions (6)

• Proposition 3.1: darrin2024textttcosmicmutualinformationtaskagnostic
• Corollary 3.2: Training objective
• Remark 3.3
• Proposition A.1
• Proposition 1: Upper bound
• proof