Dimension-Agnostic Similarity-Distribution Distillation

Compact embeddings need teacher guidance without inheriting its vector size; match the teacher’s query-candidate similarity distribution while retaining contrastive supervision.

The short version

What it is

Candidate preference matching For each query, the student learns to assign candidates relative importance similar to the teacher, rather than copying the teacher’s embedding vectors.
Two-part training signal The student keeps its ordinary contrastive target and adds a penalty for disagreeing with the teacher’s candidate preferences.

Why it matters

The student can learn from a teacher with a different embedding size or internal design. This transfers retrieval behavior without requiring both models to use the same vector coordinates.

Before the notation, here is the idea: the student does not need to become the teacher in vector space. It can learn which candidates the teacher considers more similar to a query, while still getting the usual contrastive signal telling it which candidate is the labeled positive. That matters when you want a compact model but do not want to force it into the teacher’s embedding dimensions or internals.

Representation mismatch: The canonical example is comparing a query against a shared list of retrieval candidates. The teacher may produce one score per candidate using a large embedding, and the student may produce those scores using a smaller, differently sized embedding. The list gives both models a common object to compare: candidate preferences. This is not a claim that their vectors can be substituted for one another; it is only about their scores on that list. So what should be transferred: coordinates, or the order and relative emphasis among candidates?
A distribution, not a copied vector: For one query, imagine the teacher’s scores as marks of attention across candidate cards: a strong candidate gets more probability, a weak one less. The student is trained to put a similar amount of probability on each card. The analogy is limited: this is a normalized distribution over the batch’s candidates, not a general measure of importance outside that set. The method uses a KL divergence between teacher and student distributions over identical supports. In words, it penalizes the student when its distribution over the shared candidate list differs from the teacher’s. Why keep the labeled positive in the picture if the teacher already supplies preferences?
Two signals with different jobs: The labeled contrastive objective still tells the student which candidate is the positive; distillation adds how the teacher distributes preference across candidates. The combined loss is $L=L_{InfoNCE}+\lambda L_{KD}$ , meaning direct contrastive learning and teacher-distribution matching both shape the student. For the reported distilled stages, the weight is $0.3$ ; the teacher-distribution temperature is $0.05$ , and the contrastive temperature is $0.02$ . Distillation is applied to the 480M student in retrieval and multitask fine-tuning, not its pre-training, while the larger listed models are trained without it. So the real question is whether the extra teacher signal is worth tuning and carrying alongside the ordinary target.

Read the deep dive

Why People Wanted It

A compact embedding model has to preserve useful retrieval behavior while spending less capacity than its teacher. One tempting route is to make the student reproduce the teacher’s embedding vectors. But that asks the two models to share a coordinate system and dimension, even though retrieval ultimately uses similarities between queries and candidates. The real question is: can the student inherit the teacher’s ranking signal without copying its representation space?

The Core Move

Match the teacher’s distribution over candidate similarities, while keeping the student’s own contrastive objective. For query $i$ and candidate $j$ , let $t_{ij}$ and $s_{ij}$ be teacher and student cosine similarities. Their candidate distributions are $p^T_{ij}=\operatorname{softmax}_j(t_{ij}/\tau_{KD})$ and $p^S_{ij}=\operatorname{softmax}_j(s_{ij}/\tau_{KD})$ ; the distillation term is $D_{KL}(p^T_i\|p^S_i)$ .

Quiz

What does similarity-distribution distillation align between teacher and student?

Their network routing decisions

Their probability distributions over candidate similarity scores for each query

Their embedding coordinates, element by element

Their parameter counts

Why the Match Lives Over Candidates

Now the version for readers who want more of the machinery: the key constraint is that teacher and student scores must refer to the same query-candidate support, not to vectors with equal shapes. For a fixed query, each model produces one cosine score for every candidate in the shared set. Softmax turns those scores into distributions over that set, so the comparison is meaningful even when the embedding dimensions differ.

This is exactly what coordinate-level MSE would not provide: vector subtraction requires corresponding coordinates, and therefore equal dimensions or a learned projection. The distribution objective instead compares relative candidate preference. Its blind spot follows from the same choice: it transfers relationships among the candidates presented, not a universal coordinate system.

Why the Objective Keeps Contrastive Supervision

Here is the part that gets technical: distillation is not a replacement for the student’s retrieval target. The batch-mean InfoNCE term continues to supervise query-positive discrimination directly; a batch-mean KL term asks the student to reproduce the teacher’s distribution on the same supports. The combined objective is $L=L_{InfoNCE}+\lambda L_{KD}$ , meaning the student balances direct contrastive learning against matching the teacher’s candidate preferences.

In the reported distilled fine-tuning stages, $\lambda=0.3$ , $\tau_{KD}=0.05$ , and the contrastive temperature is $0.02$ . Distillation is used for the 480M student during retrieval fine-tuning and multitask fine-tuning, but not in its pre-training; the larger listed models are trained without it. This makes the method a targeted training signal, not a requirement that every model in a family share one objective.

Minimal Causal Code

import torch
import torch.nn.functional as F

def loss_for_batch(student_q, student_c, teacher_q, teacher_c, labels,
                   kd_temperature=0.05, contrastive_temperature=0.02,
                   weight=0.3):
    # Each candidate column is shared by teacher and student for each query.
    student_scores = F.normalize(student_q, dim=-1) @ F.normalize(student_c, dim=-1).T
    teacher_scores = F.normalize(teacher_q, dim=-1) @ F.normalize(teacher_c, dim=-1).T

    # Direct contrastive supervision preserves the labeled positive target.
    contrastive = F.cross_entropy(student_scores / contrastive_temperature, labels)

    # Teacher and student distributions are over the same candidate columns.
    teacher_dist = F.softmax(teacher_scores / kd_temperature, dim=-1)
    student_log_dist = F.log_softmax(student_scores / kd_temperature, dim=-1)
    distillation = F.kl_div(student_log_dist, teacher_dist, reduction="batchmean")

    # The two terms pull on different signals: target labels and teacher preferences.
    return contrastive + weight * distillation

What You Gain, and What You Lose

You can transfer relative retrieval preferences without matching embedding dimensions, copying internal routing, or adding a vector projection. Keeping InfoNCE also preserves direct supervision from the labeled positives. But this does not make the student’s vectors or internal computation teacher-like: only the final candidate similarity distribution is matched. The signal is tied to the candidate supports used for each batch, and the extra KL term introduces a weighting and temperature choice alongside ordinary contrastive training.

Core Takeaway

Copy the teacher’s candidate preferences, not its coordinates.