중국 알리바바의 통이실험실(通义实验室)이 대형 언어모델의 추론과 검색 능력을 크게 향상시키는 새로운 오픈소스 프레임워크 MaskSearch를 공개했다. 이 혁신적인 사전훈련 프레임워크는 소형 모델이 대형 모델에 버금가는 성능을 발휘할 수 있게 하며, 여러 개방형 도메인 질의응답 데이터셋에서 현저한 성능 향상을 보여주고 있다.
MaskSearch의 핵심 개념과 기술적 혁신
MaskSearch는 BERT 모델의 마스크(Mask) 예측 작업에서 영감을 받아 개발된 검색 증강형 마스크 예측 작업을 도입했다. 이는 모델이 검색 도구를 활용하여 텍스트에서 가려진 부분을 예측하도록 하는 방식으로, 사전훈련 과정에서 범용적인 작업 분해, 추론 등의 지능형 에이전트 전략을 효과적으로 학습하게 한다.
통이실험실 검색팀은 이전에 ZeroSearch, OmniSearch 등의 연구를 통해 특정 하위 작업에서 강화학습 훈련을 진행해왔다. 하지만 특정 작업에서만 훈련할 경우 모델의 일반화 능력이 제한되어 더 많은 시나리오의 검색 추론 작업에 적응하기 어렵다는 한계가 있었다. MaskSearch는 이러한 문제점을 해결하기 위해 설계되었다.
RAMP 사전훈련 작업의 구조와 특징
검색 증강형 마스크 예측(RAMP)이 MaskSearch의 핵심 사전훈련 작업으로, 입력된 텍스트 시퀀스에서 핵심 정보를 마스크 처리하고 모델이 외부 지식베이스를 활용해 검색 도구를 호출하여 가려진 텍스트 조각을 예측하도록 한다.
마스크된 부분의 난이도를 높이기 위해, 기존 마스크 예측 작업에서 흔히 가려지던 명명된 개체(인명, 지명, 조직명 등), 날짜, 숫자 외에도 다음과 같은 핵심 정보 유형들을 추가로 고려했다:
– 온톨로지 지식: 텍스트에 관련된 분류 체계나 지식 체계의 핵심 개념
– 전문 용어: 특정 영역이나 주제의 전문 용어
– 수치값: 통계 데이터, 측정값 등 구체적인 수치
이러한 접근 방식은 작업의 난이도를 증가시킬 뿐만 아니라 모델이 검색과 추론 과정에서 정보를 더욱 정교하게 처리하도록 유도하여, 다영역 작업에서의 적응력과 일반화 능력을 향상시킨다.
다양한 훈련 방법론의 통합
MaskSearch의 가장 큰 장점 중 하나는 지도학습 미세조정(SFT)과 강화학습(RL) 두 가지 훈련 방법과의 호환성이다.
지도학습 미세조정 접근법
지도학습을 위한 사고의 연쇄(CoT) 데이터 생성을 위해 에이전트 합성과 증류를 결합한 데이터 생성 방법을 제안했다. 먼저 계획, 검색 재작성, 관찰 분석 등의 역할을 포함한 다중 지능형 에이전트 시스템을 구축하여 협력적으로 사고의 연쇄를 생성한다. 최종적으로 LLM이 답변 판단을 담당하여 정확한 답변의 사고의 연쇄만을 보존한다.
강화학습 최적화 전략
강화학습 부분에서는 동적 샘플링 전략 최적화(DAPO) 알고리즘을 채택하고, 하이브리드 보상 시스템을 구축했다. 형식 보상은 모델 출력이 지정된 형식에 부합하는지 검사하고, 답변 보상은 생성된 답변과 표준 답변의 일치성을 평가한다.
커리큘럼 학습의 도입
쉬운 것부터 어려운 것까지 단계적으로 학습할 수 있도록, 마스크 수량에 따라 훈련 샘플의 난이도를 분류하는 방식을 제안했다. 모델이 먼저 간단한 샘플을 통해 기초 추론 기술을 학습한 후, 점진적으로 능력을 향상시켜 더욱 도전적인 시나리오에 대응할 수 있게 한다.
실험 결과와 성능 분석
주요 실험 성과
연구진은 다양한 크기의 Qwen과 LLaMA 모델을 기반으로 한 실험을 통해 2단계 MaskSearch 훈련 프레임워크가 대형 모델의 검색과 추론 능력을 현저히 향상시킨다는 것을 증명했다. RAMP를 사전훈련 작업으로, HotpotQA 데이터셋을 하위 작업으로 하는 훈련 과정을 따라, MaskSearch는 도메인 내 데이터셋에서 모델 재현율을 안정적으로 향상시켰다.
특히 주목할 만한 점은 Bamboogle 등 도메인 외 데이터셋에서 성능 향상이 더욱 현저했으며, 소형 모델도 대형 모델에 버금가는 성능을 보여줘 RAMP가 확장 가능한 학습 신호로서의 효과를 입증했다는 것이다.
훈련 방식별 성능 비교
실험 결과 지도학습(SFT)과 강화학습(RL) 두 훈련 방식 모두 MaskSearch 프레임워크와의 호환성을 확인했다. 그중 RL은 RAMP 작업에서 더 높은 성능 상한선을 보여줬으며, 특히 HotpotQA 등 도메인 내 작업에서 모든 크기의 Qwen 모델에서 최적의 효과를 달성했다.
이는 RL이 동적 샘플링 전략과 하이브리드 보상 메커니즘을 통해 모델의 다단계 검색과 추론 과정을 더욱 정확하게 최적화할 수 있음을 시사하며, 검색 증강 모델의 적응성 향상을 위한 더 강력한 훈련 패러다임을 제공한다.
확장성 성능 검증
지도학습 시나리오에서 다양한 훈련 단계 실험을 통해 MaskSearch의 확장성을 검증한 결과, 소형 모델(1B)은 사전훈련 후 성능 향상이 현저했으며, 대형 모델(7B)은 자기진화 데이터의 다양성 제한으로 성능 증가가 상대적으로 완만했지만, 재현율 점수는 여전히 미세조정만 한 모델 대비 성장을 보였다.
기술적 세부사항과 분석
마스크 전략의 영향
마스크 전략은 RAMP 사전훈련 작업의 난이도에 영향을 미치는 또 다른 중요한 요소다. 연구진은 랜덤 마스크와 복잡도(PPL) 기반의 난이도 지향 마스크 전략을 비교했다. PPL 전략은 모델이 마스크를 복원할 때의 손실값(즉, 복잡도)을 계산하여 복원 난이도가 높은 부분을 우선적으로 선택해 가린다.
실험 결과 PPL 전략은 FanoutQA 데이터셋에서 모델 재현율을 향상시켰지만, 다른 데이터셋에서는 과도한 난이도 추구로 인해 성능이 하락하기도 했다. 이는 작업 난이도가 여전히 모델의 현재 검색과 추론 능력과 일치해야 함을 보여준다.
강화학습 보상 함수의 영향
강화학습 훈련 과정에서 다양한 보상 함수가 모델 성능에 미치는 영향이 각각 다르다는 것을 확인했다. Qwen2.5-7b 모델을 예로 들면, 토큰 수준 재현율 기반 보상 함수는 모델이 재현율 향상을 위해 답변에 대량의 무관한 정보를 축적하도록 유도하여 답변 길이가 대폭 증가하고, 다른 RL 보상 함수 대비 실제 성능이 현저히 하락했다.
상대적으로 모델 기반 보상 함수는 최적의 성능을 보여줬으며, 모델이 생성한 답변 길이, 토큰 수준 재현율, Qwen72b 모델 평가 점수에서 모두 다른 두 보상 방법보다 우수했다.
연구의 의의와 향후 전망
MaskSearch는 대형 언어모델(LLM)의 지능형 에이전트 추론과 검색 능력 향상을 목표로 한다. 이 프레임워크는 검색 증강형 마스크 예측(RAMP) 사전훈련 작업을 기반으로 모델이 자주적으로 다단계 검색과 추론을 실행하여 텍스트의 마스크된 공백을 채우고 외부 지식의 심층적 통합을 실현한다.
지도학습 미세조정(SFT)과 강화학습(RL) 이중 훈련 경로를 통한 훈련과 커리큘럼 학습 전략 도입으로, MaskSearch는 도메인 내 및 크로스 도메인 개방형 도메인 질의응답 작업에서 모두 기존 방법 대비 현저한 성능 향상을 달성했다.
이번 연구는 AI 모델의 검색과 추론 능력을 통합적으로 향상시키는 새로운 패러다임을 제시했으며, 특히 소형 모델도 대형 모델급 성능을 구현할 수 있다는 점에서 AI 모델의 효율성과 접근성 향상에 중요한 기여를 할 것으로 예상된다. 알리바바가 이 기술을 오픈소스로 공개함으로써 전체 AI 커뮤니티의 발전에도 긍정적인 영향을 미칠 것으로 보인다.
[참고 기사]
https://www.qbitai.com/2025/05/289933.html
