Prédiction autorégressive améliorée par masquage : Moins d'attention pour apprendre davantage

Mask-Enhanced Autoregressive Prediction: Pay Less Attention to Learn More

February 11, 2025

Auteurs: Xialie Zhuang, Zhikai Jia, Jianjin Li, Zhenyu Zhang, Li Shen, Zheng Cao, Shiwei Liu

cs.AI

Résumé

Les grands modèles de langage (LLM) sont découverts comme souffrant de difficultés à récupérer précisément des informations clés. Pour remédier à cela, nous proposons Mask-Enhanced Autoregressive Prediction (MEAP), un paradigme d'entraînement simple mais efficace qui intègre de manière transparente le Modèle de Langage Masqué (MLM) dans la Prédiction du Prochain Token (NTP) pour améliorer les capacités de récupération contextuelle de ce dernier. Plus précisément, MEAP masque d'abord de manière aléatoire une petite fraction des tokens d'entrée, puis effectue directement la prédiction du prochain token standard de manière autorégressive en utilisant un Transformer à décodeur uniquement. MEAP élimine le besoin d'une attention bidirectionnelle ou d'architectures encodeur-décodeur pour le MLM, sans entraîner de surcharge computationnelle supplémentaire lors de la pré-formation ou de l'inférence. Des expériences intensives démontrent que MEAP surpasse considérablement NTP sur la récupération d'informations clés et les tâches de raisonnement à long contexte, tout en performant aussi bien, voire mieux, sur les tâches de raisonnement de bon sens. Les avantages de MEAP s'étendent également au fine-tuning supervisé, où il présente des avantages remarquables dans les scénarios de perte au milieu, surpassant NTP de 11,77 points de pourcentage. Notre analyse indique que l'efficacité de MEAP provient de sa capacité à promouvoir des scores d'attention plus discernables en se concentrant sur un ensemble réduit de tokens non masqués. Ce mécanisme améliore la focalisation du modèle sur les signaux pertinents pour la tâche tout en atténuant l'influence du contexte périphérique. Ces résultats placent MEAP comme un paradigme d'entraînement prometteur pour les grands modèles de langage.

English

Large Language Models (LLMs) are discovered to suffer from accurately retrieving key information. To address this, we propose Mask-Enhanced Autoregressive Prediction (MEAP), a simple yet effective training paradigm that seamlessly integrates Masked Language Modeling (MLM) into Next-Token Prediction (NTP) to enhance the latter's in-context retrieval capabilities. Specifically, MEAP first randomly masks a small fraction of input tokens and then directly performs the standard next-token prediction autoregressive using a decoder-only Transformer. MEAP eliminates the need for bidirectional attention or encoder-decoder architectures for MLM, incurring no additional computational overhead during pre-training or inference. Intensive experiments demonstrate that MEAP substantially outperforms NTP on key information retrieval and long-context reasoning tasks, while performing on par or better on commonsense reasoning tasks. The benefits of MEAP also extend to supervised fine-tuning, where it shows remarkable advantages in lost-in-the-middle scenarios, outperforming NTP by 11.77 percentage points. Our analysis indicates that MEAP's effectiveness arises from its ability to promote more distinguishable attention scores by concentrating on a reduced set of non-masked tokens. This mechanism improves the model's focus on task-relevant signals while mitigating the influence of peripheral context. These findings position MEAP as a promising training paradigm for large language models.