Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training

이 논문은 지속적 사전 학습 (Continual Pre-training) 과정에서 발생하는 재앙적 망각 (Catastrophic Forgetting) 문제를 해결하기 위해, 데이터 믹싱 에이전트 (Data Mixing Agent) 라는 새로운 모델을 제안합니다. 이 에이전트는 도메인별 데이터 가중치를 학습을 통해 자동으로 조정하여, 새로운 타겟 도메인 (예: 수학 추론, 코드 생성) 에 대한 성능을 향상시키면서도 원본 도메인의 능력을 유지하는 균형을 찾습니다.

주요 내용은 다음과 같습니다.

1. 문제 정의 (Problem)

• 배경: 대규모 언어 모델 (LLM) 은 일반적인 사전 학습을 거치지만, 특정 지식 집약적 분야 (수학, 코드 등) 에 적용하기 위해서는 해당 분야의 데이터로 지속적 사전 학습이 필요합니다.
• 문제점: 타겟 도메인 데이터만 직접적으로 학습하면, 모델이 원본 도메인의 능력을 잃는 재앙적 망각이 발생합니다.
• 기존 접근법의 한계: 기존에는 소스 (원본) 데이터와 타겟 데이터를 혼합하여 학습하는 전략을 사용하지만, 이 혼합 비율 (가중치) 을 결정하는 것은 주로 인간의 직관이나 경험적 휴리스틱에 의존했습니다. 이는 최적의 데이터 배분을 보장하지 못하며, 새로운 시나리오에 적용하기 어렵습니다.

2. 방법론 (Methodology)

저자들은 데이터 믹싱 에이전트를 제안하며, 이를 강화 학습 (Reinforcement Learning, RL) 을 통해 도메인 가중치를 학습하는 모델 기반 프레임워크로 정의합니다.

• MDP (Markov Decision Process) 모델링:
  • 상태 (State): 이전 도메인 재가중치 단계들의 데이터 분포 이력.
  • 행동 (Action): 현재 단계에서의 도메인별 데이터 비율 (확률 분포).
  • 보상 (Reward): 경량화된 평가 환경 (Evaluation Environment) 에서 얻은 모델 성능 지표 (일반 능력 및 타겟 도메인 능력).
• 학습 파이프라인:
  1. 궤적 샘플링 (Trajectory Sampling): 다양한 데이터 믹싱 궤적을 무작위 및 유도적 인덕션 (Inductive Biases) 을 통해 생성합니다.
  2. 프록시 모델 학습 및 평가: 각 궤적에 대해 작은 프록시 모델을 학습시키고, 경량 평가 환경에서 성능을 측정하여 피드백을 수집합니다.
  3. 에이전트 최적화 (CQL): 수집된 궤적과 피드백을 바탕으로 보수적 Q-학습 (Conservative Q-Learning, CQL) 알고리즘을 사용하여 에이전트를 학습시킵니다. CQL 은 분포 외 (Out-of-Distribution) 행동에 대한 Q 값 과대평가를 방지하여 안정적인 학습을 가능하게 합니다.
• 에이전트 구조: Transformer 디코더 아키텍처를 사용하며, 매우 경량화되어 (약 210 만 파라미터) 추론 시 지연 시간을 최소화합니다.

3. 주요 기여 (Key Contributions)

1. 최초의 모델 기반 도메인 재가중치 프레임워크: 휴리스틱을 수동으로 설계하는 대신, 에이전트가 데이터 믹싱 궤적과 피드백을 통해 일반화 가능한 휴리스틱을 학습하도록 한 최초의 엔드 - 투 - 엔드 방법론입니다.
2. 강화 학습을 통한 휴리스틱 학습: 단순한 모방 학습 (SFT) 이 아닌 CQL 기반의 강화 학습을 통해, 성능을 향상시키는 행동과 저하시키는 행동을 모두 학습하여 최적의 믹싱 전략을 도출합니다.
3. 높은 일반화 능력: 에이전트는 학습된 도메인 (수학) 에서 훈련되었음에도 불구하고, 재학습 없이 다른 소스 데이터, 다른 타겟 모델, 다른 도메인 공간 (52 차원), 그리고 완전히 새로운 타겟 도메인 (코드 생성) 에서도 우수한 성능을 발휘합니다.

4. 실험 결과 (Results)

• 수학 추론 (Math Reasoning) 평가:
  • 8 개의 일반 벤치마크와 4 개의 수학 벤치마크에서 기존 최첨단 방법 (RegMix, DBL 등) 보다 평균적으로 3.02% 이상의 성능 향상을 보였습니다.
  • 특히, DataAgentRL은 일반 능력과 수학 능력 모두에서 균형을 이루며, Naive Training(타겟 데이터만 학습) 에 비해 일반 능력 저하를 크게 완화했습니다.
• 코드 생성 (Code Generation) 평가:
  • 수학 도메인에서 학습된 에이전트를 코드 생성 도메인에 직접 적용했을 때, 해당 도메인에 최적화된 베이스라인보다도 우수한 성능을 보여주며 도메인 간 일반화 가능성을 입증했습니다.
• 데이터 효율성:
  • 에이전트는 소스 도메인 데이터에 대한 의존도를 줄이면서도 더 높은 성능을 달성하여, 전체 학습 토큰 예산을 최대 21.4 억 토큰까지 절감할 수 있었습니다.
  • 학습 시간 (GPU 시간) 측면에서도 베이스라인 방법들보다 효율적이었습니다.

5. 의의 및 결론 (Significance)

• 자동화된 데이터 큐레이션: 인간의 직관에 의존하던 데이터 믹싱 전략을 데이터와 모델에 독립적인 (Model-and-Data-Agnostic) 에이전트가 자동으로 학습하고 최적화함으로써, 지속적 사전 학습의 효율성과 효과성을 극대화했습니다.
• 재사용 가능성: 한 번 학습된 에이전트는 다양한 시나리오에 재학습 없이 적용 가능하므로, 새로운 도메인 적응 시 발생하는 막대한 계산 비용을 절감할 수 있습니다.
• 휴리스틱의 발견: 에이전트가 학습한 전략은 인간의 직관 (예: 과학/기술 도메인이 일반 성능에 긍정적임) 과 일치하면서도, 인간이 발견하지 못했던 미세한 조정 전략 (예: 특정 도메인에서의 데이터 감소) 을 발견하여 모델 성능을 향상시켰습니다.

결론적으로, 이 논문은 데이터 믹싱 에이전트를 통해 지속적 사전 학습의 핵심 난제인 '재앙적 망각'과 '데이터 혼합 최적화' 문제를 해결하는 강력한 새로운 패러다임을 제시합니다.