Towards flexible perception with visual memory

Dear Reviewer X1Pn,

Thank you for your positive review, we’re happy to hear that you found the paper “well-written” with “excellent presentation”, and appreciated the approach as a “promising strategy for visual recognition” highlighting the “potential of using visual memory in deep learning models to adapt to dynamic data and evolving knowledge”, with “strong and interesting” Gemini re-ranking results.

Please find detailed answers to your questions below.

One limitation is that … the accuracy of recognition decreases as the memory database size grows
This might be a misunderstanding. The scaling analysis from Figure 3 shows that the error rate decreases as the memory size grows, i.e. we actually observe increasing accuracy with scale - thus this is not a limitation but an advantage. Apologies if that wasn’t clear; we now added “i.e., recognition accuracy increases with memory scale” to the figure caption.

The cost of storage and inference increases with memory scale / regarding the computation cost, I wonder if there is a rough estimate of the inference time when the database size increases to 1 billion images.
Great question. Please see https://openreview.net/forum?id=HoyKFRhwMS&noteId=YIxTk8VKZg for storage & latency estimates; we're also working to get a plot of retrieval latency as a function of memory scale which we will add.

Dino representations contain rich geometric information; but for concepts where class label correlates more with semantics rather than geometry, I wonder if the adaptation capacity still holds.
This is an interesting point. One could, for instance, think of a case where a model is tasked to retrieve “smiling cats” but not “sad cats” as an example of a semantic (rather than visual/geometric) similarity. While our current work focuses on image encoders with strong visual embeddings like DinoV2, the memory framework itself is agnostic to the encoder used: for use cases where visual similarity matters more, current encoders work well; for cases where semantic similarity (like smiling/sad cats) is more important, future encoders that may better capture semantic information could readily be integrated into our system. Presumably, the only case that wouldn’t work well is expecting a model to handle both semantic and visual similarity at the same time, without telling the model which of these similarity measures to prioritize in case they disagree.

The result of using Gemini re-ranking is quite strong and interesting. I wonder if this is due to VLM's in-context learning capacity. A discussion on this can be helpful.
Great question. We believe it is indeed the case that this strong result is due to Gemini’s in-context learning capacity, since the resulting predictions from Gemini re-ranking are better than Dino nearest neighbor predictions (83.6% on IN-val), and also better than Gemini’s predictions when no in-context neighbors are provided (69.6%). In fact, using only 10 in-context neighbors improves Gemini’s accuracy from 69.6% to 87.4%, and using 50 in-context neighbors leads to 88.5% accuracy. We now explicitly mention this connection to in-context learning (added to the paragraph on Gemini re-ranking at the end of Section 2.2).

Kindly let us know whether we have addressed your questions.

The contribution of the paper is not really clear (long history of using nearest neighbor methods for image classification; a lot of research on aggregation methods)
First of all - thanks for the additional references; we have added them to our history & related work section in the introduction. We entirely agree with the reviewer that there is a long history of using kNN for core machine learning tasks. In many ways, our goal is to remind the deep learning field - which sometimes has a tendency to be myopic - that there is a host of well established techniques that could be applied to address core and central challenges in deep learning. Our work asks how such well established techniques work in a modern deep learning setting, specifically for building a visual memory. We demonstrate that one can address many of these core challenges while also providing a substantial boost in terms of a performant and practical system. For example, deep neural networks have a static knowledge representation that is hard to update, hard to unlearn, and it is hard to understand how a decision is made. We built a working proof-of-concept alternative: By building on the long history of fast nearest neighbor methods, and “marrying” them with a powerful deep learning representation (SSL features from DinoV2) and a billion-scale visual memory. Thus our main message isn’t “use RankVoting instead of SoftmaxVoting” - that’s just a technical choice. Our main message could be summarized as “Deep learning should have a flexible visual memory, because this could solve many of its problems (static knowledge, hard to update, hard to unlearn, hard to understand how a decision is made).” Our paper then experimentally demonstrates the potential of a flexible memory for those challenges. As we note in the outlook, “while the specific approach we employ here might well be improved [...], we hope that the flexible capabilities we demonstrated might inspire and contribute to a conversation on how knowledge ought to be represented in vision models.” In response to your feedback - that the main contribution needs to be articulated more clearly - we have made the following changes to the paper:

1. Expanding related work as described above.
2. Toning down description of RankVoting, instead of writing “we invent RankVoting” we note that this is similar (albeit not identical to) the approach of Guo et al. 2011 in a non-deep learning context and rephrase as: we propose using RankVoting as an alternative to the widely used SoftmaxVoting method that is considered the SOTA method in deep learning and e.g. employed by common SSL models.
3. Clarifying main goal by adding paragraph to end of introduction: “We argue that the way current deep learning models represent knowledge (static knowledge representation, hard to update, hard to unlearn, hard to understand how a decision is made) is problematic. As an alternative, we built a working proof-of-concept: By building on the long history of nearest neighbor methods, and “marrying” them with a powerful deep learning representation (such as SSL features from DinoV2) and a billion-scale visual memory.

Does that address your concern?

That new “memories” can be added/removed seems “obvious”
We agree that unlearning in the context of a nearest neighbor visual memory is straightforward, but in deep learning the unlearning problem is considered a major challenge as e.g. evidenced by the NeurIPS unlearning competition last year - https://unlearning-challenge.github.io/ - with 5K registrations and 1.9K submissions. For this reason, we believe that this may be an important opportunity for deep learning to take inspiration from the retrieval literature to build a flexible visual memory, enabling unlearning as described in Section 3.4.

The pruning experiments are interesting, but are they not rather highlighting problems with the data (labels)? Presumably, more generally, this could be an approach for improving data quality rather than specifically improving the approach presented (although of course if does reduce computational complexity).
We’re glad to hear you appreciated the pruning experiments. Indeed, a valuable byproduct of memory pruning is its potential for improving data quality. One possibility would be to use the pruning weights obtained from memory pruning as a data selection criterion in the context of dataset pruning, where recent work has shown that given a sufficiently powerful data selection metric it is possible to improve over power-law scaling (https://arxiv.org/pdf/2206.14486). We have this potential future direction to the manuscript.

Kindly let us know whether we have addressed your concerns / answered your questions.

Dear Reviewer AsSs,

Thank you for taking the time to review our manuscript. We’re happy to hear that you “like the experiments that look at OOD performance”, found the paper “well written” and believe that it “does highlight well some of the problems with current approaches to image classification”.

Linear probe setup: What is the set-up for the linear probes? How was this validated?
In order to avoid the possibility of comparing against a straw man / weak baseline, we directly used the linear probe results that were reported in the DinoV2 and CLIP papers - this wasn’t clear from the manuscript but we have now added this information. For DinoV2, the authors froze the model backbone and trained the linear layers for 12500 iterations using SGD. Instead of training a single time, they performed a full grid search sweep over three settings (output layers in {1, 4}; pooling token concatenation in {yes, no}, and 13 different learning rates), resulting in 52 linear probes. Then, the authors evaluated the ImageNet validation accuracy for all of those 52 probes and only reported the highest one, as described in Appendix B.3 of the DinoV2 paper. Some call this test set tuning or double dipping; the DinoV2 authors call it “common practice”. In contrast, the CLIP results are based on a logistic regression classifier learned using scikit-learn’s L-BFGS implementation, and hyperparameter sweeps are performed on a held-out set not used for evaluation, according to the CLIP paper. We now added a section to Appendix I (“linear probe details”) explaining the setup and link it from Table 2.

Complexity and computational efficiency: How does it scale? How long does a query take? How much storage is required?
Those are excellent questions. Please see https://openreview.net/forum?id=HoyKFRhwMS&noteId=YIxTk8VKZg for storage & latency estimates.

Does it make sense to use a fixed value of k given class imbalance / the effect of class-imbalance in the training data and what this means with respect to k and potential accuracy?
That’s an important consideration. First of all, we fully agree that it is important to develop a system that works well if we don’t know the optimal fixed value of k, and it would be bad if the system suffered from poor performance in the case of imbalanced data. We therefore analyzed whether class-imbalance has an effect on ImageNet-1K validation accuracy (when using ImageNet-1K train features from DinoV2 ViT-L/14 in in memory). The result is visualized here: https://ibb.co/t2Bj5BB. While the accuracy of a class varies depending on the class difficulty, we do not observe a systematic disadvantage for classes with fewer images, suggesting that class-imbalance is not an issue here. This pattern holds irrespective of k, at least for the range of k within [1, 100] that we consider for ImageNet. Naturally, it is not advisable to choose k > num_images_per_class (for ImageNet, this would mean k > 732).

Many of the presented results also do not show a clear win for any particular method.
Tables 4, 5, 6, 7 and 8 show very systematic wins for RankVoting; they were moved to the appendix for space reasons. As you mention in your review, choosing a fixed k is often a bit of an arbitrary choice. While we cannot choose k flexibly depending on the class - if we did, we would have to know the class in advance - we believe it’s important to have methods that generalize across different choices for k, which is what RankVoting achieves. SoftmaxVoting precisely only on par with RankVoting if we know which k to pick - which is hard to know a priori. Existing SSL literature sometimes treats k as an hyperparameter which is optimized / swept over (e.g. DinoV2 evaluates a range of different settings - on the ImageNet validation set which is also used to report final performance), which can be avoided since RankVoting is very stable and leads to high performance across different choices of k (as shown in Tables 4–8).

Could you provide more details on how users might interact with or control the decision-making process?
Definitely. As a practical example, let’s say a visual memory based system is deployed in a clinical setting, and a radiologist uploads an X-ray scan. In contrast to a typical deep learning model, where the system would output a black-box prediction without explanation, a memory-based system could say “I think there are osseous fragments, because the scan looks very similar to the following 10 scans that I’ve seen previously. Here are these similar scans.” When looking at the evidence, the radiologist could then click an “exclude from prediction” button for any of the similar scans where their domain knowledge suggests that they shouldn’t be used in the decision-making process. This would then remove or downweigh the sample, just like the effect of hard vs. soft memory pruning described in Section 3.5 of our article, thereby altering the decision boundary.

Insufficient exploration of non-image applications (e.g., text-based memory) within the flexible memory paradigm.
A recent paper (Scaling Retrieval-Based Language Models with a Trillion-Token Datastore) explores the exciting possibility of text-based memory. As implied by the title, the focus of our work is vision, but we now include a reference to this line of work.

What would be the computational cost for re-ranking with Gemini, particularly on constrained devices?
Great question. The computational cost for VLM-based re-ranking is a single VLM call. Naturally, different VLMs have a different computational footprint. Solutions for constrained devices include: 1.) using a VLM optimized for on-device inference, such as https://github.com/Meituan-AutoML/MobileVLM; 2.) hosting a more powerful VLM like Gemini on a server which is called from the device; in this case the bottleneck would no longer be the device itself but the communication bandwidth between device and server.

Could the proposed system be combined with incremental learning techniques to improve distribution shift resilience?
Thanks, that’s a great question. TL/DR: yes. There are at least three options: (1.) Since our approach uses standard model embeddings (e.g., DinoV2, CLIP) in a plug-and-play fashion, any incremental learning technique developed for standard models can simply be used to update the embedding for improved distribution shift resilience. (2.) That said, in contrast to existing models, adding new data to a visual memory does not come at the cost of poorer performance on previously seen data: since new knowledge can be added without fine-tuning, catastrophic forgetting of old data is not an issue since the data simply gets combined and the model maintains ready access to both new and old data - a clear advantage of a visual memory. (3.) Additionally, it may be possible to extend the memory in an ensemble fashion such that it can utilize information from a variety of trained models - including some updated through incremental learning techniques - to aggregate information for downstream tasks without requiring extra post-training.

How might RankVoting's performance vary with different embedding models?
Figures 2b (ImageNet) and 8b (iNaturalist) indicate that embedding model quality leads to a performance offset, but doesn’t change the quality / stability of scaling w.r.t. neighbors k. Furthermore, Tables 4–8 in the appendix show RankVoting performance across five different models. Does that address the question?

What measures could prevent sensitive or incorrect data from being memorized?
The same measures that can be applied to avoid training on sensitive or incorrect data - this is independent of the approach (classical training vs. building a memory). However, the clear advantage of using a visual memory is that it’s possible to easily remove data after the system has been trained / developed / deployed.

Could memory pruning methods be customized for specific tasks, improving model performance further?
Definitely, memory pruning can be customized. For instance, in medical imaging, one might prioritize pruning samples that are outdated, while in autonomous driving, one might prioritize pruning samples that lead to misclassification of pedestrians. Tailoring pruning strategies to specific tasks could further improve model performance and address domain-specific challenges.

How does the model handle cases where similar but incorrect neighbors dominate?
If the model is queried with a cat image, and 8 out of the nearest 10 neighbors are very similar cat images, then the model would respond with the label of those 8 images. If this label is incorrect, then that suggests either a dataset (=label) problem or an embedding model problem, both of which are independent of our approach - our visual memory benefits from better data and better embeddings.

Thanks again for the questions & suggestions!

Dear Reviewer S9w1,

Thank you for your detailed review and questions. We’re happy to hear you appreciated the “high accuracy” resulting from a combination of a database with deep learning, and its application to machine unlearning and model control at scale.

How might the visual memory approach be adapted for tasks like object detection or segmentation?
Great question - we believe it is possible to extend to other tasks. In this paper, we represent each image with a single embedding, and this single embedding is obtained by pooling features from a spatial feature map at the top layers of an image encoder in models such as Dino or CLIP. While attention pooling, average pooling, or median pooling have been used in previous work to create a single embedding from these spatial feature maps, it is possible to pool these features into multiple embedding clusters instead of a single cluster using classical or learned clustering methods. As a proof of concept, we tested object segmentation based on a visual memory of DinoV2 features in response to your question. The result is visualized here: https://ibb.co/92dM0B2. This puts features from a single image into memory (a car in the example) based on 8 feature clusters and uses these to identify similar features in the second (test) image, thereby creating a segmentation mask. Such multi-vector representations of images (which could be expanded to image + text) are a simple and natural extension to our work in this paper, and will enable tasks such as object retrieval or detection from images containing multiple objects, or provide coarse patch-level semantic segmentation of images. Even finer-grained segmentation masks could be obtained with further training of the pooling methods. We hope this provides an intuition for how other tasks could be approached.

Reliance on a fixed pretrained embedding model may limit adaptability under significant distribution shifts.
Agreed, as we note in the discussion section, strong distribution shifts may require updating the embedding just as they would require re-training or fine-tuning a traditional deep learning model without memory. While our approach works well for typical distribution shifts (cf. Table 2), if the shift would become stronger, updating the embedding model may be advantageous. Possible options include incremental fine-tuning on new data and techniques from the adaptation literature (e.g., elastic weight consolidation).

Sensitivity of SoftmaxVoting to hyperparameters could impact performance stability.
You’re absolutely right, SoftmaxVoting - the current standard in the literature - is unstable and heavily depends on choosing a good setting for k (as we quantify in Tables 4–8 in the appendix) and also for the temperature (as quantified in Figure 9c). Therefore, e.g. the DinoV2 paper performs substantial hyperparameter sweeps to find hyperparameters that work well. For this reason, we propose to use RankVoting as a cheap, stable, and well-performing alternative that is much less sensitive to hyperparameters (as shown in Tables 4–8 and Figure 9a).

Retrieval storage & latency / Potential limitations in scalability for different embedding models with high memory footprints / Are there optimizations to manage storage costs of billion-scale datasets without affecting retrieval accuracy?
Please see https://openreview.net/forum?id=HoyKFRhwMS&noteId=YIxTk8VKZg showing that the memory footprint is only 1-3% of the cost of storing a dataset. Regarding optimizations, if storage is an issue then one could simply keep features while deleting the original dataset, saving 97-99% of storage cost. If one would like to optimize further, it may be worth training an embedding model with reduced feature dimensionality which would directly translate into storage savings.

Lack of performance benchmarking across diverse dataset distributions.
Could you elaborate what you’d like to see? We evaluate on ImageNet, iNaturalist, NINCO, InageNet-A, ImageNet-R, ImageNet-Sketch, ImageNet-V2, ImageNet-ReaL, i.e. eight different distributions currently.

It would be valuable to see an analysis of how the system performs when the memory contains corrupted samples
Thanks for the excellent suggestion. We conducted an analysis of RankVoting accuracy when the memory contains corrupted samples. The result is visualized here: https://ibb.co/X5d2PVg, showing that when 10%, 20%, 30%, …, all the way to 60% (!) of the labels in the memory database are corrupted (i.e., assigned to a random class), performance stays almost the same. We’ve added this plot to the manuscript in a dedicated appendix section (App. D, “Robustness towards label corruption”) and reference it from the main paper in Section 3.1. Thanks again for this suggestion which improves the paper.

Dear Reviewer pVvh,

Thank you for your review, we are happy to hear that you found the paper presenting “strong, well-articulated arguments” for visual memory, “stimulating important discourse that could shape future research directions”, with “excellent soundness & presentation”.

Regarding your question about potential issues with the ground-truth labels in the ImageNet-A dataset, and whether it may be possible to quantify this problem: That’s a great suggestion. To address this, we performed a human experiment on a randomly selected subset of ImageNet-A (N=100) where the dataset label and the prediction from DinoV2 ViT-L14 with JFT memory disagree. We presented the image alongside the original ImageNet-A label and our model-predicted label to three human observers, asking them to identify which of the labels best describes the image (of course, without telling them which of the labels is the dataset label). The result was that in 39.3% (!) of cases (std: +/- 1.25%), the DinoV2 label was assessed as being better/more suitable than the original dataset label - i.e., roughly 2 out of 5 model “errors” are in fact dataset label errors, quantifying the ImageNet-A label quality issue we alluded to in Figure 6. This percentage can be used to estimate how correcting problematic labels influences performance. Instead of the original model’s 61.1% accuracy on ImageNet-A, due to label errors the ‘corrected’ accuracy is instead 76.4% (a delta of +15.3% in absolute terms or +25.0% in relative terms). We believe that the issue is indeed limited to ImageNet-A label quality. Thanks again for the suggestion; we have added a discussion on this important issue to a dedicated appendix section (App. M - "ImageNet-A error analysis") and link the human experiment and result from Section 3.7 of the main paper.