Multi-Task GRPO: Reliable LLM Reasoning Across Tasks

당신이 수학 문제부터 논리 수수께끼까지 온갖 종류의 퍼즐을 풀 수 있는 초스마트 로봇 비서를 훈련시킨다고 상상해 보세요. 당신은 이 로봇이 어떤 문제가 던져지더라도 혼란에 빠지거나 어려운 부분에서 실패하지 않는 진정한 제너럴리스트(generalist)가 되기를 원합니다. 인공지능의 세계에서 이것을 '추론(reasoning)'이라고 부르며, 로봇은 '거대 언어 모델(LLM)'이라 불립니다. 이 모델들을 더 낫게 만들기 위해 과학자들은 '강화 학습(Reinforcement Learning)'이라는 기술을 사용하는데, 이는 로봇이 정답을 맞혔을 때 하이파이브(보상)를 해주고 틀렸을 때 "다시 해봐"라고 부드럽게 말해주는 것과 같습니다. 이 방법 중 하나로 GRPO라는 아주 영리한 방식이 있는데, 이는 로봇이 각 시도로부터 정확히 얼마나 배워야 하는지를 결정하는 방법입니다. 하지만 여기에는 함정이 있습니다. 만약 당신이 로봇에게 한 번에 열 가지 서로 다른 유형의 퍼즐을 배우라고 요구하면, 로봇은 종종 길을 잃고 맙니다. 쉬운 퍼즐에는 매우 능숙해질 수 있지만 어려운 퍼st는 완전히 무시할 수도 있고, 특정 종류의 수수께끼에 매몰되어 다른 것들을 하는 법을 잊어버릴 수도 있습니다. 큰 질문은 이것입니다. 어떻게 하면 하나의 뇌가 모든 것에 똑같이 능숙해지도록, 즉 단 하나의 기술도 뒤처지지 않도록 훈련할 수 있을까요?

이것이 바로 **MT-GRPO(Multi-Task GRPO)**라고 불리는 새로운 방법이 해결하고자 하는 문제입니다. 표준적인 훈련 과정을 학습 속도가 매우 다른 학생들로 구성된 학급을 가르치는 교사라고 생각해 보세요. 만약 교사가 단순히 '평균적인' 학생에게만 집중한다면, 똑똑한 아이들은 지루해하고 힘들어하는 아이들은 점점 더 뒤처지게 될 것입니다. 기존의 방식은 쉬운 작업들이 훈련을 '하이재킹(탈취)'하도록 내버려 두어, 모델이 평균적으로는 좋아 보이게 만들지만 실제로는 어려운 과제들에서는 실패하게 만들곤 합니다. 이 논문의 저자들은 단순히 점수를 평균 내는 것만으로는 신뢰할 수 있는 로봇을 만들 수 없다는 것을 깨달았습니다. 당신은 학급에서 가장 성적이 낮은 학생도 함께 나아지고 있는지 확인해야 합니다.

이를 해결하기 위해 연구진은 두 부분으로 된 전략을 발명했습니다. 첫째, 그들은 동적인 '가중치 시스템'을 만들었습니다. 이는 전광판을 끊임없이 지켜보는 코치를 상상하는 것과 같습니다. 만약 로봇이 '카운트다운'(수학 게임)은 잘하지만 '제브라 퍼즐'(논리 수수께끼)에는 엉망이라면, 코치는 즉시 초점을 바꿉니다. 코치는 로봇에게 수학 카테고리의 문제보다 제브라 카테고리의 연습 문제를 더 많이 제공하고 수학 카테고리의 문제는 줄이기 시작합니다. 하지만 까다로운 점이 있습니다. 때때로 로봇이 제브라 퍼즐을 시도했지만 제대로 풀려는 노력조차 하지 않아 '0점'을 받는 경우가 있는데, 이는 아무것도 배우지 못했음을 의미합니다. 만약 코치가 단순히 문제의 개수만 센다면, 코치는 제브라 퍼즐을 가르치고 있다고 생각할 수도 있지만, 실제로는 로봇이 수학 문제들로부터만 배우고 있는 상태, 즉 '제로 그래디언트(zero-gradients, 학습이 일어나지 않는 상태)'인 상황이 됩니다.

이 문제를 해결하기 위해 그들의 발명품의 두 번째 부분인 '비율 보존 샘플러(Ratio-Preserving Sampler)'가 등장합니다. 이것은 매우 엄격한 사서와 같아서, 설령 로봇이 쓸모없는 연습 시도들을 버리더라도, 최종적으로 남은 유용한 연습 문제들의 구성이 수학과 논리 퍼즐의 정확한 비율을 유지하도록 보장합니다. 만약 코치가 연습의 30%를 논리 퍼즐로 채우길 원한다면, 사서는 쓸모없는 것들을 걸러낸 후에도 남은 유용한 문제들 중 30%가 여전히 논리 퍼즐이 되도록 확실히 합니다. 이는 쉬운 작업들이 실패한 시도가 많다는 이유로 의도치 않게 훈련 세션을 장악하는 것을 방지합니다.

이 새로운 방법의 결과는 매우 유망합니다. 연구팀이 세 가지 다른 추론 작업의 혼합물로 테스트했을 때, 이 새로운 MT-GRPO 방식은 표준 방식에 비해 최하위 성능을 보이는 작업의 정확도를 16~28% 향상시켰으며, DAPO라고 불리는 또 다른 고급 방식보다도 6% 더 높은 성능을 보였습니다. 더욱 인상적인 것은, 경쟁 모델보다 50% 적은 훈련 단계를 사용하여 가장 어려운 작업에서 **50%**의 정확도 임계값에 도달했다는 점입니다. 9가지의 서로 다른 작업을 대상으로 한 더 큰 규모의 테스트에서도 이 방법은 다른 방식들을 계속해서 능가하며, 더 복잡하고 다양한 도전 과제들을 처리하기 위해 확장될 수 있음을 보여주었습니다. 저자들은 '노브(knob)'( $\lambda$ 라고 불리는 파라미터)를 조절함으로써, 모델이 단순히 평균적으로 좋아지는 것보다 자신의 가장 약한 기술을 고치는 데 얼마나 우선순위를 둘지 제어할 수 있다는 것을 발견했습니다.

요약하자면, 이 논문은 로봇에게 어떤 문제를 연습하게 할지 더 똑똑하게 결정하고, 그 문제들이 실제로 훈련 세션에 도록 보장함으로써, 우리가 AI를 단순히 쉬운 것들에만 능숙한 것이 아니라 전반적으로 신뢰할 수 있게 유능하게 만들 수 있다는 것을 시사합니다. 이는 AI가 단순히 하나의 좁은 분야에 특화되는 것이 아니라, 어떤 기술도 뒤처지게 하지 않고 매우 다양한 실세계 문제들을 통해 진정으로 추론할 수 있는 존재가 되는 것을 향한 한 걸음입니다.