New Open-Source Toolkit Puts Reproducibility First in Embedding Distillation

When a product-search engine matches a photograph of a shoe against millions of catalogue images, or a hospital system identifies a medication from a single snapshot, the software is not running a classifier. It is comparing vectors, embeddings, in a learned metric space where distance encodes similarity. Deploying such systems on phones, edge devices or hospital servers demands compact models, and knowledge distillation has long been the standard route to them: a small student network learns to mimic a larger teacher, recovering much of the teacher’s accuracy at a fraction of the computational cost. Yet while distillation research has flourished for a decade, the open-source tooling available to practitioners has remained stubbornly focused on classification, leaving the researchers who compress retrieval models to assemble their own pipelines by hand, often with results that cannot be reliably reproduced.

A team led by Linh Do Nhat, Dzi Lam Tran Tuan and Vinh Truong Hoang has now published EmbedKD, an open-source Python toolkit designed specifically for distilling metric embeddings, in the journal SoftwareX. The software, released under the MIT license and available on GitHub, deliberately contributes no new distillation objective. Instead, its authors argue, the real gap in the field is the lifecycle around the methods: the diagnostics that reveal whether a teacher–student pair is compatible before training, the retrieval protocols that measure whether distillation actually helped, and the deployment checks that verify the model leaving the pipeline still behaves like the one that was trained. Every component, they note, could be rebuilt inside existing frameworks with moderate effort; what EmbedKD provides is that workflow, verified stage against stage.

The distinction matters because existing tools were built for a different problem. Frameworks such as torchdistill, mdistiller and KD-Lib offer config-driven pipelines, official implementations and combinations with pruning and quantization, but none ships a retrieval protocol with mean average precision and recall metrics, a gallery-query evaluation, pre- and post-training diagnostic reports, or a deployment benchmark with a parity check. Retrieval-oriented distillation methods themselves are well developed, from multi-view distillation for re-identification to pairwise-difference relational distillation, but each is released as code accompanying a single paper’s experiments. EmbedKD instead wraps the process in a six-command lifecycle: diagnose, fit, eval, extract, deploy and reproduce, each of which can be run from the command line or through a shared Python API that guarantees the two paths produce identical numbers.

Reproducibility is treated as a hard requirement rather than an aspiration. Every execution records its full state, including configuration, random seed, software versions, code commit and hardware, and all published results trace back to these logs. The reproduction command grades its output against tolerance-checked expected values and exits with a distinct error code on any metric drift, allowing continuous-integration systems to gate on it. In one demonstration, retraining the reference experiment returned results identical to the original run in every digit rather than merely within tolerance. The released version is archived on Zenodo with checksums for every checkpoint, and the reference environment, down to driver and package versions, ships with the repository.

Architecturally, the toolkit rests on three decisions. A registry system lets users add distillation objectives, task losses and dataset adapters as first-class citizens of the configuration file without touching the core. A single declarative configuration schema merges defaults, YAML files and command-line overrides into one validated, fail-fast configuration that is dumped verbatim into the run fingerprint. And an explicit backbone policy distinguishes five validated teacher and student architectures tested in continuous integration, nineteen models that pass smoke tests, and any other timm backbone, which runs only under an explicit experimental flag with a warning rather than a silent claim of support.

The toolkit also surfaces numerical subtleties that are easy to get wrong. Because the supported objectives operate on L2-normalized embeddings, mean squared error and cosine distance are the same function up to a constant, but the reduction over the embedding dimension means that at a dimensionality of 512, the same nominal weight makes the MSE distillation term 256 times weaker than the cosine term. Experiments on the CUB-200-2011 bird dataset confirmed the algebra: at the nominal weight the MSE objective changed nothing, while at the scale-matched weight it matched cosine distillation almost exactly. The toolkit prints the effective weight of each objective when a run starts, computes all losses in fp32 outside the autocast region because geometric losses underflow in fp16 and silently freeze training, and skips and counts non-finite loss components rather than poisoning a training step.

Perhaps the most striking findings come from the diagnostics themselves, which the authors report with unusual candour. In a controlled series where only teacher quality varied, a teacher weaker than the standalone student cost the student 0.023 mAP while its representation became more aligned with the teacher, a failure mode the team calls aligned-but-worse. The sign of the distillation effect changed near a teacher–student gap of 0.02 to 0.03 and saturated after roughly 0.05. But across different student architectures the gap failed to predict the magnitude of the gain, and a stronger ConvNeXt-Tiny teacher yielded smaller gains than ResNet50 for all three students tested. Most sobering, the pre-distillation risk label, intended to warn of harmful pairs, read LOW on every rung of the controlled series including the harmful one, and its underlying similarity score shifted enough between probe splits to cross its own threshold.

The cross-domain results were equally cautionary. When students trained on one dataset were evaluated on another without fine-tuning, distillation hurt in both directions tested: a CUB-trained ResNet18 lost 0.0074 mAP on Cars196, and a Cars-trained MobileNetV3 lost 0.0320 on CUB-200, both far outside the seed spread, even though both pairs had gained in domain. The distillation term, the authors suggest, pulls the student toward a teacher fitted to the source domain, and the pull that helps within a domain costs generality when the domain changes. For the deployment scenario, however, the news was better: a distilled MobileNetV3 kept 72.4 percent of its teacher’s Recall@1 while being 5.7 times smaller and eight times faster on CPU, with ONNX export verified by a mandatory parity check that recomputes retrieval metrics through the exported graph.

The authors are explicit about the limits of their evidence base: four demonstrations, one controlled series and one six-pair grid, all on a single machine, with several comparisons resting on single runs. The toolkit targets single-GPU image workloads, and the shipped training schedule trains past the best checkpoint for every backbone measured, requiring validation-based checkpoint selection. Yet the broader message is one the field is increasingly willing to hear. Whether a given embedding pair is worth distilling is today settled by trial and error, one configuration at a time; EmbedKD makes it cheap to settle by measurement instead, turning questions usually answered by folklore into routine, reproducible experiments, and, just as importantly, documenting where its own instruments cannot see.

Subject of Research: A reproducibility-first open-source toolkit for knowledge distillation of metric embeddings in image retrieval.

Article Title: EmbedKD: a reproducibility-first toolkit for knowledge distillation of metric embeddings

Article References: Do Nhat, L., Lam Tran Tuan, D., & Truong Hoang, V. (2026). EmbedKD: a reproducibility-first toolkit for knowledge distillation of metric embeddings. SoftwareX, 36, Article 103029. https://doi.org/10.1016/j.softx.2026.103029

Image Credits: AI Generated

DOI: 10.1016/j.softx.2026.103029

Keywords: knowledge distillation, metric embeddings, image retrieval, reproducibility, model compression, open-source software, deep metric learning, teacher-student networks, ONNX deployment, negative transfer, representation similarity, fine-grained recognition