[논문 리뷰] On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting

Tomasz Korbak, Hady Elsahar외 2명|arXiv (Cornell University)|2022. 06. 01.

Topic Modeling인용 수 6

한 줄 요약

이 논문은 유해한 기억 상실 없이 언어 모델의 피드백 조정을 위해 보상 최적화(RM)와 분포 매칭(DM)을 통합한다. RM 내의 KL-제어를 DM의 한 형태로 재해석하고 DM에 기준선(baseline)을 도입함으로써, 감정, 주제, 성별 제어와 같은 제어 가능한 텍스트 생성 작업에서 샘플 효율성, 안정성, 제약 조건 이행성을 향상시킨다.

ABSTRACT

The availability of large pre-trained models is changing the landscape of Machine Learning research and practice, moving from a training-from-scratch to a fine-tuning paradigm. While in some applications the goal is to "nudge" the pre-trained distribution towards preferred outputs, in others it is to steer it towards a different distribution over the sample space. Two main paradigms have emerged to tackle this challenge: Reward Maximization (RM) and, more recently, Distribution Matching (DM). RM applies standard Reinforcement Learning (RL) techniques, such as Policy Gradients, to gradually increase the reward signal. DM prescribes to first make explicit the target distribution that the model is fine-tuned to approximate. Here we explore the theoretical connections between the two paradigms, and show that methods such as KL-control developed for RM can also be construed as belonging to DM. We further observe that while DM differs from RM, it can suffer from similar training difficulties, such as high gradient variance. We leverage connections between the two paradigms to import the concept of baseline into DM methods. We empirically validate the benefits of adding a baseline on an array of controllable language generation tasks such as constraining topic, sentiment, and gender distributions in texts sampled from a language model. We observe superior performance in terms of constraint satisfaction, stability and sample efficiency.

연구 동기 및 목표

• 새로운 선호도에 적응할 때 원래 능력을 상실하는 유해한 기억 상실 문제를 해결하기 위해 언어 모델의 피드백 조정에 초점을 맞춘다.
• 두 주요 접근법인 보상 최적화(RM)와 분포 매칭(DM)을 통합하여 각각의 장점을 활용한다.
• 강화학습에서 영감을 얻은 기준선 개념을 도입함으로써, DM 방법의 높은 기울기 분산 문제를 완화하여 학습 안정성을 향상시킨다.
• 시퀀스 수준의 보상으로는 구현이 어려운 분포 제약(예: 성별 또는 주제 균형)을 가능하게 한다.
• 다양한 제어 가능한 생성 작업에서 샘플 효율성, 제약 조건 이행성, 학습 안정성 측면에서 뛰어난 성능을 입증한다.

제안 방법

• KL-제어를 RM 내에서 분포 매칭(DM)의 한 형태로 재해석하여, 보상 제약 최적화와 정방향 KL 발산 최소화 사이의 이론적 동치성을 입증한다.
• 강화학습의 정책 그래เดียน트 방법에서 영감을 얻어, 기울기 분산을 줄이기 위해 분포 매칭 알고리즘에 기준선을 도입한다.
• 감정, 주제, 성별 균형 등 원하는 선호도를 표현하는 에너지 기반 모델(EBM)을 사용해 목표 분포를 정의한다.
• 목표 분포와 모델의 출력 분포 사이의 정방향 KL 발산을 최소화하기 위해 분포 정책 그래디언트(DPG)를 사용한다.
• 지표 함수와 목표 평균을 활용해 다중 특징(예: 여성 언급 비율 50% 및 스포츠 용어 포함)을 조합한 하이브리드 제약 조건을 정의한다.
• 기준선 빼기 기반 정책 그래디언트 업데이트를 조합하여, DM 기반 최적화에서 학습을 안정화시킨다.

실험 결과

연구 질문

• RQ1 보상 최적화(RM) 내의 KL-제어는 분포 매칭(DM)의 한 형태로 해석될 수 있는가?
• RQ2 언어 모델 피드백 조정 맥락에서 RM과 DM 간의 이론적 및 실용적 연결 고리는 무엇인가?
• RQ3 DM 방법에 기준선을 도입하면 학습 안정성과 샘플 효율성이 향상되는가?
• RQ4 DM 기반 방법은 텍스트 생성에서 복잡한 분포 제약(예: 성별 또는 주제 균형)을 효과적으로 이행할 수 있는가?
• RQ5 제안된 기준선 강화된 DM 방법은 표준 RM 및 DM 기준선 대비 제약 조건 이행성과 안정성 측면에서 어떻게 비교되는가?

주요 결과

• RM 내의 KL-제어는 사전 훈련된 모델과 목표 분포 사이의 정방향 KL 발산 최소화와 이론적으로 동치되며, 이는 DM 프레임워크 내에 포함됨을 보여준다.
• 제안된 DM 내 기준선은 기울기 분산을 크게 줄여, 표준 DM 방법 대비 더 안정적인 학습 동역학을 제공한다.
• 모델은 생성 텍스트에서 여성 언급 비율을 50% 유지하는 등 뛰어난 제약 조건 이행성을 달성하면서도 유창성과 다양성을 유지한다.
• 성별 균형 + 주제 포함 등의 하이브리드 제약 조건을 포함한 작업에서는, 기준선 강화된 DM 모델이 표준 DM 및 RM 기준선 대비 샘플 효율성과 안정성 측면에서 뛰어난 성능을 보인다.
• 실증 결과에 따르면, 이 방법은 표준 강화학습 기반 피드백 조정보다 원래 모델의 분포적 성질을 더 잘 유지하며, 유해한 기억 상실을 방지한다.
• 이 접근법은 감정 제어, 주제 제어, 성별 균형 생성 등 다양한 제어 가능한 생성 작업에 일반화되며, 모든 지표에서 일관된 향상을 보인다.