Nous Research — 토큰 중첩(Token Superposition)을 통한 효율적인 사전 학습: 기술 메커니즘 및 벤치마크 분석

1. 개요 및 배경

거대 언어 모델(LLM)의 사전 학습(Pre-training)은 데이터 스케일링 패러다임 하에서 막대한 연산 비용과 자원을 요구합니다. 높은 데이터 처리량(Data Throughput)을 달성하기 위해 다양한 기법들이 연구되어 왔으나, 대부분 분산 병렬화 처리나 모델 아키텍처 자체를 복잡하게 수정해야 하는 한계가 있었습니다.

Nous Research 연구진이 제안한 토큰 중첩 학습(Token-Superposition Training, 이하 TST)은 기존의 병렬 처리 방식, 최적화 알고리즘(Optimizer), 토크나이저, 데이터셋 및 모델 아키텍처를 전혀 수정하지 않고 그대로 사용할 수 있는 플러그인(Drop-in) 방식의 가속 방법론입니다. TST는 학습 초기 단계에서 여러 개의 연속된 토큰을 하나의 잠재적 표상으로 압축해 연산 효율을 극대화하는 ‘중첩 단계’와 표준 인과적 언어 모델링(Causal Language Modeling)으로 돌아가는 ‘회복 단계’의 2단계 스케줄링을 통해 사전 학습 시간을 대폭 단축합니다.

2. 핵심 개념 및 아키텍처 메커니즘

TST의 핵심 논리는 입력 토큰의 단위를 가상으로 묶어 처리 속도를 높이는 기법과 다중 토큰 예측 신호 메커니즘을 결합한 데 있습니다. 전체 파이프라인의 핵심 컴포넌트와 작동 매커니즘은 다음과 같습니다.

2.1 입력부 중첩 메커니즘: 토큰 임베딩 가방 (Bag-of-Token-Embeddings)

• 데이터 분할 및 압축: 연속된 토큰 데이터 시퀀스를 겹치지 않는 일정한 크기의 토큰 묶음인 ‘가방(Bag)’ 단위로 나눕니다. 여기서 가방의 크기를 중첩 가방 크기(Superposition Bag Size, $s$)라고 부릅니다.
• s-토큰(s-tokens) 생성: 모델의 임베딩 레이어에서는 하나의 가방 내부에 포함된 연속된 토큰 임베딩 벡터들을 단순 산술 평균(Average)하여 하나의 단일 잠재 벡터를 생성합니다. 이렇게 압축된 하나의 표상을 ‘s-토큰’이라고 정의합니다.
• 동일 플롭스(Equal-FLOPs) 제어: 모델이 압축된 형태의 표상 위에서 연산하기 때문에, 각 스텝당 학습 베이스라인과 완전히 동일한 연산량(Equal-FLOPs)을 유지하기 위해서는 입력 데이터 시퀀스의 길이를 의도적으로 가방 크기($s$)배만큼 늘려주어야 합니다. 이 과정을 통해 동일한 연산 비용을 쓰면서도 토큰 소비 속도(Ingestion Rate)를 비약적으로 끌어올릴 수 있습니다.

2.2 출력부 중첩 메커니즘: 멀티핫 크로스 엔트로피 손실 (Multi-hot Cross-Entropy Loss)

• 출력 헤드 예측 구조: 단일 출력 헤드로 다음 가방 전체(미래의 토큰 집합)를 동시에 예측하기 위해 표준 원핫 크로스 엔트로피 손실 함수를 멀티핫 크로스 엔트로피(Multi-hot Cross-Entropy, MCE) 손실 함수로 변형하여 적용합니다. 가방 내 모든 유효 레이블에 균등한 확률 가중치를 부여하는 방식입니다.
• 기존 라이브러리와의 호환성: 실제 학습 역학(Training Dynamics) 관점에서 그래디언트가 0이 되는 상수항(가방 크기에 따른 로그 항)을 제거함으로써 계산식을 단순화합니다. 덕분에 대중적인 사전 학습 라이브러리 내 고도로 최적화된 표준 크로스 엔트로피 커널을 큰 수정 없이 그대로 재사용할 수 있습니다.
• 인과성(Causality) 보존: 준인과적(Semi-causal) 및 준자동회귀적(Semi-autoregressive) 특징을 유지하기 위해, 데이터 분할 전에 표준 차기 토큰 레이블을 왼쪽으로 (가방 크기 — 1)만큼 시프트(Shift)합니다. 이로써 특정 위치의 가방 정보가 인과적으로 바로 다음 위치의 가방 토큰들을 엄밀히 예측하도록 보장합니다.

2.3 2단계 페이즈 오케스트레이션 및 트레이드오프

• TST 중첩 단계 (Phase 1): 전체 학습 스텝 중 개발자가 설정한 일정 비율 동안 s-토큰 변환 및 MCE 손실 함수를 적용해 빠르게 토큰을 소비하며 학습을 진행합니다. 다만 이 단계만 거친 모델은 가방 내부 토큰들의 정확한 순서 정보(Ordering)가 소실되기 때문에, 그대로 추론(Inference)에 사용하면 확률이 뒤섞인 비정상적인 문장을 출력하는 트레이드오프가 발생합니다.
• 회복 단계 (Phase 2): 이 문제를 해결하기 위해, 중첩 단계를 마친 후 TST 모듈과 코드를 완전히 제거하고 표준 차기 토큰 예측 모드로 복귀하여 남은 학습을 진행합니다. 이 회복 단계를 거치면서 모델의 미세한 표현력과 문장 생성 능력이 완벽히 복원됩니다.

3. 실험 결과 및 벤치마크 분석

실험은 270M 소형 모델부터 10B 규모의 혼합 전문가(MoE) 모델에 이르기까지 다양하게 진행되었으며, 분산 컴퓨팅 라이브러리인 TorchTitan 기반 FSDP 병렬화를 적용해 NVIDIA B200 GPU 인프라에서 측정되었습니다.

3.1 벤치마크 데이터 요약

모든 평가는 제로샷(0-shot) 조건 하에 수행되었으며, 핵심 결과는 다음과 같습니다.

• 270M 파라미터 모델 (SmolLM2 기반 변경 아키텍처)
• 기본 베이스라인: 20,000 스텝 학습 (소비 토큰: 42B) $\rightarrow$ 최종 학습 손실: 3.212, HellaSwag: 36.3%, ARC-Challenge: 24.9%
• TST 적용 모델: 20,000 스텝 학습 (소비 토큰: 105B) $\rightarrow$ 최종 학습 손실: 3.142, HellaSwag: 38.6%, ARC-Challenge: 26.4%
• 더 장기적인 100,000 스텝 학습 조건에서도 TST를 적용한 모델이 최종 학습 손실 3.048, HellaSwag 42.6%를 기록하며 베이스라인(최종 손실 3.092, HellaSwag 40.2%)을 일관되게 앞섰습니다.
• 600M 파라미터 모델
• 기본 베이스라인: 최종 학습 손실: 3.019, HellaSwag: 43.5%, ARC-Challenge: 25.5%
• TST 적용 모델: 최종 학습 손실: 2.943, HellaSwag: 48.2%, ARC-Challenge: 26.9%
• 3B 파라미터 모델
• 기본 베이스라인 (50,000 스텝): 최종 학습 손실: 2.640, HellaSwag: 63.9%, ARC-Challenge: 36.8%, MMLU: 33.3%
• TST 적용 모델 (20,000 스텝 단축): 최종 학습 손실: 2.676, HellaSwag: 62.4%, ARC-Challenge: 36.0%, MMLU: 32.8%
• TST를 적용하여 학습 스텝을 2.5배 단축했음에도 불구하고, 베이스라인이 50,000 스텝 동안 도달한 성능과 거의 대등한 수준의 지표를 달성했습니다.
• 10B 혼합 전문가 모델 (MoE, 활성 파라미터 1B)
• MoE 베이스라인: 12,311 GPU-Hours 소모 $\rightarrow$ 최종 학습 손실: 2.252, HellaSwag: 70.1%, ARC-Challenge: 46.3%, MMLU: 37.4%
• MoE TST 적용 모델: 4,768 GPU-Hours 소모 $\rightarrow$ 최종 학습 손실: 2.236, HellaSwag: 71.2%, ARC-Challenge: 47.3%, MMLU: 39.0%
• 결과 분석: 10B 규모의 MoE 모델에서 동일한 최종 손실값(Equal-Loss)에 도달하는 기준을 비교했을 때, TST는 총 사전 학습 시간을 수치상 약 2.58배(논문 공식 표기 2.5배) 단축하는 혁신적인 비용 절감 효과를 증명했습니다.

3.2 아키텍처 실패 사례 및 대안 검증 (Failure Cases)

• 손실 함수 대안 실패: 연구진은 멀티핫 크로스 엔트로피(MCE) 외에도 힌지 손실(Hinge Loss)이나 이진 크로스 엔트로피(Binary Cross-Entropy, BCE) 손실 함수를 가방 예측에 적용하는 실험을 진행했습니다. 그러나 이 두 기법은 모델의 수렴 성능을 심각하게 저하시켰으며, TST를 아예 적용하지 않은 베이스라인보다도 나쁜 최종 지표를 기록하여 최종 폐기되었습니다.
• 가방 크기별 최적 분포 트레이드오프: 가방 내부 토큰들의 손실 가중치를 제어할 때, 대형 윈도우($s=16$) 환경에서는 토큰 거리에 따른 상호 정보량 감쇠를 반영한 ‘파워 로(Power-law)’ 분포가 수렴에 유리했습니다. 반면 소형 윈도우($s=4$) 환경에서는 모든 토큰에 동일한 가중치를 두는 균등(Uniform) 분포가 더 우수한 결과를 보였습니다. 즉, 고정된 단일 손실 가중치 구조는 가방 크기 전반을 완벽히 커버하지 못하는 구조적 한계가 존재합니다.

4. 연구 방법론 및 평가 프로토콜

4.1 평가 프로토콜 및 데이터셋 구성

• 중소형 모델 벤치마크: 270M, 600M, 3B 모델은 SmolLM 및 SmolLM2의 아키텍처 설정을 응용하여 Llama3 모델링 코드 상에서 사전 학습을 수행했습니다. 사전 학습 데이터셋으로는 DCLM 데이터셋을 채택하였고, 스텝당 2M 토큰의 배치 크기를 적용했습니다. 입력 임베딩 가중치와 출력 헤드 가중치를 공유하지 않는 방식을 적용해 파라미터 무결성을 분리했습니다.
• 대형 MoE 모델 검증: Qwen3 모델 구조를 다운스케일링하여 총 파라미터 10B, 활성 파라미터 1B 수준의 가상 MoE 환경을 구축했습니다. 사전 학습 소스로 FineWeb-Edu와 DCLM을 50:50으로 정밀 블렌딩한 토큰 세트를 사용했으며, 스텝당 배치 크기는 8M 토큰으로 확장 통제했습니다.

4.2 데이터 오염 및 데이터 누수 제어 (Contamination Risks)

• 플러그인 코드 완전 격리 프로토콜: TST 기법이 다운스트림 평가 벤치마크(HellaSwag, ARC, MMLU 등)에 인위적인 편향이나 무결성 오염을 일으킬 잠재적 리스크를 원천 차단하기 위해 엄격한 코드 격리를 수행했습니다. 중첩 페이즈가 종료되는 즉시 학습 인프라 내 소스 코드 레벨에서 TST 가방 분할 및 MCE 연산 모듈을 물리적으로 완전히 언로드(Fully Removed)한 뒤 순수 베이스라인 가중치 상태로 회복 단계를 전개함으로써 실험 환경 오염 가능성을 배제했습니다.

5. 전략적 시사점

5.1 엔터프라이즈 사전 학습 인프라의 경제성 확보

기업에서 자체 도메인 특화 LLM을 구축할 때 발생하는 기하급수적인 연산 비용(Compute Cost) 부담을 최대 2.5배 이상 낮출 수 있는 실질적인 경로가 확보되었습니다. 하드웨어 가속기 인프라를 무리하게 확장하지 않고도 동일 자원 대비 고성능 모델을 빠르게 확보할 수 있다는 점에서 기업들의 AI 총소유비용(TCO)을 획기적으로 개선할 수 있습니다.

5.2 RAG 및 고성능 인코더 모델 아키텍처로의 확장성

본 연구에서 입증된 TST의 잠재적 표상 압축력(Compressive Latent Properties)은 단순 사전 학습 가속을 넘어 향후 검색 증강 생성(RAG)의 컨텍스트 압축이나 멀티 토큰 동시 검증 기법으로 진화할 수 있는 이론적 토대를 제공합니다. 여러 개의 연속된 컨텍스트를 정보 손실 없이 단일 s-토큰 벡터 공간으로 압축하여 다룰 수 있으므로, 대규모 지식 베이스를 운용하는 인프라의 메모리 및 연산 효율성을 극대화하는 압축 프라이어(Compressive Prior) 모델 설계가 가능해질 전망입니다.

References:

1. Efficient Pre-Training with Token Superposition (uploaded:2605.06546v1.pdf)

Tags:

NousResearch #TokenSuperposition #PreTrainingEfficiency #MultiHotCrossEntropy #PreTraining #TransformerArchitecture #누스리서치 #토큰중첩학습 #사전학습최적화 #멀티핫크로스엔트로피 #연산비용절감 #트랜스포머아키텍처