데이터와 모델을 동시에 최적화하는 새로운 접근법 JoBS

초록

본 논문은 대규모 언어 모델(LLM) 학습 시 데이터 혼합 비율과 파라미터 효율 파인튜닝(PEFT) 설정을 동시에 최적화하는 문제를 “닭과 달걀 딜레마”로 정의하고, 소규모 학습 단계에서 얻은 성능을 예측하는 스케일링 법칙 기반 예측기를 활용해 베이지안 최적화(BO)를 가속화하는 JoBS 알고리즘을 제안한다. 예산을 예측기 학습과 BO에 적절히 배분함으로써 전체 학습 비용을 크게 절감하면서도 기존 데이터‑전용, 모델‑전용, 다중‑신뢰도 BO 방법들을 능가하는 성능을 입증한다.

상세 분석

JoBS는 두 가지 핵심 아이디어를 결합한다. 첫째, LLM의 최종 성능 L(x) (여기서 x는 데이터 혼합 비율 X와 PEFT 구성 M을 포함한 연속형 벡터) 를 부드러운 함수로 가정하고, 가우시안 프로세스(GP)를 사용해 베이지안 최적화의 서브시스트럼을 구축한다. 이는 기존 연구에서 관찰된 “성능 지형이 매끄럽다”는 경험적 사실에 기반한다. 둘째, 전체 학습 단계 B(예: 1000 스텝) 대신 작은 단계 B_small(예: 100 스텝)에서 얻은 손실/정확도 정보를 입력으로 하는 신경망 기반 성능 예측기를 학습한다. 이 예측기는 스케일링 법칙을 일반화한 형태로, 다양한 데이터·모델 조합에 대해 학습된 파라미터를 공유함으로써 높은 표현력을 갖는다. 중요한 점은 예측기가 완벽할 필요가 없으며, 예측 오차는 GP의 관측 노이즈로 처리되어 BO가 수렴하도록 설계되었다는 것이다.

예산 C는 전체 학습 스텝 수이며, JoBS는 C를 두 부분으로 나눈다. 첫 번째 부분은 다양한 (X, M) 조합에 대해 B 스텝 전체 학습을 수행해 라벨을 얻고, 이를 통해 예측기를 학습한다. 두 번째 부분은 학습된 예측기를 사용해 B_small 스텝만 수행하고, 예측값을 GP에 입력해 BO를 진행한다. 논문은 평균 후회(R_T)를 분석하여, 예측기 학습에 할당하는 스텝 수와 BO 반복 횟수 사이의 최적 균형을 수식적으로 도출한다(정리 5.1). 이론적 결과는 예측기 오차 σ가 작을수록 후회가 선형적으로 감소함을 보여주며, 실제 실험에서도 예측기 정확도가 0.8 R² 이상일 때 기존 다중‑신뢰도 BO보다 2~3배 빠른 수렴을 달성한다.

또한, 실험에서는 Llama‑3‑8B‑Instruct를 GSM‑8K, MMLU, TruthfulQA 등 5개의 다운스트림 태스크에 적용하였다. 데이터 측면에서는 8개의 도메인(위키피디아, 뉴스, 코드 등)에서 혼합 비율을 최적화했고, 모델 측면에서는 LoRA 레이어, 랭크, α, dropout 등을 탐색했다. JoBS는 동일 예산(C = 50 k 스텝) 하에서 기존 데이터‑전용 최적화(예: 데이터 혼합 비율만 탐색)와 모델‑전용 최적화(예: LoRA 파라미터만 탐색)보다 평균 7.3%~12.5% 높은 정확도를 기록했다. 다중‑신뢰도 BO(Freeze‑Thaw, Early‑Stopping)와 비교했을 때도 후회 감소율이 30% 이상 우수했다.

이 논문의 주요 공헌은 (1) 데이터와 모델을 동시에 고려한 공동 최적화 프레임워크를 제시, (2) 스케일링 법칙을 신경망 예측기로 일반화해 저비용 관측을 가능하게 함, (3) 예산 할당에 대한 이론적 최적화 분석을 제공, (4) 다양한 실제 LLM 태스크에서 실증적으로 우수성을 입증한 점이다. 한계로는 예측기 학습을 위해 초기 전수 학습이 필요하고, 매우 큰 모델(수십억 파라미터)에서는 B_small 단계에서도 GPU 메모리 부담이 있을 수 있다. 향후 연구에서는 메타‑러닝 기반 예측기 초기화와 분산 학습을 결합해 초기 비용을 더욱 감소시키고, 다른 파인튜닝 기법(예: 프롬프트 튜닝)까지 확장하는 방향이 기대된다.