Predict before you train: Scaling Laws for particle physics foundation models

완벽한 케이크를 굽고 싶지만, 밀가루나 설탕 또는 열이 얼마나 필요한지 모른다고 상상해 보십시오. 입자 물리학의 세계에서 과학자들은 입자 충돌에서 발생하는 혼란스러운 파편들을 분석하기 위해 끊임없이 "디지털 뇌"(머신러닝 모델)를 구축하고 있으며, 이를 통해 새로운 자연 법칙을 밝혀낼 숨겨진 패턴을 찾고자 노력하고 있습니다. 이 디지털 뇌들은 거대한 오븐과 같습니다. 오븐이 더 크고(모델) 더 많은 재료를 투입할수록(데이터), 보통 케이크의 맛은 더 좋아집니다. 하지만 이 케이크를 굽는 일은 믿기 힘들 정도로 비쌉니다. 엄청난 양의 전기와 슈퍼컴퓨터를 필요로 하기 때문입니다. 수년 동안 과학자들은 베이킹을 시작하기 전에 자신의 오븐이 얼마나 커야 할지를 추측해야만 했습니다. 그들은 수백만 달러를 들여 거대한 오븐을 만들었지만, 나중에 알고 보니 중간 크기의 오븐으로도 충분했거나, 혹은 원하는 결과를 얻기 위해 훨씬 더 큰 오븐이 필요했다는 사실을 깨닫곤 했습니다. 큰 질문은 이것이었습니다: 오븐을 켜기도 전에 케이크의 맛을 예측할 수 있을까?

"Predict before you train"(훈련하기 전에 예측하라)이라는 제목의 이 논문은 바로 그 문제를 다룹니다. 연구진은 실제로 거대한 모델을 먼저 구축하지 않고도 특정 결과를 얻기 위해 정확히 얼마만큼의 컴퓨팅 파워가 필요한지 알려주는 "레시피 규칙"(스케일링 법칙이라고 불림)을 찾아낼 수 있는지 확인하고자 했습니다. 그들은 주로 텍스트를 읽는 데 사용되는 "트랜스포머(transformer)"라는 유형의 디지털 뇌에 집중했습니다. 여기서는 이 모델에게 충돌기에서 튀어나오는 입자들을 읽도록 가르쳤습니다. 이들은 작은 규모의 저렴한 모델들을 먼저 훈련함으로써, 모델을 100배 더 크게 만들었을 때 어떤 일이 일어날지 예측하는 그래프 상의 선을 그리려고 시도했습니다. 또한, "사전 훈련(pre-training)"(일반적인 패턴을 학습하는 과정)이 더 나은 사전 훈련이 나중에 탑 쿼크(top quark)를 찾아내거나 쿼크와 글루온을 구별하는 것과 같은 특정 물리 작업에 실제로 도움이 되는지도 알고 싶었습니다.

입자 물리학을 위한 "수정구슬"

스탠퍼드와 SLAC 연구진이 이끄는 팀은 대담한 아이디어를 테스트하기로 했습니다. 아주 작고 저렴한 모델의 결과만을 보고 거대한 머신러닝 모델의 성능을 예측할 수 있을까? 이를 위해 그들은 약 10억 개의 시뮬레이션된 입자 제트(입자들의 분사)를 포함하는 "OmniLearned"라는 데이터셋을 사용했습니다. 그들은 ParticleViT라고 불리는 모델 제품군을 구축했습니다. 이 모델들을 일련의 디지털 탐정이라고 생각하십시오. 가장 작은 모델은 신참 탐정이고, 가장 큰 모델은 방대한 기억력을 가진 베테로 수사관입니다.

연구진은 적은 양의 컴퓨팅 파워(계산 단위인 $10^{19}$ FLOPs 미만)로 신참 탐정들을 훈련시키기 시작했습니다. 그들은 이 작은 탐정들이 얼마나 잘 학습하는지 관찰했습니다. 그런 다음, 수학적 공식(스케일링 법칙)을 사용하여 기억력이 100배 더 큰 탐정을 훈련시킨다면 어떤 일이 벌어질지 예측하는 그래프 상의 선을 그렸습니다.

예측은 실현되었다
그 결과는 마치 수정구슬을 가진 것과 같았습니다. 그들이 실제로 거대 모델을 훈련시켰을 때(최대 $3.4 \times 10^{20}$ FLOPs 사용), 결과는 예측했던 지점에 거의 정확하게 안착했습니다. 실제 성능은 예측치의 1퍼센트 이내였습니다. 이는 지금까지 입자 물리학 분야에서 아직 훈련되지 않은 모델의 성능을 성공적으로 예측한 사례가 없었다는 점에서 매우 중요한 성과입니다. 즉, 과학자들은 이제 거대하고 비싼 훈련 과정을 실제로 진행하기 전에, 정확히 얼마만큼의 컴퓨팅 파워를 구매해야 하는지 알아내기 위해 작고 저렴한 모델로 "예비 실행(dry run)"을 할 수 있게 된 것입니다.

"정보 밀도"의 놀라움

가장 흥ente로운 발견 중 하나는 실제로 얼마나 많은 "두뇌 능력"(모델 크기)이 필요한지에 관한 것이었습니다. 에세이를 쓰거나 여러분과 대화하는 언어 모델(LLM)의 세계에서는 많은 데이터를 처리하기 위해 거대한 모델이 필요하다는 것이 일반적인 규칙입니다. 하지만 입자 물리학의 경우, 연구진은 다른 점을 발견했습니다.

그들은 제트 내의 단일 입자가 문장 속의 단일 단어보다 훨씬 적은 "정보"를 운반한다는 것을 발견했습니다. 단어는 맥락에 따라 수백만 가지의 의미를 가질 수 있지만, 입자는 몇 가지 특정 숫자가 붙어 있는 그저 입자일 뿐입니다. 데이터의 "정보 밀도"가 낮기 때문에, 이 10억 개 입자 데이터셋에 대한 최적의 모델 크기는 언어 전문가들이 추측할 법한 것보다 훨씬 작습니다. 논문은 이 특정 작업의 경우, 거대한 두뇌가 필요한 것이 아니라 단순히 많은 사례를 뇌에 먹여주는 것이 중요하다고 제안합니다. 모델 크기에 대한 "스위트 스팟(최적 지점)"은 언어 모델의 표준 규칙보다 훨씬 작게 나타났으며, 이는 입자 물리학의 경우 더 큰 모델을 만드는 것보다 더 많은 데이터를 제공하는 데 예산을 쓰는 것이 더 낫다는 것을 시사합니다.

더 나은 "일반 지식"이 도움이 되는가?

팀은 또한 일반적인 사전 훈련 작업(모든 입자에 대해 학습하는 것)을 잘 수행하는 것이 나중에 특정 물리 작업에서 더 나은 성능을 내는 데 도움이 되는지 확인했습니다. 그들은 두 가지 일반적인 작업을 테스트했습니다:

1. 탑 태깅(Top Tagging): 입자 제트가 무거운 "탑 쿼크"로부터 왔는지 식별하는 것.
2. 쿼크/글루온 태깅(Quark/Gluon Tagging): 쿼크로 만들어진 제트와 글루온으로 만들어진 제트를 구별하는 것.

그들은 명확하고 직선적인 관계를 발견했습니다. 사전 훈련 중에 일반적인 패턴을 더 잘 학습한(즉, "손실(loss)" 값이 낮은) 모델일수록, 미세 조정(fine-tuning) 후 특정 작업에서도 더 나은 성능을 보였습니다. 다시 말해, 여러분의 디지털 탐정이 일반 입자 물리학에 천재라면, 탑 쿼크를 찾아내는 데 있어서도 천재가 될 것이라는 뜻입니다. 이는 "컴퓨팅 예산"이 기대되는 물리적 성능으로 직접 변환될 수 있음을 확인시켜 줍니다. 여러분이 얼마나 많은 컴퓨팅 파워를 가지고 있는지 알면, 이제 여러분의 모델이 새로운 물리학을 찾아내는 데 얼마나 뛰어날지 예측할 수 있습니다.

최종 대결

마지막으로, 연구진은 자신들의 새로운 범용 "ParticleViT" 모델을 물리 법칙(상대성 이론 등)이 설계에 내장된 것처럼 특화된 고도의 전문 모델인 OmniLearned와 비교했습니다. 결과는 놀라웠습니다. 특별한 물리 규칙이 없는 범용 ParticleViT 모델들이 대부분의 지표에서 특화된 OmniLearned 모델과 대등한 성능을 보였습니다.

특화된 모델이 약간의 우위를 점한 유일한 곳은 매우 어렵고 극단적인 경우인 "탑 태깅(high-purity tail)" 상황이었는데, 여기서 배경 잡음을 거르는 능력이 약 5% 더 좋았습니다. 하지만 그 외 거의 모든 부분에서, 처음부터 학습한 단순한 범용 모델이 똑같이 효과적이었습니다. 이는 많은 물리 작업에 있어, 반드시 복잡하고 물리 특화적인 구조를 만들 필요는 없으며, 충분한 데이터로 훈련된 표준적인 스케일링 트랜스포머만으로도 충분히 효과적으로 임무를 수행할 수 있음을 시사합니다.

이것이 미래에 갖는 의미

논문은 이 다섯 가지 모델(소형부터 초대형까지)과 전체 코드를 공개하며 다른 과학자들이 사용할 수 있도록 합니다. 핵심적인 결론은 이 분야에서 "추측하고 희망하는" 시대는 끝났다는 것입니다. 과학자들은 이제 작고 저렴한 실험을 통해 거대하고 비싼 실험의 결과를 예측할 수 있습니다. 그들은 특정 목표에 도달하기 위해 정확히 얼마만큼의 컴퓨팅 파워가 필요한지 계산할 수 있으며, 이를 통해 시간과 자원을 절약할 수 있습니다. 이 논문이 물리학의 모든 문제를 해결했다고 주장하는 것은 아니지만, 머신러닝 모델을 훈련하는 비용이 많이 드는 환경을 항해할 수 있는 신뢰할 수 있는 지도를 제공했으며, 때로는 작게 시작하는 것이 미래를 보는 가장 좋은 방법임을 증명했습니다.