[논문 리뷰]Training Language Models to Self-Correct via Reinforcement Learning

하고 있는 프로젝트에 적용할만한 방법론을 찾는 중 해당 논문을 발견했다.
결과적으로 적용하기 애매하다는 판단을 내렸지만 흥미로운 논문이었기에 기록한다.

본 논문의 핵심 내용은 LLM 을 활용한 self-correction 학습법이다.
구글 딥마인드에서 일주일 전 공개한 논문으로 강화학습을 활용한 언어 모델의 self-correction 능력 향상을 위한 연구이다.

목차

1. 논문 핵심 내용 요약

2. 장점 & 단점 정리

3. 코드 유무 확인

---

 

1. 논문 핵심 내용 요약

a. 문제 제기

LLM 은 올바른 답변을 할 수 있는 기본 지식이 포함되어 있지만 올바른 답변을 이끌어내지 못할때가 있다. 수학 증명의 경우, 증명할 수 있는 지식이 있지만 올바른 추론을 도출하지 못하곤 한다.

이를 해결하기 위해 이전에 수행된 self-correction 연구 방법들은 LLM 에서는 거의 효과가 없었다.

☞ [우리꺼 좋아요] 본 논문의 방법론으로 학습된 LLM 인 SCoRe 는 self-correction 에서 좋은 성능을 보였다.

 

b. 해결 방안 및 결과

• 단일 모델만으로 self-correction 을 통해 LLM이 자체적으로 오류를 감지하고 수정하는 방법
  • [Multi-turn reinforcement learning] 초기 응답을 수정하고 개선된 최종 응답을 생성하는 방법
• 수학 문제 풀이, 코딩 task 에서 LLM 의 능력을 크게 확장할 수 있었다.

 

c. 학습 방법

LLM 학습 시 self-correction 을 위해 2 Stage approach 사용

SCoRe figure

솔직히 위의 figure 와 논문의 글을 보면 이해가 안 된다. 왜 Stage 1 의 figure 가 저렇게 생겼을까 싶다....

그래서 논문의 글과 어떤 한국인의 코드를 참고하여 파악한 학습 Process 를 그렸다.

 

SCoRe Training Process

1. Stage 1: Stage 2 학습 시 사용할 original_answer 의 성능을 높이기 위해 SFT 학습 진행
   • 논문에서는 이를 두 번째 응답에서 높은 보상을 얻을 수 있도록 모델 initialization 를 학습 하는 거라고 표현한다.
     • original_answer 와 Stage 2에서 두 번째 답변의 편차를 줄여 멀티턴 RL 이 최적화되도록 학습하기 위해 Stage 1 을 진행한다고 함.
   • SFT 학습 시 Loss 구성
     • 일반적인 next token prediction 에 대한 cross entropy loss 와 KL-divergence loss 를 합한 값을 total loss 로 사용
2. Stage 2: Multi-turn reinforcement learning 사용하여 Self-correction 을 위한 학습
   
   • 이를 통해 학습된 모델(SCoRe Model)은 초기 응답과 두 번째 응답의 성능이 개선된다고 함
   • reward bonus 는 모델이 초기 응답에 비해 두 번째 응답을 개선하는 영향을 준다고 함
     • 논문에서 such as a string-matching based answer checking function 라고 표현하는데 이를 보면 reward bonus 를 얻기 위해 answer 에 대한 exact match 를 하는 것 같음. ⇒ reward 계산을 위한 다른 모델을 사용하지 않음
     • reward bonus 는 두 번째 응답이 gt_answer 와 동일하면 1보다 큰 int 값을 준다고 함
       만약, 두 번째 응답이 초기 응답과 동일하면 패널티를 준다고 함 (음의 정수지 않을까 생각)

 

d. 성능 평가 결과

단순 질문만 했을 때에도 기존 방법론 대비 성능 향상이 있었고, Self-Correction 을 통해서 추가적인 성능 향상이 있음을 확인했다.

> Accuracy@t1: question 만 있을때의 성능

> Accuracy@t2: question + response + instruct(수정해서 답해라) 를 줬을때의 성능

• 수학, 코딩 task benchmark 에서 self-correction 성능이 SOTA 달성하였음
  • MATH dataset 에서 base model 대비 15.6% 향상 (4.4 - (-11.2) = 15.6)
  • HumanEval coding task 에서 base model 대비 9.1% 향상
• 두 번째 응답이 초기 응답보다 일관되게 좋았음을 의미한다고 함.
  • 논문에서는 t1 도 성능 개선이 많이 된다. 하지만 t2 는 더 잘 된다. 비록 correct → incorrect case 도 있지만 incorrect → correct case 가 더 많다. 라고 언급함.

---

 

2. 장점 & 단점 정리

내가 생각하는 본 논문의 장점과 단점이다.

 

a. 장점

• 수학 문제 풀이, 코딩과 같은 추론 관련 문제에서 성능 향상을 보였다.
• 학습에 필요한 게 적다. (단일 모델, 데이터셋에 gt만 있으면 됨)

 

b. 단점

• 추론 관련 task 외에는 성능 향상이 있었는지 자료가 없다.
  • Context-QA 중 검색, 요약 등과 같은 task 에서도 좋은 성능을 보일지 미지수다.
• 충분한 max_length 가 필요하다.
  • 초기 응답을 포함하여 input 이 들어가고 output 으로 두 번째 응답이 나오기 때문에 길이기 가니 context 에 대해서는 max_length 가 많이 필요할 수 있다.

---

 

3. 코드 유무 확인

우연히 방문한 링크드인에서 Dasol Kang 님의 글을 통해 해당 논문을 기반으로 작성한 코드를 확인할 수 있었다.

코드 링크: https://github.com/daje0601/Google_SCoRe/tree/main

 

코드는 직관적이고 간단하게 구성되어 있어 금방 파악할 수 있을 것이다.

---

 

마무리,,

요즘 sLLM 에 대해 많은 생각을 한다.
공공기관 및 금융기관 등 다양한 기관 및 Domain 에서 망분리를 하는 이 시점에 sLLM 을 잘 만든다는 것이 어느 정도의 메리트가 있을 것인가? 라는 생각 말이다.
해당 논문도 비슷한 맥락으로 피드백해보면 과연 어느 정도의 파급력이 있을 것인가 싶다.

그럼에도 지치지 않고 할 수 있는 것, 해야 할 것에 집중하는게 중요하니...

오늘도 열심히 학습하고 실험하며 하루를 보낸다....

 

사요나라...