MSFT: 멀티태스크 SFT에서 데이터셋 혼합물의 이질적 과적합 문제 해결

기존 멀티태스크 SFT는 모든 데이터셋에 동일한 학습량을 할당하여 일부는 과적합되고 일부는 덜 학습되는 불균형 문제가 있었다. MSFT는 각 데이터셋의 학습 속도에 맞춰 동적으로 학습을 중단하고 최적 시점으로 되돌리는 방식을 통해 모델 성능을 극대화하고 연산 비용을 절감한다.

대형 언어 모델의 SFT 단계에서는 수학, 대화, 상식 등 성격이 다른 여러 데이터셋을 섞어 학습한다. 이때 모든 데이터를 동일한 횟수만큼 반복 학습시키면, 배우기 쉬운 데이터는 일찍 과적합되어 성능이 떨어지기 시작하고 어려운 데이터는 아직 충분히 배우지 못한 상태가 된다. 이는 마치 마라톤에서 모든 선수가 가장 느린 선수의 속도에 맞춰 뛰거나, 가장 빠른 선수가 지칠 때까지 모두가 계속 뛰어야 하는 비효율적인 상황과 같다.

MSFT는 이 문제를 해결하기 위해 '가장 먼저 지치는 주자'를 찾아내어 코스에서 하차시키는 전략을 쓴다. 먼저 전체 데이터를 짧은 구간만큼 학습시킨 후 각 데이터셋의 성적을 확인한다. 만약 어떤 데이터셋의 성적이 정점을 찍고 내려가기 시작했다면, 그 데이터셋이 가장 좋은 성적을 냈던 과거 시점으로 모델의 상태를 되돌린다(Roll-back). 그 후 과적합된 데이터셋만 빼고 남은 데이터들로 다시 학습을 이어간다.

단순히 데이터를 빼는 것보다 과거로 되돌리는 것이 중요한 이유는 데이터 구성이 바뀌면 모델이 나아가야 할 최적의 방향도 바뀌기 때문이다. MSFT는 이 과정을 반복하여 각 데이터셋이 자신의 '골든 타임'에서 학습을 멈출 수 있게 유도한다. 결과적으로 모델은 모든 작업에서 고르게 높은 성능을 내며, 불필요한 사후 학습을 줄여 전체 연산량까지 아끼게 된다.

MSFT는 반복적인 Roll-out과 Roll-back 과정을 통해 최적의 데이터 혼합 학습 경로를 탐색한다. 전체 데이터셋 D에서 제외 세트 E를 관리하며, 활성 데이터셋 D\E에 대해 설정된 연산 예산 C만큼 모델을 학습시킨다. 각 단계에서 모든 서브 데이터셋의 정확도를 기록하여 최적의 연산 지점 c를 찾는다. [각 데이터셋의 학습 단계별 정확도 입력] → [최대 정확도 시점 추출 연산] → [각 작업의 개별 최적 에포크 결과 도출]. 이 중 가장 작은 c를 가진 데이터셋을 제외 대상으로 선정한다.

가장 빠른 과적합 지점 c_min이 현재 예산 C보다 작을 경우, 모델을 해당 c_min 시점의 체크포인트로 롤백한다. [현재 모델 파라미터와 저장된 체크포인트 입력] → [c_min 시점의 가중치 로드 연산] → [최적 상태로 복구된 모델 결과 획득]. 이후 해당 데이터셋을 제외 세트 E에 추가하고 다음 반복을 수행한다. 이 과정은 모든 데이터셋이 제외되거나 학습이 완료될 때까지 지속된다.

효율적인 자원 관리를 위해 동적 체크포인트 가지치기(Dynamic Checkpoint Pruning)를 도입했다. 각 단계에서 식별된 데이터셋별 피크 체크포인트만 유지하고 나머지는 삭제함으로써 저장 공간 오버헤드를 줄인다. [전체 생성된 체크포인트 입력] → [데이터셋별 최적 지점 외 삭제 연산] → [SFT 대비 약 4.44배 수준의 저장 공간 유지 결과]. 이를 통해 연산 대비 미미한 비용으로 알고리즘을 실행한다.

6개의 베이스 모델(OLMo 2 1B, Qwen2.5 0.5B~7B, Qwen3 8B)을 대상으로 10개 벤치마크에서 실험한 결과, MSFT는 모든 모델에서 기존 SFT 및 최신 베이스라인(DynamixSFT, IES 등)을 능가하는 평균 정확도를 기록했다. Qwen3 8B 모델의 경우 표준 SFT 대비 평균 1.8%p 이상의 성능 향상을 보였다. 도메인별로는 과학 및 지식(+0.7%), 상식 및 언어(+2.4%), 수학 및 정량적 추론(+3.0%) 등 모든 주요 영역에서 일관된 개선을 입증했다.

효율성 분석에서 연산 예산 C=1로 설정했을 때, 표준 SFT 대비 성능은 3.4% 향상시키면서도 전체 연산량(FLOPs)은 평균 120.3 PFLOPs 절감했다. 또한 벤치마크 간 성능 편차(Standard Deviation)가 가장 낮게 나타나, 특정 작업에 치우치지 않고 전반적인 능력을 고르게 개선함을 확인했다. 이는 MSFT가 성능의 하한선과 상한선을 동시에 높이는 안정적인 알고리즘임을 시사한다.

MSFT는 데이터 스케줄링 레벨에서 작동하며, 모델 내부의 그래디언트 조작이나 복잡한 손실 함수 재설계 없이 하드 제외(Hard Exclusion)와 롤백을 사용한다. 이는 기존의 그래디언트 기반 방식들이 가지는 계산 오버헤드와 민감한 하이퍼파라미터 튜닝 문제를 회피하게 해준다. 연구진은 데이터셋 제외 시 발생하는 파라미터 발산 현상을 수학적으로 분석하여, 단순 제외가 아닌 롤백 후 재학습이 최적화 궤적을 유지하는 데 필수적임을 규명했다.

실험 결과 MSFT는 표준 SFT보다 낮은 트레이닝 로스(Training Loss)를 달성한다. 이는 과적합된 데이터셋이 제공하는 노이즈 섞인 그래디언트를 제거함으로써, 옵티마이저가 남은 미학습 작업들에 모델 용량을 더 효율적으로 할당할 수 있게 돕기 때문이다. 구현 측면에서는 C라는 단일 하이퍼파라미터만 추가되며, 체크포인트 관리 알고리즘을 통해 디스크 사용량을 모델 크기의 약 4.44배 이내로 억제하여 대규모 모델 학습 시의 실무적 제약 사항을 해결했다.

MSFT의 주요 한계점은 롤-아웃 단계에서 각 데이터셋의 피크 시점을 찾기 위해 중간 체크포인트를 저장해야 하므로 추가적인 디스크 저장 공간 오버헤드가 발생한다는 점이다. 동적 가지치기 알고리즘으로 이를 완화했으나 여전히 표준 SFT보다는 많은 저장 자원을 요구한다.