Token Superposition Training(TST): LLM 사전학습 시간을 최대 2.5배 단축하는 Nous Research의 새로운 학습 기법

이 글에서는 대규모 언어 모델(LLM) 사전학습 비용과 시간을 획기적으로 줄일 수 있는 새로운 학습 기법인 Token Superposition Training(TST)을 소개합니다. TST는 모델 구조나 추론 방식은 그대로 유지하면서, 학습 루프만을 수정해 동일한 연산량(FLOPs) 대비 더 많은 텍스트를 학습하게 만드는 방법입니다. Nous Research가 공개한 이 기법은 270M부터 10B 파라미터 모델까지 검증되었으며, 실제로 최대 2.5배의 사전학습 시간 단축 효과를 보여주었습니다.

반응형

---

대규모 LLM 사전학습의 현실적인 문제

최근 LLM 사전학습은 모델 규모와 데이터 양이 급격히 증가하면서 막대한 비용과 시간이 소요됩니다. 특히 동일한 연산 자원에서 얼마나 많은 텍스트를 처리할 수 있는지가 핵심 지표가 되었습니다.
기존에는 BPE와 같은 서브워드 토크나이저로 시퀀스를 줄여 처리량을 높여왔지만, TST는 토크나이저를 바꾸지 않고도 이 처리량을 더 끌어올릴 수 있는 방법을 제시합니다.

---

Token Superposition Training이란 무엇인가

Token Superposition Training(TST)은 두 단계로 구성된 사전학습 방법입니다. 핵심 아이디어는 일정 구간 동안 여러 토큰을 하나의 표현으로 묶어 학습하고, 이후 다시 표준 학습 방식으로 복귀하는 것입니다.

• 모델 아키텍처 변경 없음
• 옵티마이저, 토크나이저, 병렬화 전략 그대로 유지
• 추론 단계는 기존 모델과 완전히 동일

즉, 학습 효율만 개선하고 결과 모델은 기존과 동일하게 사용할 수 있습니다.

---

TST의 동작 방식: 두 단계 학습 구조

1단계: Superposition 단계

전체 학습 스텝 중 일부(r 비율, 보통 0.2~0.4) 동안 적용됩니다.

• 입력 시퀀스 길이 L을 s개의 연속 토큰 묶음(bag)으로 분할
• 각 묶음의 임베딩을 평균 내 하나의 “s-token”으로 변환
• Transformer는 길이 L/s의 시퀀스를 처리
• 대신 시퀀스 길이를 s배 늘려 동일 FLOPs를 유지

출력에서는 다음 토큰 하나가 아니라 다음 s개 토큰 묶음을 예측합니다. 이를 위해 기존 cross-entropy 대신 Multi-hot Cross-Entropy(MCE) 손실을 사용하며, 구현은 기존 CE 커널을 그대로 활용할 수 있습니다.

2단계: Recovery 단계

Superposition 단계 종료 후 체크포인트에서 학습을 재개합니다.

• TST 관련 코드는 완전히 제거
• 표준 next-token prediction 방식으로 학습
• 전환 직후 일시적인 손실 증가(1~2 nats)가 있으나 수천 스텝 내 회복

중요한 점은 임베딩과 LM 헤드를 재초기화하지 않는 것입니다. 두 단계가 동일한 표현을 공유해야 TST 효과가 유지됩니다.

---

실험 결과로 본 TST의 효과

TST는 270M, 600M, 3B dense 모델과 10B-A1B MoE 모델에서 검증되었습니다.

• 3B 모델
  • TST 20k 스텝: 손실 2.676
  • 베이스라인 36k 스텝: 손실 2.677
  • GPU 사용 시간 약 1.8배 절감
• 10B-A1B MoE 모델
  • TST: 4,768 B200 GPU-hours
  • 베이스라인: 12,311 B200 GPU-hours
  • 약 2.5배 학습 시간 단축
  • HellaSwag, ARC, MMLU 등 모든 벤치마크에서 성능 우수

동일 FLOPs 혹은 동일 최종 손실 기준에서는 TST가 일관되게 더 효율적이었습니다. 다만 동일 토큰 수 기준 비교에서는 기존 방식이 유리합니다.

---

입력과 출력, 두 가지 독립적인 메커니즘

연구진은 TST를 입력 측(superposition 임베딩)과 출력 측(다중 토큰 예측) 메커니즘으로 분리해 실험했습니다. 두 방식은 각각 단독으로도 성능 향상을 보였으며, 결합 시 추가 이득을 제공했습니다.
특히 출력 측 방식은 기존 Multi-Token Prediction과 유사하지만, 추가 헤드나 파라미터 없이 구현된다는 점에서 비용 효율성이 높습니다.

---

간단한 구현 개념 예시

TST 구현은 학습 루프에서 세 가지 변경만 필요합니다.

1. 입력 토큰을 s 단위로 묶어 재구성
2. 임베딩 레이어에서 토큰 임베딩 평균 계산
3. 손실 함수에서 MCE 방식으로 s개 토큰에 대한 평균 CE 계산

이 과정은 PyTorch 기반 기존 사전학습 코드에 비교적 간단히 적용할 수 있습니다.

---

언제 TST를 사용해야 하는가

적합한 경우

• 연산 자원이 병목인 환경
• 충분한 학습 데이터가 있는 경우
• 동일 FLOPs에서 더 낮은 손실을 원하는 경우

적합하지 않은 경우

• 데이터 자체가 병목인 환경
• 동일 토큰 수 기준으로 성능 비교가 필요한 경우

---

728x90

Token Superposition Training은 모델 구조를 바꾸지 않고도 LLM 사전학습 효율을 크게 높일 수 있는 실용적인 기법입니다. 학습 시간과 비용이 중요한 현실적인 환경에서, 동일한 추론 모델을 유지하면서 최대 2.5배의 시간 절감 효과를 얻을 수 있다는 점은 매우 큰 의미를 가집니다.
앞으로 대규모 모델 사전학습이 더욱 일반화되는 상황에서, TST는 연산 자원 활용도를 높이는 중요한 선택지로 자리 잡을 가능성이 큽니다.

300x250

https://www.marktechpost.com/2026/05/13/nous-research-releases-token-superposition-training-to-speed-up-llm-pre-training-by-up-to-2-5x-across-270m-to-10b-parameter-models/?fbclid=IwY2xjawRzlHtleHRuA2FlbQIxMQBzcnRjBmFwcF9pZBAyMjIwMzkxNzg4MjAwODkyAAEe5Eh0VFfTu-zX9A_K3ns8PzHSbE52AuLgV6TZjFfjgfgrxwQk_3FGfwpjFX0_aem_N8AhwUwXRnkl0VpbTbxr-g

Nous Research Releases Token Superposition Training to Speed Up LLM Pre-Training by Up to 2.5x Across 270M to 10B Parameter Mode