[논문 리뷰] KNN-BERT: Fine-Tuning Pre-Trained Models with KNN Classifier

Linyang Li, Demin Song외 3명|arXiv (Cornell University)|2021. 10. 06.

Adversarial Robustness in Machine Learning참고 문헌 45인용 수 12

한 줄 요약

이 논문은 사전 훈련된 BERT의 미세조정을 향상시키기 위해 표준 선형 분류기를 K-최근접 이웃(KNN) 분류기로 대체하면서, 클래스별 양성 예시와 모멘텀 대비 학습 프레임워크를 사용하여 강건하고 대비적으로 학습된 표현을 훈련하는 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 제안된 방법은 저자원 설정에서 정확도를 크게 향상시키고 강력한 악성 공격에 대한 방어 능력을 갖출 수 있는가?

주요 결과

• KNN-BERT는 풍부한 소스 설정에서 RTE 작업에서 94.3%의 테스트 정확도를 달성하여 표준 BERT 미세조정을 능가한다.
• 소수 샘플 설정의 IMDB 데이터셋에서 표준 미세조정 대비 정확도가 10% 이상 향상되어 강력한 소수 샘플 일반화 능력을 입증한다.
• Textfooler 및 Bert-Attack 악성 공격 이후에도 모델은 42.7%의 정확도를 유지하며, 표준 BERT와 선형 분류기보다 뚜렷하게 뛰어난 성능을 보인다.
• 절단 실험 결과, 가장 유사하고 가장 비유사한 양성 예를 모두 사용할 경우 클러스터링과 성능이 향상되며, 특히 RTE와 같은 다양한 작업에서 두드러진다.
• KNN에서 최적의 K는 다양한 설정에서 강건하며, 큰 K가 필요로 하지 않아 계산 비용을 줄일 수 있다.
• 융합 비율 φ = 1.0(순수 KNN)가 가장 우수한 성능을 보였지만, KNN과 선형 분류기를 융합한 경우(φ = 0.5) 추가적인 성능 향상이 나타나 상호 보완적인 강점을 보여준다.