MIT의 LLM 훈련 효율 2배 향상 방법론: 토큰 가중치부터 중복 연산 제거까지

표준적인 LLM 훈련은 모든 훈련 샘플을 동등하게 취급하고, 모든 레이어의 모든 파라미터를 매 스텝마다 동일한 강도로 업데이트한다. MIT Computer Science and Artificial Intelligence Laboratory(CSAIL) 연구팀은 이 방식이 근본적으로 비효율적이라는 데서 출발해, 대형 언어 모델 훈련의 효율을 두 배로 향상시키는 방법론을 발표했다. 동일한 규모의 모델을 동일한 성능 수준으로 훈련하는 데 필요한 컴퓨팅 자원을 절반으로 줄일 수 있다는 이 연구는, 수억 달러의 훈련 비용이 필요한 최첨단 AI 모델 개발의 경제성을 근본적으로 바꿀 수 있는 잠재력을 담고 있다.

이미 학습된 패턴에 반복 투입되는 연산

현재의 표준적인 LLM 훈련은 Stochastic Gradient Descent(SGD)와 그 변형인 Adam 옵티마이저를 사용한다. 이 방식의 핵심 특징은 균일성이다. 각 미니배치 내의 모든 토큰이 동일한 가중치로 손실 함수에 기여하고, 모든 레이어의 파라미터가 동일한 학습률로 업데이트된다.

그러나 실제 훈련 데이터는 균일하지 않다. 자주 등장하는 흔한 패턴과 드물지만 중요한 패턴이 섞여 있고, 모델의 학습이 진행될수록 쉬운 패턴은 이미 잘 학습되어 있어 추가 학습이 거의 필요 없는 상태가 된다. 이런 이미 학습된 패턴에 대한 연산이 전체 훈련 비용의 상당 부분을 차지하지만, 모델 성능에는 거의 기여하지 않는다.

⤢✕

GPT-4 수준의 모델을 훈련하는 데 드는 비용은 수천만에서 수억 달러에 달한다. A100 또는 H100 GPU 수만 개를 수주에서 수개월간 구동하는 이 규모의 연산은 전력 소비와 탄소 배출 면에서도 상당한 환경적 부담을 만든다. 훈련 비용의 2배 절감은 이 모든 수치를 절반으로 줄인다는 의미다.

연구진은 7B에서 70B 파라미터 규모의 모델들을 대상으로 광범위한 실험을 수행했으며, 제안된 방법론이 다양한 모델 아키텍처와 훈련 데이터셋에서 일관된 2배 효율 향상을 달성함을 보였다. 이를 뒷받침하는 것이 다음 세 가지 핵심 기법이다.

모델이 어려워하는 토큰에 학습을 몰아준다

첫 번째 기법은 훈련 중 각 토큰의 중요도를 동적으로 평가하고, 중요도에 따라 손실 함수에서의 가중치를 조정하는 동적 토큰 가중치 조정이다. 연구팀은 모델이 이미 높은 확신도로 예측할 수 있는 토큰들은 훈련에서의 가중치를 줄이고, 모델이 불확실해하거나 오류를 범하는 토큰들에 더 많은 학습 신호를 집중하는 알고리즘을 개발했다. 이 방법론의 핵심 지표는 퍼플렉시티(Perplexity)다 — 모델이 다음 토큰을 예측하는 데 얼마나 불확실한지를 나타내는 이 지표를 실시간으로 계산하고, 불확실도가 높은 토큰에 학습 자원을 집중하는 방식이다.

⤢✕

두 번째 기법인 계층별 적응형 학습률은 Transformer 모델의 각 레이어가 훈련 과정에서 서로 다른 속도로 수렴한다는 관찰에 기반한다. 초기 레이어(기본적인 언어 패턴 담당)는 훈련 초반에 빠르게 수렴하는 반면, 상위 레이어(복잡한 추론 담당)는 더 오랜 학습이 필요하다. 표준 훈련에서는 모든 레이어에 동일한 학습률이 적용되어 이미 수렴한 레이어에도 지속적인 업데이트가 이뤄진다. MIT의 방법은 각 레이어의 수렴 상태를 실시간으로 모니터링하고, 수렴이 진행된 레이어의 학습률을 점진적으로 낮추는 방식으로 불필요한 업데이트를 줄인다.

세 번째 기법인 중복 연산 제거는 Forward Pass와 Backward Pass에서 발생하는 중복 연산을 식별하고 제거하는 방법이다. 미니배치 내의 유사한 샘플들은 중간 계산 결과를 공유할 수 있는 경우가 많지만, 표준 구현에서는 각 샘플이 독립적으로 처리된다. 연구팀은 배치 내 샘플 간의 유사도를 효율적으로 측정하고, 유사도가 높은 샘플 그룹에서 계산을 공유하는 알고리즘을 구현했다. 이 기법은 특히 긴 문맥의 시퀀스를 처리할 때 효과적이며, Self-Attention 연산의 중복을 크게 줄인다.

벤치마크로 확인한 결과와 한계

연구팀은 제안된 방법론을 LLaMA, Mistral, Falcon 아키텍처에 적용해 표준 훈련 방식과 비교했다. 결과는 세 가지 척도에서 일관되게 나타났다. 동일한 컴퓨팅 예산으로 훈련했을 때 새로운 방법론이 더 높은 최종 성능을 달성했다 — MMLU, HellaSwag, ARC 등 주요 언어 모델 벤치마크에서 평균 3-5% 성능 향상이 관찰됐다. 동일한 목표 성능에 도달하는 데 필요한 GPU 시간이 평균 48% 감소했는데, 이것이 "2배 효율 향상"의 근거다. 방법론의 효과는 모델 규모에 무관하게 일관성 있게 나타났다 — 7B, 13B, 34B, 70B 파라미터 모델 모두에서 유사한 수준의 효율 향상이 확인됐다.

연구팀은 재현성을 위해 전체 코드와 훈련 설정을 공개했으며, 독립적인 연구 그룹들이 결과를 재현하는 검증 과정이 진행 중이다. 다만 연구팀 자신도 몇 가지 한계를 인정한다. 동적 토큰 가중치 계산과 계층별 학습률 모니터링에는 추가적인 메모리와 연산 오버헤드가 발생하는데, 이 오버헤드가 효율 향상의 약 10-15%를 상쇄한다. 또한 매우 작은 규모(3B 이하)의 모델에서는 효과가 다소 줄어드는 경향이 관찰됐다.

진입 장벽과 탄소 발자국을 함께 낮춘다

이 연구의 가장 중요한 함의는 AI 모델 개발의 진입 장벽을 낮춘다는 점이다. 현재 최첨단 AI 모델 훈련은 대형 빅테크 기업과 잘 자금 지원된 스타트업만이 감당할 수 있는 비용을 요구한다. 훈련 비용이 절반으로 줄어든다면 중소 기업, 대학, 연구기관도 경쟁력 있는 모델을 훈련할 수 있는 환경이 만들어진다.

⤢✕

AI 훈련의 탄소 발자국은 꾸준히 비판받아 온 문제다. GPT-3 훈련은 미국 가정 120가구의 1년 전력 소비에 해당하는 에너지를 소비한 것으로 알려져 있다. 훈련 효율의 2배 향상은 이 환경적 비용을 절반으로 줄이는 직접적인 효과를 가진다.

기존 효율화 기법과는 대체가 아니라 보완 관계

MIT의 연구는 LLM 훈련 효율화를 목표로 하는 광범위한 연구 흐름의 일부다. 비교적 최근의 주요 접근법으로는 데이터와 파라미터의 최적 비율을 연구하는 Chinchilla 스케일링 법칙, 필요한 파라미터만 활성화하는 MoE(Mixture of Experts) 아키텍처, 메모리 효율적인 어텐션 연산인 Flash Attention이 있다. MIT의 방법론은 이 기법들과 독립적으로 적용 가능하며, 실제로 MoE 아키텍처와 Flash Attention을 함께 사용했을 때 각 기법의 효과가 더해지는 시너지가 관찰됐다. 즉 MIT의 방법론은 기존 효율화 기법을 대체하는 것이 아니라 보완하는 위치에 있다.

MIT 연구팀은 후속 연구 방향으로 세 가지를 제시했다. 현재 방법론이 최적화된 사전 훈련(Pre-training)을 넘어 RLHF 같은 정렬 훈련(Alignment Training) 단계로 확장하는 파인튜닝 효율화, 텍스트만 처리하는 언어 모델이 아닌 이미지·오디오·비디오를 함께 처리하는 멀티모달 모델로의 확장, 스마트폰이나 엣지 디바이스에서의 소형 모델 훈련에 적용하는 온디바이스 모델 훈련 적용이다. 마지막 방향이 의미 있는 효율 향상을 가져올 수 있다면 개인화된 AI의 보급에 크게 기여할 수 있다.

MIT CSAIL의 LLM 훈련 효율 2배 향상 연구는 AI 모델 개발의 경제성과 접근성을 동시에 개선하는 실용적 의미를 가진다. 동적 토큰 가중치 조정, 계층별 적응형 학습률, 중복 연산 제거라는 세 기법의 조합은 기존 훈련 패러다임의 근본적인 비효율을 겨냥한다.