Les jetons critiques comptent : l'estimation contrastive au niveau des jetons améliore la capacité de raisonnement des LLM.

Critical Tokens Matter: Token-Level Contrastive Estimation Enhence LLM's Reasoning Capability

November 29, 2024

Auteurs: Zicheng Lin, Tian Liang, Jiahao Xu, Xing Wang, Ruilin Luo, Chufan Shi, Siheng Li, Yujiu Yang, Zhaopeng Tu

cs.AI

Résumé

Les grands modèles de langage (LLMs) ont montré des performances remarquables dans les tâches de raisonnement. Ils utilisent la génération de jetons autorégressive pour construire des trajectoires de raisonnement, permettant le développement d'une chaîne de pensée cohérente. Dans ce travail, nous explorons l'impact des jetons individuels sur les résultats finaux des tâches de raisonnement. Nous identifions l'existence de "jetons critiques" qui conduisent à des trajectoires de raisonnement incorrectes dans les LLMs. Plus précisément, nous constatons que les LLMs ont tendance à produire des résultats positifs lorsqu'ils sont contraints de décoder d'autres jetons au lieu des jetons critiques. Motivés par cette observation, nous proposons une nouvelle approche - cDPO - conçue pour reconnaître automatiquement et mener des récompenses au niveau du jeton pour les jetons critiques pendant le processus d'alignement. Plus précisément, nous développons une approche d'estimation contrastive pour identifier automatiquement les jetons critiques. Cela est réalisé en comparant la probabilité de génération des modèles positif et négatif. Pour ce faire, nous affinons séparément les modèles positif et négatif sur diverses trajectoires de raisonnement, leur permettant ainsi d'identifier les jetons critiques au sein des trajectoires incorrectes qui contribuent à des résultats erronés. De plus, pour aligner davantage le modèle avec les informations des jetons critiques pendant le processus d'alignement, nous étendons les algorithmes DPO conventionnels au niveau du jeton et utilisons la probabilité différentielle des modèles positif et négatif susmentionnés comme poids important pour l'apprentissage du DPO au niveau du jeton. Les résultats expérimentaux sur les référentiels GSM8K et MATH500 avec les modèles largement utilisés Llama-3 (8B et 70B) et deepseek-math (7B) démontrent l'efficacité de l'approche proposée cDPO.

English

Large Language Models (LLMs) have exhibited remarkable performance on reasoning tasks. They utilize autoregressive token generation to construct reasoning trajectories, enabling the development of a coherent chain of thought. In this work, we explore the impact of individual tokens on the final outcomes of reasoning tasks. We identify the existence of ``critical tokens'' that lead to incorrect reasoning trajectories in LLMs. Specifically, we find that LLMs tend to produce positive outcomes when forced to decode other tokens instead of critical tokens. Motivated by this observation, we propose a novel approach - cDPO - designed to automatically recognize and conduct token-level rewards for the critical tokens during the alignment process. Specifically, we develop a contrastive estimation approach to automatically identify critical tokens. It is achieved by comparing the generation likelihood of positive and negative models. To achieve this, we separately fine-tune the positive and negative models on various reasoning trajectories, consequently, they are capable of identifying identify critical tokens within incorrect trajectories that contribute to erroneous outcomes. Moreover, to further align the model with the critical token information during the alignment process, we extend the conventional DPO algorithms to token-level DPO and utilize the differential likelihood from the aforementioned positive and negative model as important weight for token-level DPO learning.Experimental results on GSM8K and MATH500 benchmarks with two-widely used models Llama-3 (8B and 70B) and deepseek-math (7B) demonstrate the effectiveness of the propsoed approach cDPO.