Replanteando la Optimización y la Arquitectura para Modelos de Lenguaje Pequeños

Rethinking Optimization and Architecture for Tiny Language Models

February 5, 2024

Autores: Yehui Tang, Fangcheng Liu, Yunsheng Ni, Yuchuan Tian, Zheyuan Bai, Yi-Qi Hu, Sichao Liu, Shangling Jui, Kai Han, Yunhe Wang

cs.AI

Resumen

El poder de los modelos de lenguaje de gran escala (LLMs) ha sido demostrado a través de numerosos recursos de datos y computación. Sin embargo, la aplicación de modelos de lenguaje en dispositivos móviles enfrenta un gran desafío en cuanto a los costos de computación y memoria, es decir, se requieren urgentemente modelos de lenguaje pequeños con alto rendimiento. Limitados por el proceso de entrenamiento altamente complejo, hay muchos detalles para optimizar modelos de lenguaje que rara vez se estudian cuidadosamente. En este estudio, basado en un modelo de lenguaje pequeño con 1B de parámetros, diseñamos cuidadosamente una serie de estudios empíricos para analizar el efecto de cada componente. Se discuten principalmente tres perspectivas, es decir, la arquitectura neuronal, la inicialización de parámetros y la estrategia de optimización. Varias fórmulas de diseño se demuestran empíricamente especialmente efectivas para modelos de lenguaje pequeños, incluyendo la compresión del tokenizador, el ajuste de la arquitectura, la herencia de parámetros y el entrenamiento en múltiples rondas. Luego entrenamos PanGu-pi-1B Pro y PanGu-pi-1.5B Pro en 1.6T de corpus multilingües, siguiendo las fórmulas establecidas. Los resultados experimentales demuestran que la optimización y arquitectura mejoradas producen una notable mejora promedio de 8.87 en los conjuntos de evaluación de referencia para PanGu-pi-1B Pro. Además, PanGu-pi-1.5B Pro supera a una variedad de modelos SOTA con tamaños de modelo más grandes, validando su rendimiento superior. El código se lanzará pronto (https://github.com/YuchuanTian/RethinkTinyLM).

English

The power of large language models (LLMs) has been demonstrated through numerous data and computing resources. However, the application of language models on mobile devices is facing huge challenge on the computation and memory costs, that is, tiny language models with high performance are urgently required. Limited by the highly complex training process, there are many details for optimizing language models that are seldom studied carefully. In this study, based on a tiny language model with 1B parameters, we carefully design a series of empirical study to analyze the effect of each component. Three perspectives are mainly discussed, i.e., neural architecture, parameter initialization, and optimization strategy. Several design formulas are empirically proved especially effective for tiny language models, including tokenizer compression, architecture tweaking, parameter inheritance and multiple-round training. Then we train PanGu-pi-1B Pro and PanGu-pi-1.5B Pro on 1.6T multilingual corpora, following the established formulas. Experimental results demonstrate the improved optimization and architecture yield a notable average improvement of 8.87 on benchmark evaluation sets for PanGu-pi-1B Pro. Besides, PanGu-pi-1.5B Pro surpasses a range of SOTA models with larger model sizes, validating its superior performance. The code will be released soon (https://github.com/YuchuanTian/RethinkTinyLM).