Rho-1: Not All Tokens Are What You Need (2404.07965v3)

Abstract: Previous LLM pre-training methods have uniformly applied a next-token prediction loss to all training tokens. Challenging this norm, we posit that ''9l training''. Our initial analysis examines token-level training dynamics of LLM, revealing distinct loss patterns for different tokens. Leveraging these insights, we introduce a new LLM called Rho-1. Unlike traditional LMs that learn to predict every next token in a corpus, Rho-1 employs Selective LLMing (SLM), which selectively trains on useful tokens that aligned with the desired distribution. This approach involves scoring pretraining tokens using a reference model, and then training the LLM with a focused loss on tokens with higher scores. When continual pretraining on 15B OpenWebMath corpus, Rho-1 yields an absolute improvement in few-shot accuracy of up to 30% in 9 math tasks. After fine-tuning, Rho-1-1B and 7B achieved state-of-the-art results of 40.6% and 51.8% on MATH dataset, respectively - matching DeepSeekMath with only 3% of the pretraining tokens. Furthermore, when continual pretraining on 80B general tokens, Rho-1 achieves 6.8% average enhancement across 15 diverse tasks, increasing both efficiency and performance of the LLM pre-training.

• The paper introduces Selective Language Modeling (SLM), emphasizing that prioritizing high-loss tokens enhances training efficiency.
• Empirical results show up to a 30% absolute few-shot accuracy gain on math tasks, with state-of-the-art performance on the MATH dataset.
• The study underscores resource-efficient training and a nuanced understanding of token-level dynamics, paving the way for adaptive training strategies.

Introducing Rho-1: Advancing Efficiency in LLM Training with Selective LLMing

Overview of Selective LLMing (SLM)

The innovation presented in this research pivots on the introduction of Selective LLMing (SLM), a methodology that deviates from the traditional approach of training LLMs by treating every token within the training corpus with equal importance. This paper propounds a discerning stance, advocating that not all tokens contribute equally towards the efficacious training of LLMs. By deploying a reference model to evaluate and select tokens based on their utility and alignment with the desired distribution, SLM focuses training efforts on tokens exhibiting higher excess loss. This approach marks a strategic shift towards optimizing training processes, underscoring efficiency and targeted learning.

Empirical Validation and Results

Empirical results substantiate the efficacy of SLM, showcasing significant improvements across various tasks. When applied to the mathematical domain through continuous pretraining on the 15B OpenWebMath corpus, the Rho-1 models demonstrated an absolute few-shot accuracy improvement of up to 30% on nine math tasks. Furthermore, after fine-tuning, the Rho-1 models (1B and 7B versions) surpassed state-of-the-art results on the MATH dataset, achieving accuracies of 40.6% and 51.8% respectively, with notably lesser pretraining tokens compared to DeepSeekMath. In general domain pretraining on 80 billion tokens, Rho-1 underscored its utility by delivering a 6.8% average improvement across fifteen diverse tasks. These numerical results reinforce the premise that SLM not only enhances model performance but also operationalizes a more resource-efficient training procedure.

Implications and Theoretical Contributions

The SLM's methodological contributions extend beyond empirical success, presenting several theoretical and practical implications:

• Efficiency in Training: By pinpointing and prioritizing tokens that are pivotal for model learning, SLM conserves computational resources and accelerates the training cycle.
• Token Dynamics Understanding: The differentiation between "easy" and "hard" tokens introduces a nuanced understanding of token-level learning dynamics, providing insights into how models interact with diverse data subsets during training.
• Strategic Data Utilization: SLM embodies a strategic approach to data utilization, ensuring that training efforts are concentrated on data segments that promise the greatest returns in model performance.

Future Directions

The promising results of SLM open avenues for further exploration and refinement. Future work could delve into the optimization of token selection criteria, exploring dynamic or adaptive mechanisms that evolve with the model's learning trajectory. Moreover, the application of SLM across broader domains and model architectures presents an interesting frontier, potentially unveiling domain-specific insights and customization strategies for model training.

Conclusion

The introduction of Selective LLMing (SLM) heralds a thoughtful reconsideration of how training resources are allocated in the development of LLMs. By privileging the quality of tokens over quantity, SLM achieves remarkable efficiency and efficacy, casting a new light on the path towards optimizing LLM training. This research enriches the tapestry of generative AI and machine learning with a method that succinctly aligns training focus with the most beneficial data points, marking a step forward in the journey towards more intelligent and resource-aware computational models.