Discrete Key–Value Bottleneck (DKVB)

• DKVB is a neural architecture that uses a discrete key–value codebook to insulate a frozen encoder while enabling sparse, localized adaptation.
• It mitigates catastrophic forgetting by restricting updates to tunable value vectors and demonstrates robust performance across vision and NLP continual learning tasks.
• By quantizing encoder features into discrete regions, DKVB reduces hypothesis complexity and improves generalization under distribution shifts.

A Discrete Key–Value Bottleneck (DKVB) is a neural architecture mechanism that interposes a discrete, fixed codebook of key–value pairs between a frozen encoder and a lightweight decoder. Originally devised to mitigate catastrophic forgetting in continual learning and non-i.i.d. input streams, DKVB enforces that all adaptation occurs locally via sparse value updates, while the high-capacity encoder remains immutable. It achieves this by quantizing encoder features into discrete regions ("keys") and associating each with a tunable value vector. DKVB has been instantiated for both vision and NLP continual learning, demonstrating strong empirical performance and theoretical robustness to distributional shifts (Träuble et al., 2022, Diera et al., 2024).

1. Formal Definition and Architectural Components

A DKVB-based predictor is defined as a composition

$$f_\theta(x) = d_\delta \circ v_\gamma \circ k_\beta \circ z_\alpha(x),$$

where

• $z_\alpha$ is a pre-trained encoder $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$ (frozen during task adaptation),
• $k_\beta$ is a quantization operator mapping $z$ to key indices per "head,"
• $v_\gamma$ is a value lookup retrieving values associated with those keys,
• $d_\delta$ is a decoder producing outputs in $\mathcal{Y}$.

Key–value codebooks are defined as follows:

• $C$ independent codebooks (heads), each of size $M$.
• Each codebook $z_\alpha$0 contains keys $z_\alpha$1 (fixed after initialization) and values $z_\alpha$2 (learnable).
• Encoded features $z_\alpha$3 are projected into $z_\alpha$4 subspaces ($z_\alpha$5–dimensional each); $z_\alpha$6 with $z_\alpha$7 random Gaussian.
• Each $z_\alpha$8 is quantized to its nearest key by

$z_\alpha$9

• Corresponding values $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$0 are aggregated (commonly by averaging) and decoded for prediction.

This mechanism applies directly to frozen vision backbones (e.g., ResNet50, ViT-B/32, ConvMixer) (Träuble et al., 2022) and, with minor modifications, to frozen transformer encoders (e.g., BERT) for NLP (Diera et al., 2024). For LLMs, DKVB can partition the encoder outputs along the hidden dimension or the token dimension; variants differ in pool type (CLS, mean) and decoder type (parametric, non-parametric).

2. Initialization, Training, and Optimization

Key initialization is performed offline in an unsupervised fashion using Exponential Moving Average (EMA), ensuring maximal coverage of the encoder representation space: $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$1 when $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$2 is closest to $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$3. In vision, keys are frozen post-init and initialized using generic unlabeled datasets (e.g., CIFAR-100, ImageNet, STL-10). In NLP, cross-domain corpora (e.g., English Wikipedia) are used, yielding task-agnostic keys (Diera et al., 2024).

Value vectors and additional decoder weights are the only parameters updated during task-specific training. Gradients flow through the quantization boundary using the straight-through estimator if needed, but in standard continual learning the encoder and keys are frozen (Träuble et al., 2022, Diera et al., 2024).

Training objectives: For classification, DKVB uses standard cross-entropy: $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$4 with optional label smoothing. No replay, distillation, or regularization is required; standard SGD or weight decay applies to trainable parameters.

3. Theoretical Properties and Robustness to Distribution Shift

DKVB leverages discrete quantization to partition the encoded feature space into Voronoi cells, with only value vectors subject to adaptation. The key theoretical result is as follows (Träuble et al., 2022):

Given an arbitrary covariate shift $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$5 on $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$6, both standard and DKVB-predictors satisfy

$z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$7

where:

• $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$8 quantifies how often covariate shifts move inputs between key regions,
• $z_\alpha : \mathcal{X} \to \mathbb{R}^{m_z}$9 (Rademacher complexity) vanishes for DKVB,
• If shift $k_\beta$0 preserves assignments to key regions, $k_\beta$1.

DKVB thus provably reduces the effective hypothesis-class complexity and shields predictions from small distributional drift in the encoder space, yielding improved generalization in non-i.i.d. continual learning streams.

4. Empirical Performance in Vision and Language

Vision benchmarks (Träuble et al., 2022):

• Toy 2D tasks: DKVB perfectly retains earlier classes (zero forgetting), unlike linear probes or MLPs.
• CIFAR-10 (class-incremental, five splits × two classes, no task signals or memory):
  • DKVB final accuracies: ConvMixer 77.3% (vs. SDM-MLP 71%), ResNet50 81.5%, ViT-B/32 (CLIP) 91.3%, SwAV 66.3%, DINO 62.5%.
  • DKVB matches i.i.d. oracle performance, with almost no catastrophic forgetting.
• Large-scale: CIFAR-100 54.7% (ResNet50), 64% (CLIP), ImageNet-1K (500 splits × 2): 49.0% DKVB (class-incremental) vs. 49.9% i.i.d.
• Ablations: Key dimension optimal at 8–12; increasing codebook size and heads improves results; key initialization dataset is robust across domains.

NLP continual learning (Diera et al., 2024):

• Domain-Incremental (DIL): DKVB–NP(Generic) achieves 82.1% (CTR baseline 88.7%, frozen BERT 87.4%).
• Class-Incremental (multi-head): DKVB–NP(Oracle) 97.06%, DKVB–NP(Generic) 96.30%, on par with EWC (96.80%) and above DER++ (range 59–95%).
• Task-Type Incremental: DKVB–NP(Generic) 68.79% (CTR 72.71%).
• Single-Head CIL: R8 dataset: DKVB–NP(Generic) 81.17% (DER++ 16.7%, naïve 31.8%); R52: 47.78% (DER++ 35.8%).
• Backward Transfer: DKVB–NP(Oracle) near-zero (–0.12 on CIL), substantially better than naïve BERT (–29.8), DER++.
• Runtime efficiency: DKVB close to frozen-BERT fine-tuning, much faster (by an order of magnitude or more) than replay- or adapter-based methods (DER++, CTR).
• Overhead: Key init ~47–469s (one-off); per-epoch cost negligible.

5. Variants, Analyses, and Design Ablations

Variants explored:

• Pooling: mean vs. CLS-token, applied before or after the bottleneck (mean-pooling after bottleneck preferred in NLP).
• Axis of codebook segmentation: hidden-dimension vs. token-dimension (best results with hidden-dimension for BERT).
• Decoder: parametric (linear+softmax) vs. non-parametric (mean+softmax) — non-parametric favored for continual learning.
• Codebook size: typically 4096 per head (vision and NLP).
• Key dimension: 8–12, with ablations indicating tradeoff between separation and coverage.

Empirical ablations:

• Value vectors learned class-incrementally and i.i.d. differ by only 2–3% (mean absolute).
• UMAP analysis shows EMA keys cover the encoder feature manifold; on test data, keys cluster but do not lose class coverage.
• 70% of keys are used after training, suggesting pruning potential.
• Robustness: DKVB maintains accuracy under pixel/embedding noise up to $k_\beta$2.
• Performance is stable with generic (cross-domain) codebook initialization, reducing the need for oracle ("future-task") init (Diera et al., 2024).

6. Advantages, Limitations, and Future Directions

Advantages:

• Localization: All learning is confined to a small set of value vectors, leaving the backbone untouched.
• Sparsity: Only a small subset of keys/values updated per input; changes are input-local.
• No need for replay buffers, regularization, or distillation.
• Parameter efficiency: Number of trainable parameters scales as $k_\beta$3.
• Robustness: Provably mitigates catastrophic forgetting and generalizes under distribution shift.

Limitations:

• Underperforms in some domain-incremental (DIL) settings where knowledge transfer between domains is critical (DKVB below strong baselines in DSC).
• Requires robust, task-independent codebook initialization (generic corpora for NLP).
• Evaluated primarily in encoder-only and classification settings; extension to generative, token-level, or full encoder–decoder models is untested.

Directions for future research:

• Extension of DKVB to decoder-only and encoder–decoder models (e.g., T5).
• Dynamic, possibly data- or input-driven resizing of codebooks and dimensions per head.
• Exploration of Gumbel-softmax relaxation to allow end-to-end differentiability of keys.
• Insertion of DKVB at intermediate network layers or for token-level input/output tasks (NER, QA, MT).
• Hybrid approaches combining DKVB with light replay or parameter regularization to further improve DIL or more challenging generalization regimes (Diera et al., 2024).

7. Significance and Relationship to Broader Continual Learning

DKVB represents a modular paradigm for continual learning that addresses catastrophic forgetting through architectural rather than buffer-based or regularization means. Its hybridization of quantized, fixed encoder manifolds with localized adaptive codebooks yields provable robustness under shift and empirical state-of-the-art in class-incremental tasks. In both vision and language, DKVB achieves competitive or superior backward transfer and memory efficiency, with extremely modest additional computational or parameter overhead relative to full model fine-tuning. This positions DKVB as an influential reference point for research on modular, discrete, and bottlenecked continual learning systems (Träuble et al., 2022, Diera et al., 2024).