Predicción Autoregresiva Mejorada con Máscara: Prestar Menos Atención para Aprender Más

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

February 11, 2025

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

cs.AI

Resumen

Se ha descubierto que los Modelos de Lenguaje Grandes (LLMs) sufren de dificultades para recuperar con precisión información clave. Para abordar esto, proponemos Predicción Autoregresiva Mejorada por Máscara (MEAP), un paradigma de entrenamiento simple pero efectivo que integra de manera fluida el Modelado de Lenguaje Enmascarado (MLM) en la Predicción del Próximo Token (NTP) para mejorar las capacidades de recuperación en contexto de este último. Específicamente, MEAP primero enmascara aleatoriamente una pequeña fracción de tokens de entrada y luego realiza directamente la predicción estándar del próximo token de forma autoregresiva utilizando un Transformer solo con decodificador. MEAP elimina la necesidad de atención bidireccional o arquitecturas codificador-decodificador para MLM, sin incurrir en sobrecarga computacional adicional durante el pre-entrenamiento o la inferencia. Experimentos intensivos demuestran que MEAP supera sustancialmente a NTP en la recuperación de información clave y tareas de razonamiento de largo contexto, al tiempo que se desempeña igual o mejor en tareas de razonamiento común. Los beneficios de MEAP también se extienden al ajuste fino supervisado, donde muestra ventajas notables en escenarios de "perdido en el medio", superando a NTP en 11.77 puntos porcentuales. Nuestro análisis indica que la efectividad de MEAP surge de su capacidad para promover puntuaciones de atención más distinguibles al concentrarse en un conjunto reducido de tokens no enmascarados. Este mecanismo mejora el enfoque del modelo en señales relevantes para la tarea mientras mitiga la influencia del contexto periférico. Estos hallazgos posicionan a MEAP como un paradigma de entrenamiento prometedor para grandes modelos de lenguaje.

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.