Towards flexible perception with visual memory (2408.08172v2)

Abstract: Training a neural network is a monolithic endeavor, akin to carving knowledge into stone: once the process is completed, editing the knowledge in a network is nearly impossible, since all information is distributed across the network's weights. We here explore a simple, compelling alternative by marrying the representational power of deep neural networks with the flexibility of a database. Decomposing the task of image classification into image similarity (from a pre-trained embedding) and search (via fast nearest neighbor retrieval from a knowledge database), we build a simple and flexible visual memory that has the following key capabilities: (1.) The ability to flexibly add data across scales: from individual samples all the way to entire classes and billion-scale data; (2.) The ability to remove data through unlearning and memory pruning; (3.) An interpretable decision-mechanism on which we can intervene to control its behavior. Taken together, these capabilities comprehensively demonstrate the benefits of an explicit visual memory. We hope that it might contribute to a conversation on how knowledge should be represented in deep vision models -- beyond carving it in "stone" weights.

• The paper introduces a visual memory framework that separates image representation from stored knowledge, allowing dynamic updates.
• The methodology employs k-nearest neighbor retrieval and RankVoting to improve accuracy and scalability on benchmarks like ImageNet.
• The approach supports flexible lifelong learning and machine unlearning, offering transparent, adaptable decision-making for real-world applications.

Towards Flexible Perception with Visual Memory

The paper "Towards Flexible Perception with Visual Memory" by Geirhos et al. explores an alternative approach to the traditional static representation of knowledge in deep neural networks by incorporating a dynamic visual memory. The authors propose a system that separates image representation from visual memory, allowing for flexible and interpretable manipulation of learned knowledge.

Overview

The authors present a methodology that decomposes the task of image classification into two main components: image similarity from a pre-trained embedding and fast nearest neighbor retrieval from a knowledge database. This system offers several key capabilities:

1. Flexible Data Addition: The visual memory can integrate data at various scales, from individual samples to entire classes and billion-scale datasets.
2. Data Removal: The system supports unlearning and memory pruning, allowing the removal of specific data points.
3. Interpretable Decision Mechanism: The decision-making process is transparent, enabling interventions to control the model's behavior.

Technical Approach

Visual Memory Construction

The visual memory framework leverages a fixed pre-trained image encoder, such as DinoV2 or CLIP, to extract feature maps from training images. These feature maps, along with their corresponding labels, are stored in a database, forming the visual memory. This approach decouples the storage of image features from the model, facilitating dynamic updates to the memory without re-training.

Nearest Neighbor Retrieval for Classification

For classification, a query image is passed through the pre-trained encoder to obtain its feature map. The system retrieves the k-nearest neighbors from the visual memory based on cosine distance. The retrieved neighbors are then used to determine the query's label through several aggregation strategies, including PluralityVoting, DistanceVoting, SoftmaxVoting, and a newly proposed RankVoting.

Experimental Results

Empirical results demonstrate that RankVoting outperforms existing aggregation methods like SoftmaxVoting, sharing the advantages of robustness and improved accuracy across different pre-trained models and datasets.

• RankVoting achieves higher and stable performance over a range of neighbors (k), showing significant improvements in ImageNet validation accuracy.

The authors further validate the flexibility of this approach via several capabilities:

1. Flexible Lifelong Learning: The system can add new out-of-distribution (OOD) classes without degrading the performance on existing classes.
2. Scalability: The system scales efficiently from million-scale datasets like ImageNet to billion-scale datasets with pseudo-labeled JFT data, showing continuous improvements in both in-distribution and OOD performance.
3. Machine Unlearning: The visual memory enables straightforward and effective unlearning by removing data points, demonstrating superior performance on key unlearning metrics.
4. Memory Pruning: The system supports memory pruning, improving accuracy by removing or downweighting unreliable samples. Both hard and soft pruning methods show efficacy in enhancing image classification accuracy.
5. Progressive Classification: The system can flexibly refine its understanding of hierarchical data, improving classification accuracy at various taxonomic levels with the addition of new examples.

Implications and Future Directions

The research presented offers substantial implications for both practical applications and theoretical advancements in AI:

• Practical Applications: The ability to dynamically update and prune memory makes this approach particularly suited for real-world systems where data and requirements constantly evolve, such as personalized recommendation systems and adaptive security mechanisms.
• Theoretical Advancements: The paper introduces a paradigm shift in how knowledge representation can be made more flexible and interpretable, paving the way for more adaptive and resilient neural network models.

Future work could extend this approach to other visual tasks beyond image classification, such as object detection, segmentation, and generative modeling. Additionally, exploring methods to seamlessly update the embedding model for scenarios involving substantial distribution shifts remains an open area of research.

Conclusion

In conclusion, the incorporation of a visual memory within deep learning models as proposed by Geirhos et al. offers a promising solution to the limitations imposed by static knowledge representation. This flexibility to add, remove, and refine knowledge dynamically aligns well with the requirements of real-world applications where the constant evolution of data necessitates adaptable models. This work marks a significant step towards more intelligent, adaptable, and interpretable AI systems.