# Discrete Key-Value Bottleneck

arXiv: https://arxiv.org/abs/2207.11240 · HTML: https://mlsanity.com/paper/2207.11240/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Builds on 68 of its 96 references that are in the map.
- First posted on arXiv 22 Jul 2022.

## Builds on (most cited first, 30 of 68)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Learning without Forgetting](https://mlsanity.com/paper/1606.09282/) (2016)
- [iCaRL: Incremental Classifier and Representation Learning](https://mlsanity.com/paper/1611.07725/) (2016)
- [Learning to Prompt for Vision-Language Models](https://mlsanity.com/paper/2109.01134/) (2021)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Big Self-Supervised Models are Strong Semi-Supervised Learners](https://mlsanity.com/paper/2006.10029/) (2020)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [How transferable are features in deep neural networks?](https://mlsanity.com/paper/1411.1792/) (2014)
- [Continual Learning with Deep Generative Replay](https://mlsanity.com/paper/1705.08690/) (2017)
- [Learning to Prompt for Continual Learning](https://mlsanity.com/paper/2112.08654/) (2021)
- [SoundStream: An End-to-End Neural Audio Codec](https://mlsanity.com/paper/2107.03312/) (2021)
- [Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution](https://mlsanity.com/paper/2202.10054/) (2022)

## Built on by (most cited first, 18 of 18)

- [Aging with GRACE: Lifelong Model Editing with Discrete Key-Value Adaptors](https://mlsanity.com/paper/2211.11031/) (2022)
- [How Does Information Bottleneck Help Deep Learning?](https://mlsanity.com/paper/2305.18887/) (2023)
- [ASIF: Coupled Data Turns Unimodal Models to Multimodal Without Training](https://mlsanity.com/paper/2210.01738/) (2022)
- [Disentanglement via Latent Quantization](https://mlsanity.com/paper/2305.18378/) (2023)
- [Learning Invariant Molecular Representation in Latent Discrete Space](https://mlsanity.com/paper/2310.14170/) (2023)
- [Grounded Object Centric Learning](https://mlsanity.com/paper/2307.09437/) (2023)
- [Object-Centric Multiple Object Tracking](https://mlsanity.com/paper/2309.00233/) (2023)
- [Neural Concept Binder](https://mlsanity.com/paper/2406.09949/) (2024)
- [Unlearning via Sparse Representations](https://mlsanity.com/paper/2311.15268/) (2023)
- [An Efficient Memory Module for Graph Few-Shot Class-Incremental Learning](https://mlsanity.com/paper/2411.06659/) (2024)
- [Long-Tailed Question Answering in an Open World](https://mlsanity.com/paper/2305.06557/) (2023)
- [Skin Lesion Phenotyping via Nested Multi-modal Contrastive Learning](https://mlsanity.com/paper/2505.23709/) (2025)
- [Disentanglement of Latent Representations via Causal Interventions](https://mlsanity.com/paper/2302.00869/) (2023)
- [Interpreting the structure of multi-object representations in vision encoders](https://mlsanity.com/paper/2406.09067/) (2024)
- [On Design Choices in Similarity-Preserving Sparse Randomized Embeddings](https://mlsanity.com/paper/2501.14741/) (2024)
- [Continuous QA Learning with Structured Prompts](https://mlsanity.com/paper/2208.14602/) (2022)
- [Self-Organising Neural Discrete Representation Learning `a la Kohonen](https://mlsanity.com/paper/2302.07950/) (2023)
- [Discrete Dictionary-based Decomposition Layer for Structured Representation Learning](https://mlsanity.com/paper/2406.06976/) (2024)

## Abstract (arXiv)

Deep neural networks perform well on classification tasks where data streams are i.i.d. and labeled data is abundant. Challenges emerge with non-stationary training data streams such as continual learning. One powerful approach that has addressed this challenge involves pre-training of large encoders on volumes of readily available data, followed by task-specific tuning. Given a new task, however, updating the weights of these encoders is challenging as a large number of weights needs to be fine-tuned, and as a result, they forget information about the previous tasks. In the present work, we propose a model architecture to address this issue, building upon a discrete bottleneck containing pairs of separate and learnable key-value codes. Our paradigm will be to encode; process the representation via a discrete bottleneck; and decode. Here, the input is fed to the pre-trained encoder, the output of the encoder is used to select the nearest keys, and the corresponding values are fed to the decoder to solve the current task. The model can only fetch and re-use a sparse number of these key-value pairs during inference, enabling localized and context-dependent model updates. We theoretically investigate the ability of the discrete key-value bottleneck to minimize the effect of learning under distribution shifts and show that it reduces the complexity of the hypothesis class. We empirically verify the proposed method under challenging class-incremental learning scenarios and show that the proposed model - without any task boundaries - reduces catastrophic forgetting across a wide variety of pre-trained models, outperforming relevant baselines on this task.