Towards Engineering Scaling Laws with Pretraining Data Composition

이 논문은 고충실도 시뮬레이터가 저렴한 합성 데이터 생성을 가능하게 하는 입자 물리학 분야에서, 사전 학습 데이터셋을 전략적으로 큐레이션함으로써 하드론 제트 분류 모델의 스케일링 동작이 모델 크기보다 데이터의 다양성과 정렬을 우선시하도록 설계될 수 있음을 입증한다.

핵심

당신이 학생에게 복잡한 도시에서 다양한 종류의 차량을 인식하는 법을 가르치고 있다고 상상해 보십시오. 당신은 이 학습을 돕기 위해 두 가지 주요 방법을 사용할 수 있습니다. 학생에게 더 큰 뇌(더 큰 모델)를 주거나, 혹은 더 많은 연습 문제(더 많은 데이터)를 주는 것입니다.

오랫동안 인공지능(AI)을 연구하는 과학자들은 하나의 "황금률"이 존재한다고 믿어 왔습니다. 만약 정해진 시간과 비용(컴퓨팅 예산)이 있다면, 가장 똑똑한 학생을 만들기 위해 자원을 더 큰 뇌를 만드는 데와 더 많은 연습 문제를 주는 데 각각 대략 50/50으로 나누어 배분하는 것이 최선이라고 생각했습니다.

하지만 이 새로운 논문은 입자 물리학의 세계에서, 우리가 학생이 무엇을 먼저 배울지를 바꿈으로써 더 나은 규칙을 **설계(engineer)**할 수 있다고 제안합니다.

설정: 물리학 교실

연구진은 "제트(jets)"를 대상으로 연구하고 있습니다. 입자 물리학에서 아주 작은 입자들이 서로 충돌할 때, 이들은 다른 입자들의 흐름인 '제트'를 뿜어냅니다. 이는 마치 불꽃놀이가 터질 때 불꽃 대신 아원자 입자의 흐름이 나오는 것과 같습니다.

목표는 AI에게 이 흐름을 보여주고, "아, 이것은 특정한 종류의 폭발로부터 온 것이구나!"라고 말하도록 가르치는 것입니다.

실험: 교과서 바꾸기

연구진은 학습 규칙을 어떻게 변화시키는지 확인하기 위해 두 가지 서로 다른 "교과서"(사전 학습 데이터셋)를 테스트했습니다.

1. 지루한 교과서 (QCD 전용): 이 책에는 "표준적인" 입자 폭발 사례들만 들어 있습니다. 이는 마치 일반적인 세단 운전법만 가르치는 운전 학원과 같습니다.
2. 다양한 교과서 (BSM 강화): 이 책에는 표준적인 사례뿐만 아니라, 우리 일반 우주에서는 일어나지 않는 복잡하고 희귀하며 이색적인 폭발(시뮬레이션된 "표준 모형 너머(Beyond Standard Model)"의 물리학)까지 포함되어 있습니다. 이는 세단 운전법을 가르치면서도 레이싱 카, 트럭, 심지어 비행 차량까지 가르치는 운전 학원과 같습니다.

발견: 규칙의 재작성

연구진이 지루한 교과서를 사용하여 AI를 훈련시켰을 때, 기존의 50/50 규칙이 유효했습니다. 더 나은 결과를 얻으려면 뇌를 크게 만드는 것과 더 많은 연습 문제를 주는 것 사이의 균형을 맞춰야 했습니다.

하지만 다양한 교과서를 사용하여 AI를 훈련시켰을 때, 규칙이 완전히 바뀌었습니다. AI는 더 많은 연습 문제를 푸는 것이 훨씬 더 가치 있다는 것을 배웠습니다.

• 비유: 다양한 교과서로 훈련받은 AI는 이미 세상의 모든 종류의 차량을 본 적이 있는 학생과 같습니다. 그들에게 새로운 테스트를 줄 때, 그들은 새로운 차를 이해하기 위해 더 큰 뇌를 필요로 하는 것이 아니라, 단지 그 차에 대한 더 많은 사례를 보는 것만으로도 완벽해질 수 있습니다. 그들의 "경험"이 매우 풍부하기 때문에 "뇌"가 그렇게 빨리 성장할 필요는 없습니다.

결과: 새로운 "데이터 우선" 전략

논문은 다양한(이색적인) 데이터를 사용하여 초기 훈련을 했을 때 다음과 같은 결과를 발견했습니다:

• "더 큰 뇌" 전략의 중요성이 줄어들었습니다.
• "더 많은 데이터" 전략이 승자가 되었습니다.

실제로 연구진은 컴퓨팅 파워를 사용하는 매 단위마다, 약 78%를 더 많은 데이터를 얻는 데 쓰고, 단 22%만을 모델을 더 크게 만드는 데 써야 한다는 것을 발견했습니다. 이는 기존의 50/50 분할 방식에서 일어난 거대한 변화입니다.

이것이 물리학에서 중요한 이유:
이 논문은 물리학이 가진 독특한 강점을 강조합니다. 우리는 우리만의 데이터를 직접 만들 수 있습니다.

의학이나 언어 분야에서는 새로운 데이터를 얻는 것이 어렵고, 비용이 많이 들거나 불가능할 수 있습니다(새로운 환자를 그냥 "시뮬레이션"할 수는 없습니다). 하지만 입자 물리학에서 과학자들은 강력한 컴퓨터를 사용하여 입자 충돌을 시뮬레이션합니다. 일단 시뮬레이션이 실행되면, 이들은 무한한 양의 고품질, 고다양성 데이터를 무료로 생성할 수 있습니다.

핵심 요약:
물리학을 위한 초지능 AI를 구축하고 싶다면, 단순히 가능한 가장 큰 뇌를 만들려고 애쓰지 마십시오. 대신, AI가 먼저 배울 수 있도록 더 다양하고 풍부한 커리큘럼을 설계하는 데 시간과 비용을 투자하십시오. 일단 AI가 다양한 "이색적인" 사례들을 접하고 나면, 당신이 부여하는 특정 과업으로부터 더 빠르고 더 잘 배우게 될 것이며, 모델을 더 크게 만드는 것보다 더 많은 데이터를 제공함으로써 더 나은 결과를 얻을 수 있을 것입니다.

요약하자면: 잘 선택된 다양한 구성의 학습 데이터는 더 큰 뇌보다 더 강력합니다.

기술 요약

기술 요약: 사전 학습 데이터 구성을 통한 엔지니어링 스케일링 법칙 연구

문제 정의
뉴럴 스케일링 법칙(Neural scaling laws)은 모델의 성능이 연산량(compute), 모델 크기(model size), 데이터셋 크기(dataset size)에 따라 거듭제곱 법칙(power law)을 따르며 어떻게 개선되는지를 설명한다. 대규모 언어 모델(LLM)에서는 잘 확립되어 있으나, 입자 물리학에서도 이러한 관계가 새롭게 나타나고 있다. 기초 물리학에서의 핵심적인 차이점은, 더 큰 모델을 훈련하는 데 드는 계산 비용에 비해 상대적으로 저렴한 비용으로 시뮬레이터를 통해 고충실도(high-fidelity) 합성 데이터를 생성할 수 있다는 점이다. 이는 사전 학습 데이터셋 자체를 설계하여 스케일링 동작에 영향을 미칠 수 있는 독특한 기회를 제공한다. 본 연구가 다루는 핵심 질문은 사전 학습 데이터의 구성, 즉 데이터의 다양성과 다운스트림 태스크와의 정렬(alignment)을 설계함으로써, 연산 최적화된 스량(compute-optimal)의 중점을 더 큰 모델에서 더 큰 데이터셋으로 이동시킬 수 있는지 여부이다.

방법론
본 연구는 고에너지 입자 충돌에서 발생하는 강입자 제트(hadronic jets)를 분류하는 태스크에 초점을 맞춘다. 저자들은 제트 데이터를 포인트 클라우드(point cloud)로 처리하는 일반적인 트랜스포머 아키텍처를 활용하며, 깊이(depth)와 어텐션 헤드 차원(attention head dimensions)은 고정시킨 채 모델 크기를 약 3,000개에서 1,050만 개의 파라미터까지 변화시키며 실험을 진행하였다(세 자릿수 규모의 변화).

실험 설계는 두 단계의 훈련 프로토콜을 포함한다:

1. 사전 학습(Pretraining): 모델들은 188개의 클래스를 가진 JetClass-II 데이터셋의 하위 집합들로 사전 학습된다. 저자들은 다양성과 정렬을 조작하기 위해 네 가지 별도의 사전 학습 하위 집합을 정의하였다:
   • QCD: 가벼운 쿼크 또는 글루온에 의해 생성된 제트만 포함 (17개 클래스).
   • QCD + res2p: QCD 제트에 두 물체 붕괴(two-body decays)를 통한 BSM(Beyond the Standard Model) 공명(resonance) 제트를 추가.
   • QCD + res34p: QCD 제트에 세 물체 또는 네 물체 붕괴를 통한 BSM 공명 제트를 추가.
   • QCD + res2p + res34p: 모든 BSM 공명 붕괴를 포함한 전체 데이터셋.
   • 참고: BSM 하위 집합은 QCD 전용 데이터에 비해 더 높은 다양성(더 많은 프로세스 클래스, 더 넓은 위상 공간 커버리지)과 다운스트림 태스크에 대한 더 나은 정렬(멀티 프롱(multi-prong) 토폴로지)을 제공한다.
2. 미세 조정(Fine-tuning): 사전 학습된 모델들은 원래의 JetClass 데이터셋을 사용하여 10개 클래스 제트 분류 태스크(가벼운 쿼크/글루온, 톱 쿼크, W/Z 보존, 힉스 입자 식별)에 대해 미세 조정된다. 이 태스크는 프롱 다중도(prong multiplicity)와 질량 척도를 식별해야 하며, 이는 BSM이 증강된 사전 학습 데이터에는 잘 나타나 있지만 QCD 전용 데이터에는 거의 나타나지 않는 특성이다.

스케일링 지수(scaling exponents)는 총 연산량($C$)에 따른 최적 모델 크기($N^*$)와 데이터셋 크기($D^*$)에 거듭제곱 법칙을 피팅하여 추출한다. 연구는 "스크래치(scratch)" 훈련(사전 학습 없음)과 다양한 사전 학습 구성 간의 지수를 비교한다.

주요 결과
본 연구는 사전 학습 데이터 구성이 연산 최적화 스케일링 지수를 유의미하게 변화시킨다는 것을 입증한다:

• 스크래치 훈련: 처음부터 훈련할 경우 지수는 $a \approx 0.52$ (모델 크기) 및 $b \approx 0.48$ (데이터셋 크기)로 나타나며, 이는 LLM의 결과와 일치하게 모델 크기와 데이터 사이에 연산 자원을 대략 균등하게 배분함을 의미한다.
• QCD 전용 사전 학습: QCD 제트로만 사전 학습할 경우 지수의 변화는 미미하며 ($a \approx 0.53, b \approx 0.47$), 이는 특정 정렬이나 다양성 없이 사전 학습만으로는 스케일링 체계를 근본적으로 바꾸지 못함을 시사한다.
• BSM 증강 사전 학습: 사전 학습 코퍼스에 BSM 공명 붕괴를 포함하면 극적인 변화가 일어난다. 전체 BSM 증강 데이터셋을 사용할 경우, 지수는 $a \approx 0.22$ 및 $b \approx 0.78$로 이동한다.
  • 이는 연산 최적화 전략이 모델 크기를 키우는 것보다 데이터 크기를 키우는 쪽을 강력하게 선호하는 체계로 전환됨을 나타낸다.
  • 이 변화는 스크래치 베이스라인과 비교했을 때 모델 크기에 대한 스케일링 지수가 2.3배 감소했음을 의미한다.
  • 미세 조정 손실 곡선(loss curves)은 BSM이 강화된 사전 학습이 모든 모델 크기에서 일관되게 손실을 낮추며, 모델이 커질수록 그 이점이 증가함을 확인해 준다.

주요 기여

1. 스케일링 법칙의 엔지니어링: 본 논문은 사전 학습 데이터 구성이 기초 물리학에서 스케일링 지수를 변화시키도록 엔지니어링될 수 있음을 보여주는 최초의 체계적인 연구를 제공한다. 이는 사전 학습 코퍼스의 다양성과 다운스트림 정렬이 스케일링 체계를 데이터 선호 전략으로 이동시킬 수 있음을 입증한다.
2. 정량적 변화: 본 연구는 BSM 물리학을 사전 학습에 통합함으로써 균형 잡힌 스케일링 체계($a \approx b \approx 0.5$)에서 강력한 데이터 선호 체계($a \approx 0.22, b \approx 0.78$)로의 변화를 정량화한다.
3. 파운데이션 모델에 대한 시사점: 결과는 다양하고 정렬된 합성 데이터로 사전 학습된 파운데이션 모델이 더 작은 파라미터 수로도 최적의 성능을 달({achieve) 수 있음을 시사하며, 이를 통해 절약된 연산 예산을 추가적인 데이터 생성으로 재배정할 수 있게 한다.

의의 및 주장
저자들은 본 연구가 과학적 머신러닝의 새로운 설계 공간인 '파운데이션 모델 훈련을 위한 물리적 입력'을 식별했다고 주장한다. 데이터 큐레이션이 가용성에 의해 제한되는 자연어 또는 이미지 도메인과 달리, 기초 물리학은 저렴하고 충실도가 높은 시뮬레이터를 활용하여 스케일링 법칙을 능동적으로 형성하는 사전 학습 코퍼스를 구축할 수 있다.

논문은 BSM 구성된 코퍼스를 통한 사전 학습이 다운스트림 연산을 더 큰 모델보다는 더 많은 데이터에 쓰는 것이 최적인 체계를 허용하지만, 이러한 결과가 다른 미세 조정 태스크, 더 큰 모델 규모, 그리고 더 큰 데이터셋에 걸쳐 일반화될 수 있는지 검증하기 위해서는 추가적인 연구가 필요하다고 조심스럽게 결론짓는다. 본 연구는 스케일링의 모든 과제를 해결했다고 주장하는 것이 아니라, 과학적 파운데이션 모델의 발견 잠재력을 극대화하기 위한 미개척 레버로서 '사전 학습 구성 엔지니어링'을 강조하는 데 목적이 있다.