토큰 임베딩을 활용한 새로운 스케일링 축 JTok과 JTok M

초록

본 논문은 토큰별 파라미터를 도입해 FLOPs는 그대로 두면서 모델 용량을 확장하는 새로운 축을 제시한다. Joint‑Token(JTok)과 그 확장형인 Mixture of Joint‑Token(JTok‑M)를 통해 각 Transformer 레이어에 경량화된 모듈을 삽입하고, 토큰 ID에 기반한 임베딩 벡터를 요소별 곱으로 백본을 조절한다. 실험 결과 650M에서 61B 규모까지 검증했으며, 동일 FLOPs 조건에서 검증 손실 감소와 MMLU·ARC·CEval 등에서 평균 5~9점 상승을 기록한다. 또한 isoFLOPs 분석을 통해 JTok‑M이 기존 MoE 대비 35% 적은 연산량으로 동등한 성능을 달성함을 보였다.

상세 분석

JTok은 기존 Transformer의 Pre‑Norm 구조에 토큰‑인덱스된 임베딩 테이블 Eℓ∈ℝV×d를 추가한다. 입력 토큰 x는 각 레이어 ℓ에서 Eℓ