[论文解读] KNN-BERT: Fine-Tuning Pre-Trained Models with KNN Classifier

Linyang Li, Demin Song+3 more|arXiv (Cornell University)|Oct 6, 2021

Adversarial Robustness in Machine Learning参考文献 45被引用 12

一句话总结

本文提出 KNN-BERT，一种通过将标准线性分类器替换为 K-最近邻（KNN）分类器，并结合类别内正样本与动量对比学习框架来训练鲁棒、对比学习表示的方法，以增强预训练 BERT 的微调。该方法在少样本和对抗性设置下显著提升了准确率与鲁棒性。

ABSTRACT

Pre-trained models are widely used in fine-tuning downstream tasks with linear classifiers optimized by the cross-entropy loss, which might face robustness and stability problems. These problems can be improved by learning representations that focus on similarities in the same class and contradictions in different classes when making predictions. In this paper, we utilize the K-Nearest Neighbors Classifier in pre-trained model fine-tuning. For this KNN classifier, we introduce a supervised momentum contrastive learning framework to learn the clustered representations of the supervised downstream tasks. Extensive experiments on text classification tasks and robustness tests show that by incorporating KNNs with the traditional fine-tuning process, we can obtain significant improvements on the clean accuracy in both rich-source and few-shot settings and can improve the robustness against adversarial attacks. \footnote{all codes is available at https://github.com/LinyangLee/KNN-BERT}

研究动机与目标

• 解决标准交叉熵微调在预训练模型中出现的不稳定与泛化能力差的问题。
• 提升模型在低资源设置下对对抗性攻击与噪声标签的鲁棒性。
• 利用 KNN 分类器基于相似性的决策机制，实现更好的表示泛化。
• 开发一种基于类别内正样本的对比学习框架，以形成更紧密、更可分的聚类。
• 通过可学习的加权融合策略，将 KNN 与线性分类器结合，以提升推理性能。

提出的方法

• 在 BERT 微调中，用基于训练样本相似性的 K-最近邻（KNN）分类器替代标准线性分类器。
• 提出一种使用类别内正样本（同类别样本）而非数据增强的监督动量对比学习框架。
• 采用动量对比机制（类似 MoCo）维护负样本表示队列，并通过动量更新，实现大规模负样本挖掘。
• 使用双重采样策略：在每类内同时选择最相似和最不相似的正样本，以收紧类内聚类并扩大类间间隔。
• 在推理阶段，通过可学习的融合比例 φ 将 KNN 预测与线性分类器输出进行加权融合，以平衡鲁棒性与准确率。
• 使用对比损失进行表示训练，通过选定的正样本对最小化类内距离并最大化类间距离。

实验结果

研究问题

• RQ1 将线性分类器替换为 KNN 分类器是否能提升 BERT 微调中的泛化能力与鲁棒性？
• RQ2 在对比学习中使用类别内正样本是否能带来比基于数据增强的正样本更优的表示聚类效果？
• RQ3 通过加权融合策略结合 KNN 与线性分类器，对性能与鲁棒性有何影响？
• RQ4 同时选择最相似与最不相似正样本对表示质量与模型性能有何影响？
• RQ5 所提方法是否能在低资源设置下显著提升准确率并有效防御强对抗性攻击？

主要发现

• 在丰富源设置下的 RTE 任务中，KNN-BERT 达到 94.3% 的测试准确率，优于标准 BERT 微调方法。
• 在少样本 IMDB 数据集上，该方法相比标准微调准确率提升超过 10%，展现出强大的少样本泛化能力。
• 在 Textfooler 与 Bert-Attack 对抗性攻击后，模型仍保持 42.7% 的准确率，显著优于标准 BERT 与线性分类器。
• 消融实验表明，同时使用最相似与最不相似正样本可提升聚类质量与性能，尤其在 RTE 等多样化任务中效果显著。
• KNN 中最优 K 值在各类设置下均表现稳健，无需使用大 K 值，从而降低计算成本。
• 融合比例 φ = 1.0（纯 KNN）表现最佳，但将 KNN 与线性分类器结合（φ = 0.5）可进一步提升性能，表明二者具有互补优势。