Jan 2, 2025

Publication

rBridge: 추론 모델의 학습 비용을 100배 절감하는 방법론

Main Project Image
Main Project Image
Main Project Image
Main Project Image

요약 (TL;DR): 본 연구에서는 초소형 프록시 모델만으로도 훨씬 더 큰 언어 모델의 추론 성능을 신뢰도 높게 예측할 수 있는 새로운 방법론인 rBridge를 제안합니다. 이를 통해 사전 학습 데이터 선별 및 하이퍼파라미터 최적화 비용을 100배에서 최대 700배까지 절감할 수 있습니다. 자세한 내용은 여기에서 확인할 수 있습니다.

Trillion Labs는 AGI로 가는 길이 단순한 규모의 확장만이 아닌 지속가능한 비즈니스를 위한 효율성에 있다고 생각합니다. 딥시크가 그랬듯, 최소한의 자원과 인력을 투입해 Frontier 수준의 추론 능력을 이끌어내는 방법을 발견하는 것이 AI의 게임 체인저가 될 것입니다.

이러한 맥락에서 rBridge는 Frontier AI를 향한 경쟁의 법칙을 근본적으로 바꿉니다. 1B 이하의 초소형 프록시 모델이 해당 모델보다 13배에서 32배 더 큰 모델의 추론 성능을 안정적으로 예측할 수 있게 합니다. 이는 추론 능력을 위한 사전 학습 최적화 비용을 100배 이상 절감하는 혁신적인 방법론입니다.

과거에 수백만 달러의 비용과 수개월의 시간이 필요했던 실험들을 이제는 빠르고, 저렴하며, 보다 정확하게 탐색할 수 있습니다. 이는 추론 능력을 갖춘 AI로 향하는 근본적으로 더 효율적인 길을 열어줄 것입니다.

사전 학습의 핵심, '최적의 데이터 레시피'를 찾아라

모델이 보는 토큰의 99%, 전체 Compute(연산 자원)의 99%를 차지하는 사전 학습(Pre-training)은 모델의 기초 지능을 결정하는 가장 중요한 과정입니다. 그리고 그 핵심은 바로 '데이터 레시피'입니다. 수학, 코드, 웹 텍스트, 합성 데이터(Synthetic data) 등 방대한 데이터 소스 중 무엇을, 어떤 비율로 학습시키느냐에 따라 모델의 최종 성능은 천차만별로 달라집니다.

예를 들어, 수학 데이터를 너무 적게 넣으면 추론 능력이 떨어지고, 반대로 너무 많이 넣으면 일반 대화 능력이 저하될 수 있습니다. 특히 다른 AI 모델이 생성한 '합성 데이터'가 실제 학습에 긍정적인 효과를 주는지 사전에 검증하는 과정은 필수적입니다.

문제는 모델 학습을 완전히 끝내기 전까지는 이 레시피의 성공 여부를 확신할 수 없다는 점입니다. 수십억 원 상당의 GPU 자원을 투입한 뒤에야 "데이터 레시피가 잘못되었다"는 것을 깨닫는다면, 이는 돌이킬 수 없는 막대한 손실로 이어집니다.

이러한 리스크를 줄이기 위한 업계의 표준 해법은 '작은 프록시(Proxy) 모델'을 먼저 학습시켜 보는 것입니다. "작은 모델에서 효과적인 데이터라면, 큰 모델에서도 똑같이 작동할 것"이라는 가정하에 최적의 데이터셋을 선별하는 전략입니다. 다행히 일반적인 언어 능력의 경우 이 전략은 꽤 유효합니다. 100M 파라미터 모델에서 확인된 성능 향상이 7B, 혹은 그보다 큰 대형 모델에서도 비교적 일정하게 재현되기 때문입니다.

초등학생에게 미적분을 시킬 수 없다

하지만 여기에 치명적인 함정이 있습니다. 단순한 지식 습득과 달리, 수학이나 코딩과 같은 고차원적 추론 능력은 모델이 일정 크기 이상 커져야만 비로소 나타나는 창발적(Emergent) 특성이라는 점입니다.

1B 이하의 작은 모델에서 추론 성능을 측정하면, 데이터의 품질과 상관없이 거의 무작위에 가까운 결과가 나오거나, 심지어 데이터가 좋을수록 성능 지표가 떨어지는 역설적인 현상이 발생하기도 합니다.

MATH500 벤치마크를 통한 모델 규모별 사전 학습 진행 상황 비교. 동일한 데이터 소스(OLMo-Mix-1124)를 사용했을 때, 작은 모델(1B, 7B)은 성능 변화의 방향성조차 잘못 예측하거나 극심한 노이즈를 보여줍니다.

위 그래프에서 볼 수 있듯, 1B 모델의 정확도(Accuracy)는 노이즈에 불과해 32B 모델의 성능을 전혀 예측하지 못합니다. 이로 인해 지금까지 엔지니어들은 추론 성능을 예측하기 위해 울며 겨자 먹기로 7B 이상의 무거운 프록시 모델을 사용해야 했습니다.

문제는 실험이 거듭될수록 기하급수적으로 불어나는 비용입니다. 7B 모델을 단 한 번 학습시키는 데 약 2억 원, 21B 모델은 8억 원이 소요됩니다. 하지만 최적의 데이터 레시피를 찾기 위해 최소 10가지의 배합을 비교 실험한다고 가정해 봅시다. 7B 모델만으로도 순식간에 20억 원이 증발합니다.

결국 높은 비용 장벽은 자본이 풍부한 소수의 Big Tech Lab에게만 유리한 환경을 제공하며, 대다수 연구자와 스타트업에게는 기술 발전의 기회를 가로막는 큰 장애물이 됩니다.

Trillion Labs는 이 문제를 해결하기 위해, 1B 미만의 아주 작은 모델로도 거대 모델의 추론 성능을 정확히 예측할 수 있는 기술, rBridge를 개발했습니다.

rBridge의 핵심 철학: 정답이 아닌 풀이 과정을 평가하라

rBridge는 놀랍도록 단순한 통찰로 이 문제를 해결합니다.

소형 모델을 프록시로 사용한 평가 방식이 실패하는 이유는 소형 모델의 능력이 부족해서가 아니라, 우리가 그들을 잘못 평가하고 있기 때문입니다.

기존의 사전 학습 평가는 다음 토큰 예측을 통해 NLL(Negative Log-Likelihood)을 낮추는 과정이었습니다. 그러나 정확도나 Pass@K와 같은 평가 지표는 소형 모델이 결과적으로 '정답을 맞췄는가'에만 집중합니다. rBridge는 이 한계를 극복하기 위해 GPT-4o 등 Frontier Model의 '추론 과정 (Reasoning Trace)'을 정답으로 활용합니다.

1. Reasoning Trace 도입

단순히 "정답: 42"를 학습 데이터로 쓰는 대신, Frontier Model이 생성한 추론 과정 전체를 정답지로 사용합니다. 이 추론 과정은 단순 정답보다 훨씬 더 자연어 분포에 가까워, 작은 모델도 충분히 학습하고 모방할 수 있는 신호를 제공합니다. 즉, 작은 모델이 비록 '답'은 맞추지 못하더라도, 뛰어난 모델의 '사고방식'을 얼마나 잘 따라가는지를 측정하면 모델이 커졌을 때의 잠재력을 파악할 수 있다는 것입니다.

2. Task-Aligned Weighted NLL (토큰별 가중치 적용)

하지만 단순히 추론 과정을 따라가는 것만으로는 부족합니다. 문장 내에는 "The answer is..." 같은 형식적인 토큰과, "sum modulo 9" 같은 문제 해결에 핵심적인 토큰이 섞여 있기 때문입니다. rBridge는 Frontier Model의 확률 분포를 활용하여, 중요한 토큰에는 더 큰 가중치를, 덜 중요한 토큰에는 낮은 가중치를 부여하는 새로운 NLL 수식을 고안했습니다. 이를 통해 토크나이저가 달라도 적용 가능한 범용적인 평가 지표를 완성했습니다.


rBridge의 개요. 프론티어 모델의 추론 과정을 골드 레이블로 사용하고, 평가를 위해 가중치가 적용된 NLL을 계산합니다

압도적인 효율성, 1/100의 비용으로 더 높은 정확성 달성

Trillion Labs는 rBridge의 성능을 검증하기 위해 수학(MATH500), 코딩(HumanEval), 과학(ARC-C, MMLU Pro) 등 6개 벤치마크에서 강도 높은 실험을 진행했습니다.

1. 데이터셋 랭킹 비용 100배 절감

가장 놀라운 성과는 데이터셋 선별 비용의 절감입니다. 1.2B 모델을 위한 최적의 데이터를 찾기 위해 25개의 데이터셋 랭킹을 매기는 실험에서, rBridge는 기존 Baseline 대비 100배 이상의 연산량(FLOPs)을 절감하면서도 80.8%의 높은 정확도를 달성했습니다.

2. 프록시 모델과 타겟 모델 성능 간의 강력한 상관관계

1B 프록시 모델과 32B 타겟 모델 간의 성능 상관관계를 분석했을 때, rBridge는 기존 지표들을 압도했습니다. 이는 rBridge 점수가 높으면, 실제 거대 모델을 학습시켰을 때 추론 성능이 높을 것이라고 80% 이상의 확신을 가지고 예측할 수 있음을 의미합니다. 특히 1B 모델에서 rBridge를 사용하는 것이, 7배나 더 큰 7B 모델을 학습시켜 정확도를 측정하는 것보다 타겟 모델(32B)과의 상관관계가 더 높게 나타났습니다. 즉, 1/7도 안 되는 자원으로 더 정확하게 미래의 성능을 예측하는 것입니다.

3. Zero-shot 전이 학습

더 나아가, rBridge로 학습된 상관관계 함수는 전혀 다른 새로운 데이터셋에도 적용 가능했습니다. 이는 한 번의 실험 세팅으로 향후 추가될 수많은 데이터셋의 가치를 추가 학습 없이 즉시 평가할 수 있는 확장성을 제공합니다.

rBridge 의 의의, 사전 학습의 불확실성을 줄이다

AI 모델 개발에서 가장 큰 어려움 중 하나는 성과를 사전에 예측하기 어렵다는 점입니다. 투입되는 자본, 실험, 데이터가 실제로 기대한 수준의 성능 향상으로 이어질지, 대규모 학습을 진행하기 전까지는 확인하기 어려웠습니다. 이러한 불확실성 때문에 많은 팀과 기업은 비용과 리스크가 큰 실험을 피하거나, 제한된 범위의 시도에 머물 수밖에 없었습니다.

rBridge는 이러한 제약을 근본적으로 완화합니다. 초소형 프록시 모델을 통해 대규모 모델의 추론 성능을 사전에 예측할 수 있게 함으로써, 이전에는 감에 의존하거나 시도하기 어려웠던 사전 학습 전략과 데이터 설계를 100배 저렴한 가격으로 탐색할 수 있게 합니다. 다시 말해, 그동안 높은 비용과 불확실성으로 모델 개발을 주저했던 기업들도 이제 사전 검증을 통해 리스크를 관리하며 도메인 특화 모델 개발을 현실적인 선택지로 고려할 수 있게 되었습니다.

이처럼, rBridge는 Trillion Labs가 "어떻게 남들보다 빠르고 효율적으로, 압도적인 성능의 모델을 개발하는가"에 대한 하나의 해결책을 제시합니다.

Trillion Labs의 여정에 함께하고 싶으신가요?

Trillion Labs는 rBridge와 같이 기존의 상식을 깨는 기술로 AI의 한계를 넓혀갈 엔지니어와 연구원을 찾고 있습니다. 우리가 정의하는 문제는 어렵지만, 그만큼 임팩트는 거대합니다. 가장 효율적인 방법으로, 가장 임팩트 있는 문제를 푸는 곳. Trillion Labs는 당신의 합류를 기다립니다.


Join our newsletter

Get the latest AI news and insights every week