Predição Autoregressiva Aprimorada por Máscara: Prestar Menos Atenção para Aprender Mais

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

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) são descobertos como sofrendo para recuperar com precisão informações-chave. Para lidar com isso, propomos a Predição Autoregressiva Aprimorada por Máscara (MEAP), um paradigma de treinamento simples, porém eficaz, que integra perfeitamente a Modelagem de Linguagem Mascara (MLM) na Predição do Próximo Token (NTP) para aprimorar as capacidades de recuperação no contexto deste último. Especificamente, o MEAP primeiro mascara aleatoriamente uma pequena fração dos tokens de entrada e em seguida realiza diretamente a predição do próximo token padrão autoregressiva usando um Transformer apenas com decodificador. O MEAP elimina a necessidade de atenção bidirecional ou arquiteturas codificador-decodificador para o MLM, sem acarretar custos computacionais adicionais durante o pré-treinamento ou inferência. Experimentos intensivos demonstram que o MEAP supera substancialmente o NTP em tarefas de recuperação de informações-chave e raciocínio de longo contexto, enquanto se mantém no mesmo nível ou melhor em tarefas de raciocínio de senso comum. Os benefícios do MEAP também se estendem ao ajuste fino supervisionado, onde ele mostra notáveis vantagens em cenários de perda no meio, superando o NTP em 11,77 pontos percentuais. Nossa análise indica que a eficácia do MEAP decorre de sua capacidade de promover pontuações de atenção mais distinguíveis concentrando-se em um conjunto reduzido de tokens não mascarados. Esse mecanismo melhora o foco do modelo em sinais relevantes para a tarefa, ao mesmo tempo que mitiga a influência do contexto periférico. Essas descobertas posicionam o MEAP como um paradigma de treinamento promissor para grandes modelos de linguagem.

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.