기술 요약: 토큰 중첩을 통한 효율적인 사전 학습

문제 제기

대규모 언어 모델 (LLM) 의 사전 학습은 규모가 커질수록 비용이 급증하고 비효율적이 되어가고 있습니다. 현재의 패러다임은 종종 높은 데이터 처리량을 달성하기 위해 병렬화 전략, 옵티마이저, 토크나이저, 또는 모델 아키텍처에 복잡하고 침습적인 수정을 요구합니다. furthermore, 최근 연구는 추론 시 효율성 향상 또는 학습/추론 효율성의 분리에 초점을 맞추고 있지만, 최종 모델의 아키텍처나 추론 역학을 변경하지 않으면서 학습 시간 처리량을 획기적으로 높이는 방법들은 부족합니다. 기존 접근법들은 종종 학습 효율성과 추론 효율성을 혼동하거나, 반드시 FLOP 당 원시 토큰 처리 속도를 높이는 것은 아닌 보조 손실 (auxiliary losses) 에 의존하여 파라미터와 하이퍼파라미터를 추가합니다.

방법론: 토큰 중첩 학습 (Token Superposition Training, TST)

저자들은 병렬화, 옵티마이저, 토크나이저, 데이터, 또는 모델 아키텍처를 수정하지 않고 사전 학습 중 FLOP 당 데이터 처리량을 향상시키기 위해 설계된 "드롭인 (drop-in)" 방식인 **토큰 중첩 학습 (TST)**을 제안합니다. TST 는 두 가지 명확한 단계로 작동합니다:

1 단계: 중첩 단계 (Superposition Phase)

이 단계에서 모델은 연속된 토큰들의 "백 (bag)"을 단일 잠재 단위인"s-token"으로 처리합니다.

• 입력 중첩: $s$개의 연속된 토큰 시퀀스는 겹치지 않는 백들로 분할됩니다. 이러한 $s$개 토큰의 임베딩은 평균화되어 단일 중첩된 잠재 토큰 임베딩을 생성합니다. 이는 시퀀스 길이를 $s$배만큼 줄입니다 ( $L$에서 $\lfloor L/s \rfloor$로).
• 처리량 메커니즘: 기준 학습 (동일 FLOP) 과 공정한 비교를 보장하기 위해, 이 단계 중 데이터 시퀀스 길이 $L$은 $s$배만큼 증가됩니다. 결과적으로 모델은 단계당 동일한 계산 비용으로 $s$배 더 많은 데이터 토큰을 섭취합니다.
• 출력 중첩: 단일 다음 토큰을 예측하는 대신, 모델은 다음 $s$개 토큰의 "백"을 예측합니다. 표준 원-핫 교차 엔트로피 (CE) 손실은 멀티-핫 교차 엔트로피 (Multi-Hot Cross-Entropy, MCE) 손실로 대체됩니다. MCE 목적 함수는 대상 백을 멀티-핫 레이블로 간주하여 백 내의 각 유효 토큰의 확률을 $1/s$로 끌어당깁니다. 손실은 백 내의 각 토큰에 대한 표준 CE 손실의 평균으로 정의되며, 기울기에 영향을 주지 않는 상수 항 ($\log |y|$) 을 뺀 형태입니다.
• 인과성: 인과성을 유지하기 위해 표준 다음 토큰 예측 레이블은 겹치지 않는 백으로 분할되기 전에 $s-1$만큼 왼쪽으로 이동됩니다.

2 단계: 복구 단계 (Recovery Phase)

순수하게 TST 로만 학습된 모델은 백 내의 토큰 순서를 잃고 표준 자기회귀 샘플링 능력이 부족하여 nonsensical한 출력을 생성합니다. 이를 해결하기 위해:

• 학습은 표준 다음 토큰 인과 예측 목적 함수로 되돌아갑니다.
• 모델은 중첩 단계가 끝날 때 저장된 체크포인트에서 학습을 재개합니다.
• 최종 모델 아키텍처와 추론 역학이 기준선과 동일하도록 TST 코드는 완전히 제거됩니다.
• 중첩 단계에 소요된 단계의 비율은 $r$로 표시됩니다 (예: $r=0.3$은 전체 단계의 30% 가 TST 이고 70% 가 표준임을 의미함).

주요 기여

1. 새로운 학습 패러다임: 모델의 최종 아키텍처나 추론 행동을 변경하지 않고 특정 학습 기간 동안 토큰 처리량을 $s$배 증가시키는 TST 의 도입.
2. 이중 단계 설계: 입력 중첩 (임베딩 평균화) 과 출력 중첩 (멀티-핫 예측) 을 결합한 간단한 메커니즘과 이를 표준 자기회귀 행동으로 정렬하기 위한 복구 단계.
3. 기존 방법론과의 직교성: 저자들은 TST 를 보조 손실 방법 (예: 멀티 토큰 예측) 과 직교한다고 위치시킵니다. MTP 는 파라미터를 추가하고 FLOP 당 토큰 수를 증가시키지 않으면서 미래 토큰을 타겟으로 하는 반면, TST 는 처리량을 엄격하게 증가시키고 아키텍처는 untouched 상태로 둡니다.
4. 실증적 검증: 다양한 하이퍼파라미터 설정에서 견고성을 입증하는 270M, 600M, 3B, 10B 전문가 혼합 (Mixture-of-Experts) 모델 규모에 걸친 광범위한 평가.

결과

DCLM 및 FineWeb-Edu 데이터셋을 사용한 실험에 기반한 이 논문의 주요 발견 사항은 다음과 같습니다:

• 손실 감소: 동등한 FLOP 제약 하에서 TST 는 최종 학습 손실 측면에서 일관되게 기준 모델보다 우수한 성능을 보입니다.
• 하류 성능: TST 로 학습된 모델은 동일한 단계 수 또는 동일한 계산 비용으로 학습된 기준 모델에 비해 하류 벤치마크 (HellaSwag, ARC-E/C, MMLU 등) 에서 향상된 성능을 보입니다.
• 효율성 향상:
  • 동등 손실 설정 하에서 TST 는 10B 파라미터 규모에서 총 사전 학습 시간을 최대 2.5 배 단축합니다.
  • 10B A1B(총 10B, 활성 1B) MoE 실험에서 TST 모델은 기준 모델의 1.05T 토큰과 비교해 2T 데이터 토큰을 처리하면서도 2.236 대 2.252 로 더 낮은 최종 손실을 달성했으며, 이는 B200-시간을 4768 대 12311 로 획기적으로 줄인 결과입니다.
• 하이퍼파라미터 민감도: 이 방법은 합리적인 범위의 중첩 백 크기 ($s \in [4, 8]$) 와 단계 비율 ($r \in [0.2, 0.4]$) 내에서 견고합니다.
• Ablation 연구:
  • 입력 대 출력: 입력 전용 및 출력 전용 중첩 모두 기준선 대비 성능을 향상시키지만, 둘의 조합이 가장 좋은 결과를 내며 이는 이들이 직교하는 메커니즘임을 시사합니다.
  • 표현 정렬: 중첩 단계와 복구 단계 사이에서 입력 임베딩과 출력 LM 헤드를 무작위로 재초기화하면 TST 의 이득이 완전히 사라지므로, 단계 간 표현 정렬을 유지하는 것이 결정적으로 중요합니다.

중요성 및 주장

이 논문은 TST 가 모델의 추론 시간 아키텍처와 독립적으로 학습 시간 처리량을 극대화함으로써 사전 학습 효율성을 획기적으로 향상시킨다고 주장합니다. 저자들은 이 접근법이 "거칠기에서 정밀함 (coarse-to-fine)"의 세분화 일정이 효율적인 사전 학습을 위한 재사용 가능한 요소임을 입증함으로써 단일체 사전 학습 패러다임에 도전한다고 논합니다.

주요 주장은 다음과 같습니다:

• 비용 절감: 약 절반의 계산 비용으로 동일한 모델 품질 (손실) 달성.
• 간결성: 복잡한 아키텍처 변경, 새로운 옵티마이저, 또는 추가 파라미터가 필요 없는 실용적인 "드롭인" 솔루션.
• 효율성 분리: 학습 시간 효율성과 추론 시간 효율성을 성공적으로 분리하여, 현재 문헌에서 두 가지를 혼동하는 많은 방법론들의 간극을 해소함.
• 확장성: 작은 조밀 모델 (270M) 에서부터 대규모 전문가 혼합 모델 (10B) 까지 효과적으로 확장되어 광범위한 적용 가능성을 시사함.

저자들은 근본적인 메커니즘에 대해 겸손하게 언급하며, 입력 중첩이 임베딩 기하학의 일종의 "사전 - 사전 학습" 또는 암시적 정규화 역할을 한다고 가정하지만, 더 많은 해석 가능성 연구가 필요하다고 지적합니다. 또한, 현재 사전 학습이 데이터 제약이 아닌 계산 제약이라는 가정과 계산 제약으로 인한 대규모 Ablation 연구의 부재와 같은 한계점도 인정합니다.