Zayd M. K. Zuhri Erland Hilman Fuadi Alham Fikri Aji

초록

다중 토큰 예측(Multi-Token Prediction, MTP)은 언어 모델 학습에서 다음 토큰 예측(Next-Token Prediction, NTP) 성능을 향상시키기 위한 보조 목표로 제안되었으나, 표준 자연어 처리(NLP) 벤치마크에서 일관된 성능 향상을 보이지 못하고 오히려 성능이 떨어지는 경우가 있다. 본 연구에서는 MTP가 정확한 미래 토큰 예측을 보조 손실로 요구하는 데 있어 과도하게 어려운 문제라고 주장한다. 대신, 우리는 학습-순위화(learning-to-rank) 손실을 활용하여 향후 등장할 토큰들을 그 발생 가능성에 따라 순서대로 정렬하도록 모델을 훈련하는 토큰 순서 예측(Token Order Prediction, TOP)을 제안한다. TOP는 MTP가 요구하는 다수의 트랜스포머 계층 대비 단 하나의 추가적인 언임베딩 레이어만 필요하다. 본 연구는 NTP, MTP, TOP 목표를 사용하여 340M, 1.8B, 7B 파라미터 규모의 모델을 사전 훈련하였다. 8개의 표준 NLP 벤치마크에서의 실험 결과, TOP는 규모가 커져도 NTP 및 MTP를 모두 상회하는 전반적인 성능을 보였다. 본 연구의 코드는 https://github.com/zaydzuhri/token-order-prediction 에 공개되어 있다.