[论文解读] Discrete Key-Value Bottleneck

Frederik Träuble, Anirudh Goyal+5 more|arXiv (Cornell University)|Jul 22, 2022

Domain Adaptation and Few-Shot Learning被引用 5

一句话总结

本文提出了一种离散键值瓶颈机制，通过冻结预训练编码器并仅在推理过程中选择性地更新一组稀疏的可学习值码，实现了持续学习。通过使用离散键值对来存储和检索特定任务的知识，该模型在无任务边界的情况下，在类别增量学习基准上显著减少了灾难性遗忘，且在多种预训练模型上的表现优于现有方法。

ABSTRACT

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.

研究动机与目标

• 解决训练数据分布随时间变化的持续学习场景中的灾难性遗忘问题。
• 实现在不微调预训练编码器权重的前提下，高效且局部化的模型更新。
• 通过离散瓶颈降低假设类复杂度，并在协变量偏移下提升泛化能力。
• 设计一种可重用预训练编码器知识的模型，同时通过稀疏、上下文相关的值码更新适应新任务。
• 在无任务边界标注的多种预训练模型上，通过实证验证该方法在类别增量学习中的有效性。

提出的方法

• 模型采用三阶段流程：编码、通过离散键值瓶颈处理，以及解码。
• 预训练编码器生成连续表征，这些表征被投影到 C 个头中以进行键值检索。
• 每个头在可学习的键码本中搜索最近的键，并检索对应的连续值码。
• 训练期间仅更新值码，而键和编码器保持冻结，从而实现局部化适应。
• 通过在多样化数据分布上使用指数移动平均（EMA）初始化键值码本，以确保对特征流形的广泛覆盖。
• 通过解耦键（编码）和值（解码）码，该方法降低了假设类复杂度，限制了模型容量，从而提升了泛化能力。

实验结果

研究问题

• RQ1 在无任务边界的类别增量学习中，离散键值瓶颈能否有效减少灾难性遗忘？
• RQ2 离散瓶颈架构在输入分布偏移下如何影响泛化性能？
• RQ3 将参数更新局部化到仅值码的程度，在多大程度上能保留先前任务的知识？
• RQ4 在非独立同分布（non-i.i.d.）数据流下，该模型在多种预训练编码器上的表现如何？
• RQ5 键值对的使用模式是怎样的？未使用的容量是否表明利用不足，还是反映了模型的鲁棒性？

主要发现

• 所提出的模型在所有测试的预训练模型（ResNet50、ViT、DINO、SwAV、ConvMixer）中，均显著减少了类别增量学习中的灾难性遗忘。
• 该模型在无需任务边界标注或任务特定头重新初始化的情况下，优于相关基线方法。
• 训练过程中超过70%的键值对至少被使用了一次，其中28%的键未被使用，表明保留了容量且利用稀疏。
• 在初始化阶段，键码在数据流形上分布广泛，而在目标数据集上则变得更加聚集，反映出对领域特定特征的适应。
• 由于假设类复杂度降低，该模型在分布偏移下的泛化界优于标准模型。
• 值码更新高度局部化，大多数键被多个输入检索，从而实现了高效且上下文相关的适应。