Treinamento Contínuo Eficiente por Mitigação da Lacuna de Estabilidade

Efficient Continual Pre-training by Mitigating the Stability Gap

June 21, 2024

Autores: Yiduo Guo, Jie Fu, Huishuai Zhang, Dongyan Zhao, Yikang Shen

cs.AI

Resumo

O pré-treinamento contínuo tornou-se cada vez mais a abordagem predominante para adaptar Modelos de Linguagem de Grande Escala (LLMs) a novos domínios. Esse processo envolve atualizar o LLM pré-treinado com um corpus de um novo domínio, resultando em uma mudança na distribuição de treinamento. Para estudar o comportamento dos LLMs durante essa mudança, medimos o desempenho do modelo ao longo do processo de pré-treinamento contínuo. Observamos uma queda temporária no desempenho no início, seguida por uma fase de recuperação, um fenômeno conhecido como "lacuna de estabilidade", anteriormente observado em modelos de visão que classificam novas classes. Para lidar com esse problema e melhorar o desempenho do LLM dentro de um orçamento computacional fixo, propomos três estratégias eficazes: (1) Pré-treinar continuamente o LLM em um subconjunto com um tamanho apropriado por várias épocas, resultando em uma recuperação de desempenho mais rápida do que pré-treinar o LLM em um corpus grande em uma única época; (2) Pré-treinar o LLM apenas em um sub-corpus de alta qualidade, o que impulsiona rapidamente o desempenho no domínio; e (3) Usar uma mistura de dados semelhante aos dados de pré-treinamento para reduzir a diferença de distribuição. Realizamos vários experimentos em modelos da família Llama para validar a eficácia de nossas estratégias tanto no pré-treinamento contínuo médico quanto na afinação de instruções. Por exemplo, nossas estratégias melhoram o desempenho médio da tarefa médica do modelo OpenLlama-3B de 36,2% para 40,7% com apenas 40% do orçamento de treinamento original e aprimoram o desempenho médio da tarefa geral sem causar esquecimento. Além disso, aplicamos nossas estratégias ao modelo Llama-3-8B. O modelo resultante, Llama-3-Physician, alcança o melhor desempenho médico entre os modelos de código aberto atuais e se sai de forma comparável ou até melhor do que o GPT-4 em vários benchmarks médicos. Lançamos nossos modelos em https://huggingface.co/YiDuo1999/Llama-3-Physician-8B-Instruct.

English

Continual pre-training has increasingly become the predominant approach for adapting Large Language Models (LLMs) to new domains. This process involves updating the pre-trained LLM with a corpus from a new domain, resulting in a shift in the training distribution. To study the behavior of LLMs during this shift, we measured the model's performance throughout the continual pre-training process. we observed a temporary performance drop at the beginning, followed by a recovery phase, a phenomenon known as the "stability gap," previously noted in vision models classifying new classes. To address this issue and enhance LLM performance within a fixed compute budget, we propose three effective strategies: (1) Continually pre-training the LLM on a subset with a proper size for multiple epochs, resulting in faster performance recovery than pre-training the LLM on a large corpus in a single epoch; (2) Pre-training the LLM only on high-quality sub-corpus, which rapidly boosts domain performance; and (3) Using a data mixture similar to the pre-training data to reduce distribution gap. We conduct various experiments on Llama-family models to validate the effectiveness of our strategies in both medical continual pre-training and instruction tuning. For example, our strategies improve the average medical task performance of the OpenLlama-3B model from 36.2% to 40.7% with only 40% of the original training budget and enhance the average general task performance without causing forgetting. Furthermore, we apply our strategies to the Llama-3-8B model. The resulting model, Llama-3-Physician, achieves the best medical performance among current open-source models, and performs comparably to or even better than GPT-4 on several medical benchmarks. We release our models at https://huggingface.co/YiDuo1999/Llama-3-Physician-8B-Instruct.