TRL (Transformer Reinforcement Learning)

안녕하세요, MoonLight입니다.

​

이번 Post에서는 LLM을 Fine-tuning하고 Rreinforcement learning을 적용하는 데 사용되는 도구 모음인 TRL(Transformer Reinforcement Learning)에 대해서 알아보도록 하겠습니다.

​

 

 

 

1. LLM Alignment 필요성 ​

2017년 "Attention Is All You Need"라는 논문으로 Transformer라는 구조가 세상에 나왔고, 이 Transformer 구조를 기반으로 하는 LLM들이 방대한 데이터로 비지도 학습(Unsupervised Learning) 방식으로 사전 훈련(Pre-training)되어 광범위한 세계 지식과 추론 능력을 학습하였습니다.

​

하지만, Pre-training만으로는 LLM이 특정 다운스트림 태스크(Downstream Task)나 사용자의 복잡한 요구사항, 안전 제약 조건 등에 완벽하게 부합하도록 제어하기 어렵습니다.

​

'다운스트림(Downstream)'에 대한 내용은 아래 Post를 참고해 주시기 바랍니다

https://moonlight314.tistory.com/entry/Downstream-in-LLM

Downstream in LLM

 

​

Pre-Trained Model은 일반적인 언어 능력을 갖추지만, 특정 목적에 최적화되거나 인간의 가치 및 선호도와 정렬(Alignment)하도록 만드는 과정이 필수적입니다.

​

'정렬(Alignment)'에 대한 내용은 아래 Post를 참고해 주시기 바랍니다.

https://moonlight314.tistory.com/entry/Alignment-in-LLM

Alignment in LLM

 

​

지도 미세 조정(Supervised Fine-Tuning, SFT)은 레이블된 데이터를 사용하여 모델을 특정 태스크에 적응시키는 첫 단계이지만 , 인간의 미묘한 선호도를 반영하거나 바람직하지 않은 출력을 억제하는 등 섬세한 제어에는 한계가 있습니다.

​

예를 들면, 고객 지원 챗봇 시나리오에서 SFT만으로는 공감 능력이 부족한 건조한 답변을 생성할 수도 있습니다.

​

​

이러한 Alignment 문제를 해결하기 위해 강화학습(Reinforcement Learning, RL) 방식이 주목받고 있습니다.

​

간단하게 강화학습(RL)에 대해서 설명드리자면, RL은 에이전트(Agent)라는 객체가 환경(Environment)과 상호작용하며 상태(State)를 관찰하고 행동(Action)을 취하며, 그 결과로 받는 보상(Reward) 신호를 최대화하는 방향으로 정책(Policy)을 학습하는 순차적 의사결정 학습 방법론입니다.

​

RL의 핵심은 미지의 영역을 탐색(Exploration)하는 것과 현재 지식을 활용(Exploitation)하는 것 사이의 균형을 찾는 것입니다.

​

특히, 인간 피드백 기반 강화학습(Reinforcement Learning from Human Feedback, RLHF)은 인간이 선호하는 모델 응답에 더 높은 보상을 부여함으로써 RL Model의 인간의 의도한 방향으로 인위적으로 변경가능하다는 장점덕분에 LLM을 인간의 가치에 정렬시키는 핵심 기술로 부상했습니다.

​

RLHF는 일반적으로 SFT 모델을 기반으로, 인간 선호도 데이터를 학습한 보상 모델(Reward Model)을 구축하고, 이 보상 모델의 점수를 최대화하도록 PPO(Proximal Policy Optimization)와 같은 RL 알고리즘을 사용하여 언어 모델 정책을 미세 조정합니다.

​

이와 같이, 단지 Pre-training만된 LLM Model을 인간 피드백 기반 강화학습(Reinforcement Learning from Human Feedback, RLHF) 기법을 적용하여 좀 더 인간적(?)인 Model로 만들 필요가 있고, 이 작업을 위해서 TRL(Transformer Reinforcement Learning)이라는 라이브러리가 필요해지게 되었습니다.

​

 

 

 

2. TRL의 등장

 

2.1. 배경

​

앞서 언급하였듯이, LLM을 보다 효과적으로 정렬하기 위해 RLHF가 제안되었다.

​

일반적인 RLHF 파이프라인은 다음과 같은 과정을 거치게 됩니다.

1. Base LLM을 SFT 데이터로 미세 조정 과정을 거칩니다.
2. 먼저 인간 선호도 데이터(예: 동일한 프롬프트에 대한 두 응답 중 어떤 것이 더 나은지에 대한 레이블)를 수집하여 보상 모델(Reward Model, RM)을 훈련시킵니다.
3. 이 RM이 예측하는 보상을 최대화하도록 SFT 모델을 PPO와 같은 RL 알고리즘으로 추가 미세 조정한다.

​

​

과정만 봐도 Fine Tuning이 쉽게 되지 않을 것 같다는 느낌이 듭니다.

구체적으로 다음과 같은 이유때문에 안정적인 Fine Tuning하기가 어렵습니다.

 

● ​복잡성 : 두 개의 개별 모델(Reward Model & Policy Model)을 훈련하고 관리해야 하며, RL 최적화 과정 자체가 여러 구성 요소(롤아웃, 평가, 정책 업데이트 등)를 포함하고 있습니다.

● 불안정성 : RL 훈련은 Hyperparameter 설정에 매우 민감하며, 훈련 과정이 불안정하여 발산하거나 최적점에서 벗어나기 쉽습니다. 특히 PPO 알고리즘 자체도 튜닝이 까다로울 수 있습니다.

● 계산 비용 : RL 단계에서 Policy Model로부터 응답을 반복적으로 샘플링하여 Reward Model로 평가하고 정책 업데이트를 계산해야 하므로 상당한 계산 자원이 필요합니다.

● Hyperparameter Tuning : 최적의 성능을 달성하기 위해 학습률, 배치 크기, KL 페널티 계수 등 수많은 Hyperparameter를 신중하게 조정해야 합니다.

● RL 고유의 어려움 : 지도 학습(Supervised Learning, SL)과 비교할 때, RL은 훈련 데이터가 최적화 과정에서 계속 변하는 정책에 의해 수집되므로 비정상성(non-stationarity) 문제가 발생하며, 이는 Transformer와 같은 함수 근사기를 훈련시키는 데 독특한 어려움을 야기합니다.
​

이러한 복잡성과 어려움 때문에 RLHF의 잠재력에도 불구하고, 이를 직접 구현하고 활용하는 것은 많은 연구자와 개발자에게 높은 진입 장벽으로 작용했습니다.

​

보기만 해도 어질어질하네요.

​

​

 

2.2. 등장

TRL은 이런 어렵고 복잡한 LLM Train을 쉽게 할 수 있도록 하기 위해서 개발되었습니다.

​

TRL은 SFT, Reward Modeling, PPO, DPO 등 다양한 Trainer를 제공하며, HuggingFace와 연동이 잘 되어 있어서 다양한 Transformer 모델을 쉽게 사용할 수 있습니다.

​

TRL 개발 초기에는 RLHF가 처음 주목받기 시작했을 때 PPO를 사용했기 때문에 주로 PPO를 지원했지만, 시간이 지남에 따라 DPO(Direct Preference Optimization)와 같은 간단하고 안정적인 Trainer도 지원하기 시작했습니다.

​

DPO는 RLHF 문제를 간단한 분류 손실 문제로 재구성하여 명시적인 보상 모델 훈련과 RL 최적화 단계를 생략함으로써 안정성과 계산 효율성을 높인 방법론입니다.

​

​

​

​

 

3. 핵심 Trainers

​

TRL이 제공하는 주요 훈련 패러다임과 관련 트레이너는 다음과 같습니다.

 

트레이너	주요 목표	핵심 입력	주요 출력/효과	핵심 메커니즘	복잡도	관련 Snippet ID
SFTTrainer	모델을 특정 태스크/스타일에 적응	레이블된 데이터셋 (지침/완성)	미세 조정된 모델	지도 학습 (Supervised Learning)	낮음	4
RewardTrainer	선호도 점수 함수 학습	선호도 데이터셋 (선택/거부 쌍)	보상 모델 (스칼라 점수 출력)	분류/회귀	중간	8
PPOTrainer	RL 보상을 통한 모델 최적화	쿼리 데이터셋, 보상 모델/함수, 참조 모델	보상에 최적화된 정책 모델 (KL 제약)	강화 학습 (PPO)	높음	8
DPOTrainer	직접 선호도를 통한 모델 최적화	선호도 데이터셋 (선택/거부 쌍), SFT 모델	선호도에 최적화된 정책 모델 (암묵적 보상)	분류 손실 (DPO)	중간-낮음	3
GRPOTrainer	상대적 보상을 통한 모델 최적화	쿼리 데이터셋, 보상 함수, 참조 모델	상대적 보상에 최적화된 정책 모델	강화 학습 (GRPO)	높음	8

각각의 Trainer의 특징을 하나씩 살펴보도록 하겠습니다.

​

 

 

​

3.1. Supervised Fine-Tuning(지도 미세 조정 , SFTTrainer)

SFTTrainer는 Pre-Trained Model을 특정 Downstream 태스크나 원하는 응답 스타일에 맞게 조정하는 과정입니다.

​

레이블된 데이터셋(예: 프롬프트-응답 쌍, 지침-결과 쌍)을 사용하여 모델의 가중치를 업데이트하는데, 이 작업은 종종 RLHF나 DPO와 같은 더 복잡한 Alignment 기법을 적용하기 전의 필수적인 첫 단계로써 필요합니다.

​

SFTTrainer의 주요 기능은 모델 및 데이터셋 로딩하고, 토큰화, 훈련 루프 실행, PEFT(Parameter-Efficient Fine-Tuning) 통합 등으로 Supervised Fine-Tuning 전체 프로세스를 크게 단순화해서 쉽게 사용할 수 있도록 해줍니다.

​

​

 

3.2 보상 모델링 (RewardTrainer)

Reward Model은 주어진 텍스트 입력에 대해 인간의 선호도나 특정 기준(예: 유용성, 무해성)을 반영하는 스칼라 보상 점수를 예측하도록 훈련된 모델입니다.

​

이 보상 점수는 PPO와 같은 RL 알고리즘을 사용하여 언어 모델 정책을 미세 조정하는 데 중요한 신호로 사용됩니다.

Reward Model의 입력은 보통 동일한 프롬프트에 대해서 "선택된(chosen)" 응답과 "거부된(rejected)" 응답 쌍으로 구성된 선호도 데이터셋을 사용합니다.

​

이 모델의 출력은 일반적으로 transformers 라이브러리의 AutoModelForSequenceClassification 인스턴스이며, 입력 텍스트에 대해 단일 숫자 점수(보상)를 출력하며, 훈련 과정에서 모델은 "선택된" 텍스트에 "거부된" 텍스트보다 더 높은 보상값을 출력하도록 Train됩니다.

이 Trainer의 또 다른 특징으로는 선호도 데이터 처리를 용이하게 하고, PEFT와의 통합을 지원하여 효율적인 훈련을 가능하게 하며, Llama 2 논문에서 사용된 것과 같은 마진 손실(margin loss) 및 보상 값 중심화(reward centering)와 같은 고급 기능을 지원합니다.

​

​

 

3.3 Proximal Policy Optimization(근접 정책 최적화, PPOTrainer)

PPO 알고리즘은 언어 모델이 생성한 응답에 대해 보상 모델이나 사용자 정의 함수로부터 얻은 보상 신호를 최대화하도록 정책(언어 모델)을 미세 조정하는 데 사용됩니다.

설명이 조금 어려운데, 여러가지 절차를 거치면서 조정과정이 진행되며 그만큼 학습하기가 까다롭다고 알려져 있습니다.

​

PPO를 이루는 중요한 개념은 크게 3가지가 있습니다.

1. 롤아웃(Rollout) : 현재 정책 모델을 사용하여 주어진 프롬프트(쿼리)에 대한 응답을 생성합니다.
2. 평가(Evaluation) : 생성된 쿼리-응답 쌍에 대해 보상 모델 또는 함수를 사용하여 스칼라 보상 값을 계산합니다.
3. 최적화(Optimization) : 계산된 보상과 KL 페널티를 사용하여 PPO 알고리즘에 따라 정책 모델의 가중치를 업데이트합니다.

 

 

PPO Trainer를 이루는 구성요소는 크게 아래 3가지가 있습니다.

1. 활성 모델 (정책 모델) : 현재 훈련되고 있는 언어 모델을 말합니다.
2. 참조 모델 : 일반적으로 SFT 단계 후의 모델을 복사하여 가중치를 고정시킨 버전입니다. KL 발산 페널티를 계산하는 기준점으로 사용되며, TRL은 create_reference_model과 같은 헬퍼 함수를 제공하여 참조 모델 생성을 용이하게 하고 있습니다.
3. 보상 모델/함수 : 생성된 응답에 대한 스칼라 보상 신호를 제공합니다. 이는 훈련된 보상 모델일 수도 있고, 응답 길이, 특정 키워드 포함 여부 등 규칙 기반의 사용자 정의 함수 등 다양한 값일 수 있습니다.

 

 

이 이외에 PPO는 KL 발산 페널티라는 중요한 개념이 있습니다.

이는 PPO 기반 RLHF의 핵심 요소 중 하나이며. 정책 모델이 보상을 최대화하는 과정에서 원래의 언어 생성 능력이나 SFT 단계에서 학습한 스타일을 잃어버리는 것을 방지하는 목적으로 사용됩니다.

정책 모델과 참조 모델의 출력 확률 분포 간의 KL 발산을 계산하여 손실 함수에 페널티 항으로 추가하며, RL 기반 트레이너(PPO, GRPO)에서 참조 모델과 KL 페널티 개념이 반복적으로 등장하는 것은, 특정 보상에 최적화하면서도 일관성 있는 언어 생성을 유지해야 하는 LLM RL의 근본적인 과제를 해결하기 위한 핵심 장치입니다.

​

​

 

3.4. Direct Preference Optimization(직접 선호도 최적화 , DPOTrainer)

DPO 알고리즘은 앞서 살펴본 PPO의 단점(무지하게 복잡하고 Train하기가 힘든점)을 극복한 방법으로써, 선호도 데이터(예: 더 선호되는 응답과 덜 선호되는 응답 쌍)만 있으면 직접 LLM을 Fine Tuning할 수 있는 패러다임입니다.

​

DPO가 PPO의 단점을 해결한 핵심 아이디어는 RLHF의 목적 함수를 재구성하여, 선호되는 응답과 거부되는 응답에 대한 간단한 분류 손실 함수로 변환하고, 이를 통해 암묵적인 보상 모델에 대한 최적 정책을 직접적으로 유도할 수 있도록 했다는 점입니다.

​

이런 아이디어를 적용해서 DPO는 PPO의 단점들을 해결했는데요, RLHF에 비해 훈련이 더 안정적이고, 계산 비용이 적으며(훈련 중 LLM 샘플링 불필요), 구현이 훨씬 간단하고, 하이퍼파라미터 튜닝의 필요성이 적습니다.

​

이것만으로도 아주 훌륭한데, 이에 더해서 성능 측면에서도 기존 RLHF 방법과 비슷하거나 특정 태스크(예: 감성 제어, 요약)에서는 더 우수한 결과를 보여주고 있습니다.

​

입력으로는 선호도 데이터셋(프롬프트, 선택된 응답, 거부된 응답)과 초기 SFT 모델이 필요하며, TRL에서는 DPOTrainer 클래스를 통해서 사용할 수 있습니다.

​

​

 

 

3.5. Group Relative Policy Optimization( 그룹 상대 정책 최적화 , GRPOTrainer )

DeepSeek-R1의 Train에 사용되었다고 알려진 GRPO는 PPO의 대안으로 제시된 온라인 RL 알고리즘으로, 특히 PPO보다 메모리 효율성이 높을 수 있다는 장점이 있습니다.

​

GRPO가 다른 알고리즘과 다른점은 동일한 프롬프트에 대해 생성된 여러 응답(completion) 그룹 내에서 각 응답의 상대적인 이점(advantage)을 계산하고 정규화하여 이를 최대화하도록 정책을 최적화한다는 점입니다.

​

이런 방법을 적용하여 Test해 본 결과, 수학적 추론과 같이 복잡하고 구조화된 영역에서 언어 모델의 추론 능력을 향상시키는 데 효과적인 것으로 알려져 있습니다.

​

TRL에서 GRPO를 사용하려면 GRPOTrainer 클래스를 통해 사용할 수 있으며, 보상 함수를 정의해야 합니다.

​

​

도움이 되셨으면 좋겠네요.

읽어주셔서 감사합니다.