[論文レビュー] 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の予測と線形分類器の出力を、学習可能な比率φを用いて重み付き融合する。これにより、頑健性と精度のバランスを取る。
• 選択された陽性ペairを用いて、クラス内距離を最小化し、クラス間距離を最大化する対照損失を用いて表現を学習する。

実験結果

リサーチクエスチョン

• 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でもさらなる向上が得られ、相補的な強みが示された。