UDON: Universal Dynamic Online distillatioN for generic image representations

This paper misses a series of very important works for embedding learning [1,2]. The [1] proposed the Matryoshka representation learning (MRL) that the OpenAI's recent embedding model adopts. MRL would unfold the embedding model into multiple dimensions such as 64, 128, 512 or more. This paper simply takes the 64 and cannot be extended to higher dimensions which significantly limited its potential in applications.

The main focus of our paper is on improving universal embedding models, following the setup proposed by [44]. Note that in the experimental setup of [44] the embedding dimension must be fixed to 64, which is the reason we chose this dimensionality. Regarding MRL, it is a technique proposed to enable adaptive embedding dimensions in deployment. We agree that this is an interesting direction, but it is not directly related to the topic explored in our paper, which proposes techniques to improve universal image embeddings in general. Combining UDON and MRL in the future would be an interesting exploration.

Tab 1 is pretty confusing. There are "Off-the-shelf", "Specialist+Oracle", "ImageNet21k pretraining" and "CLIP pretraining". It is hard to understand the exact meaning of them.

The meaning of different settings is explained in section 4.1 "Compared methods", 4.2 and the caption of Table 1. The concept of the table is taken from [44 - Table 4].

Moreover, the evaluation is not convincing. It misses some recent SOTA embedding models, such as SigLIP.

As mentioned in the "Compared methods" paragraph in Subsection 4.1 "Experimental settings", the Universal Embedding task and the corresponding UnED benchmark is only recently created, and we have done our best to include as many baselines as we can into the evaluation. We appreciate the reviewer additionally pointing us to the recent SigLIP model. We evaluated the SigLIP ViT-Base off-the-shelf model, which achieves a mean P@5 of 49.0% and a mean R@1 of 62.0% on the UnED benchmark. We will add it to the final version of the paper.

Please unify the format of subtitles. In Line 253, there is "Implementation Details" but line 235 applies "Compared methods".

Thank you for this one, good catch! We will fix this by replacing "D" with "d".

There is no ablation study over the temperature and embedding dimension, which should be critical hyperparameters to explore.

The value of the temperature hyperparameter was tuned on the validation set independently of other hyperparameters and kept fixed. For the rebuttal, we provide additional numbers alternating the temperature value in the full UDON method (IN21k pre-trained):

all of which perform significantly worse than the one used in the paper. We will add this study to the final version of the paper.

For the teacher embedding dimension, there is already an ablation in the paper, see Table 2 row 5, and corresponding paragraph "Online teacher dimensionality" in Subsection 4.3. Regarding the student embedding dimensionality, we keep it fixed to 64, given the constraints of the UnED benchmark [44], as discussed above.

Thank you for the detailed response from the authors. I will maintain my original score, as the novelty and impact of the work seem limited. The concept of multi-teacher distillation is not new, and testing solely on UnED with fixed dim size may not be sufficiently convincing.

We thank the reviewer for the comment about how the UDON method works compared to the USC baseline. We will provide some clarifications and explain our interpretation of the key components that allow UDON to produce a better universal embedding than previous work, based on the findings from the experiments of the paper.

The extra classifier heads used in UDON as teachers, mentioned by the reviewer, each consist of a shallow domain-specific projection (which produces the domain-specific embeddings) plus a linear classifier. They stem from the same backbone embedding $E_b$ as the universal embedding $E_u$ does, and they are not necessarily lightweight (given the high dimensionality of the domain-specific embeddings + the large number of classes in the different subdomains). We will refer to them as domain-specific heads from now on.

Building from USC + Dyn. Sampler for a fair comparison, appending these domain-specific heads on the side of the universal embedding that share the same backbone and training them with classification training loss from domain-specific data already results in a performance boost for the universal embedding; see the experiment in Ablation studies Table 2, rows 7 and 2. We hypothesize that these extra domain heads backpropagating through the shared backbone act as a form of regularization.

On top of that, adding the distillation losses to obtain the full UDON method results in a large performance boost (Table 2, rows 7 and 3 of the table). We hypothesize that the domain-specific embeddings have additionally captured features that the universal embedding would otherwise skip. These features are used by UDON as an extra supervisory signal (teachers) to the universal embedding in the form of similarities between embeddings and learnable class prototypes.

In UDON, not allowing the teachers to utilize their own backbone is not only efficient in terms of compute, but allows for a higher-performing distillation process. Our experiments show that the UDON teachers (the domain-specific embeddings) underperform domain-specific embeddings that utilize their own backbone (8 separate fixed teachers case) when evaluated on their own domain (see Table 4 row 1 vs 3). However, the resulting student universal embedding of UDON performs better than the corresponding universal embedding produced by the distillation from 8 separate fixed teachers (see Table 3 row 1 vs 3, subsection 4.4). We hypothesize that distilling from the UDON teachers is an easier task than distilling from different networks that utilize their own backbone, as the UDON teachers stem from the same embedding and are only a shallow transformation of it, so the domain-specific spaces they capture are more compatible with each other.

[important] The paper tackles generalization for instance recognition/retrieval systems, but the scaling axis of this question has not been studied. This is lacking since in the past few years we have witnessed scalable pre-training in terms of data and parameters being a very effective solution to many generalization problems. ... an important question is whether scaling is indeed the answer for most problems. For the off-the-shelf baselines, it can be easy to test much larger backbones (e.g. DinoV2-G, Clip-Large, …). It is interesting to see if simply scaling general-purpose pre-training would be sufficient to address the issues tackled in the paper.

We examined larger variants of the off-the-shelf models used in this work, like the CLIP ViT-Large and the DINOv2 ViT-Large, both using a 1024 dimensional embedding. We present their results below:

Those larger models are still underperforming the smaller ViT-Base UDON model with 64-dimensional embedding. This result shows that simply scaling the general-purpose foundational models is not sufficient to address the issues tackled in the paper. We will add the corresponding results and discussion in the final version of the paper.

Related to the question above, given that the performance gain between a naive unified model vs UDON is significant but not enormous, I wonder if scaling the naive single model baseline’s capacity slightly would be sufficient to address the shortcomings of training one model with many fine-grained domains.

This is a valid point raised by the reviewer. In order to examine if the performance gain achieved by UDON with a ViT-Base (over the baseline USCRR) diminishes if we scale the backbone size, we performed additional experiment with the larger backbone size of ViT-Large (IN21k pre-trained). The results are shown in the following table:

We observe that the baseline with larger backbone achieves almost the same performance as its smaller ViT-Base counterpart, indicating that larger backbone size doesn't necessarily mean better performance (note that a similar observation is made in [44] Appendix A, section A.2). Additionally, the UDON trained ViT-Large achieves a performance which is a bit better but close to the ViT-Base counterpart, indicating both the effectiveness of the UDON training procedure for the larger backbone size compared to the baseline training procedure, as well as the fact that the ViT-Base achieves a very good size-performance tradeoff. We will include this discussion in the final version of the paper.

The performance gains provided by UDON seem highly sensitive to which dataset is used for the evaluation. For example, while the dynamic sampling contribution helps two datasets, it dropped the performance of the other five. This does not mean that this contribution should be dismissed but rather we might need more work to achieve its objectives without hurting the performance of other stable benchmarks.

Assume one (meaningful) sampling strategy, say Round Robin, which oversamples some domains and undersamples others. If with a new strategy some domain is sampled more than it was before, the model is expected to perform better on that domain, and the other way around. The proposed Dynamic Sampling is not perfect, but it improves the difficult domains where previous work was performing poorly and leaves the others close to the original perfomance, and most importantly improves the average performance, moving closer towards the notion of a "Universal Image Embedding", which works well on all of the different subdomains.

depending on the number of classes the magnitude of the loss can be very different. How is this handled in equation [5]?

The Dynamic Sampling is designed to sample the difficult domains more often. There is no explicit guarantee that domains with more classes must have higher learning error. Large number of well separable classes may have smaller error that a few hard-to-distinguish classes. In general, domains with higher number of classes have higher chance to contain difficult sets of classes - and thus should be sampled more often. Therefore, the number of classes is not explicitly accounted for in equation [5]. Additionally, a number of sampling strategies and variations on the dynamic sampling have been tried, all yielding similar or worse performance on the validation set. At the moment, the Dynamic sampling operates on the level of domains. Ideally, one would like to analyze the training data and sample difficult classes (and their confusers) more often, which might be time demanding. An efficient approach for such sampling is left for future work.

Some ablations of the projections architecture as well as the design of the MLP baseline would be a welcome addition.

We performed an experiment where we replace the linear layers used as the projection for both the domain-specific teachers and the universal student by a deeper network. More specifically, we use a one hidden layer MLP with layernorm and GELU activation (standard in the literature), with the same hidden dimension as the final dimension, i.e. 256 for the teachers and 64 for the student. The obtained results are shown below (IN21k pre-trained ViT-Base):

The results indicate a significant drop compared to using the proposed linear layers. We plan to include this study in the final version of the paper.

Can you elaborate more on impact of the throughput drop in real-world industry applications?

The 20% throughput drop only impacts the model training stage. The inference time of our model is exactly the same as previous state-of-the-art models from [44], which makes our inference-time comparisons fair. For industry applications, generally the model will be trained during a development cycle when compute resources are used intensively one-off to produce a high-quality model with acceptable cost. Once the development cycle ends, the finalized model will be used for inference in production for a long period of time. For this reason, the inference time cost is the main concern when developing industry models, and a 20% training cost increase is not usually a problem, given the significant quality gains.

Also, specify how practical applications could benefit - please provide specific details.

Several practical applications can directly benefit from an improved universal image embedding. In particular, generic visual recognition systems such as [46, 1, 2] need to handle queries depicting any type of object or scene. As per news reports, these systems today receive 10B+ searches per month and their popularity continues growing. This calls for a universal embedding, since it is not scalable to handle images of different domains with specialized, per-domain models. Besides, these embeddings can be useful in many other applications. For example, retrieval-augmented generation with large language models often requires access to specific visual information from an external database, which can searched with a universal image embedding. With the ever-growing number of images in many aspects of modern life, such embeddings also become critical for searching large private photo collections or street-level photos at planet-scale, for example.