Modelado de secuencias largas con XGen: Un LLM de 7B entrenado con secuencias de entrada de 8K de longitud

TLDR

Entrenamos una serie de 7B LLMs llamados XGen-7B con atención densa estándar en hasta 8K de longitud de secuencia para hasta 1.5T tokens. También afinamos los modelos en datos de instrucción de dominio público. Los principales resultados son:

• En pruebas estándar de PNL, XGen consigue resultados comparables o mejores que los LLM de código abierto más avanzados (por ejemplo, MPT, Falcon, LLaMA, Redpajama, OpenLLaMA) de tamaño de modelo similar.
• Nuestra evaluación específica sobre modelos de secuencias largas muestra los beneficios de nuestros modelos 8K-seq sobre los modelos 2K- y 4K-seq.
• XGen-7B archiva resultados igualmente sólidos tanto en texto (por ejemplo, MMLU, QA) como en tareas de código (HumanEval).
• Coste de formación de 150.000 dólares en tokens de 1T según los precios de Google Cloud para TPU-v4.

Por escrito: https://arxiv.org/abs/2309.03450
Codebase: https://github.com/salesforce/xGen
Model Checkpoint: https://huggingface.co/Salesforce/xgen-7b-8k-base

Por qué XGen-7B con 8K de longitud de secuencia

A medida que los LLM se hacen omnipresentes, sus aplicaciones a secuencias largas han sido un punto clave, especialmente para aplicaciones como el resumen de texto (potencialmente intercalado con otras fuentes de datos como tablas e imágenes), la escritura de código y la predicción de secuencias de proteínas, que requieren que el modelo considere eficazmente las dependencias estructurales de larga distancia. Un contexto amplio permite a un LLM preentrenado examinar datos de clientes (por ejemplo, documentos que el LLM no utilizó en el entrenamiento) y responde a consultas útiles de búsqueda de información.

Aún así, la mayoría de los LLM de código abierto (por ejemplo, LLaMA, MPT, Falcon) han sido entrenados con un máximo de 2K tokens de longitud de secuencia, lo que supone una limitación clave en el modelado de secuencias largas. Aún no se han evaluado soluciones de tiempo de inferencia como ALiBi para modelos más grandes (por ejemplo, MPT-7b-StoryWriter-65k+). Trabajos recientes sobre el escalado de modelos han demostrado que, para un presupuesto de computación determinado, los mejores resultados no se obtienen necesariamente con los modelos más grandes, sino con modelos más pequeños entrenados con más datos (medidos por el número de tokens). Por lo general, también se prefiere un modelo más pequeño por la eficacia de la inferencia durante el servicio, incluido el servicio en el dispositivo. En vista de ello, entrenamos una serie de LLM de 7B denominados XGen con atención densa estándar en secuencias de hasta 8K de longitud para un máximo de 1,5T tokens. También afinamos los modelos XGen en datos de instrucción de dominio público, creando sus homólogos afinados por instrucción (XGen-7B-inst).

Datos de pre-entrenamiento

Empleamos una estrategia de entrenamiento en dos etapas, en la que cada etapa utiliza una mezcla de datos diferente.

Primera etapa (1,37T tokens)

Los datos de lenguaje natural son una mezcla de datos disponibles públicamente. Los datos de código son una mezcla del subconjunto de GitHub del conjunto de datos de RedPajama y los datos de código de Apex que recopilamos.

Segunda etapa (110B tokens)

Para apoyar mejor las tareas de generación de código, en la segunda etapa mezclamos más datos de código de Starcoder con los datos de la Etapa 1.

Usamos tiktoken de OpenAI para tokenizar nuestros datos. Añadimos tokens adicionales para espacios en blanco consecutivos y tabuladores, así como los tokens especiales descritos en el artículo de Starcoder.

Detalles del entrenamiento

Los modelos XGen-7b se entrenan con nuestra biblioteca interna JaxFormer, que facilita un entrenamiento eficiente de los LLM bajo un paralelismo tanto de datos como de modelos optimizado para el hardware TPU-v4. La receta de entrenamiento y la arquitectura del modelo siguen el modelo LLaMA, mientras que nosotros realizamos dos exploraciones adicionales. En primer lugar, investigamos la aparición de los denominados «picos de pérdida» [PaLM, loss spikes] durante el entrenamiento, es decir, la pérdida estalla de repente temporalmente mientras se desconoce la causa raíz de estos picos. En segundo lugar, los modelos XGen admiten secuencias de hasta 8.192 tokens (en lugar de los 2.048 habituales), por lo que introducimos el entrenamiento por etapas.

Picos de pérdida

A medida que los modelos se escalan a tamaños mayores, el propio entrenamiento es cada vez más sensible a las inestabilidades, que provocan un rendimiento deficiente del modelo si no se abordan con cuidado. En nuestra exploración, hemos reunido pruebas de varios factores que contribuyen individualmente a un entrenamiento inestable. Estos hallazgos preliminares incluyen «circuitos secuenciales sobre paralelos», «swish-GLU sobre GeLU», «RMS-Norm sobre Layer-norm». Específicamente, los circuitos paralelos ampliamente utilizados, que paralelizan el cómputo de autoatención y feed-forward como se adoptó en [GPT-J, PaLM, CodeGen] pueden afectar a la estabilidad del entrenamiento.

La figura anterior muestra la pérdida en términos de entropía cruzada a lo largo del tiempo siguiendo las conocidas leyes de escalado. Sorprendentemente, el entrenamiento no sufre inestabilidades ni picos de pérdida. Los dos picos de pérdida representados en la figura se esperan cuando se amplía la longitud de la secuencia, por ejemplo de 2k a 4k tokens, ya que el modelo necesita adaptarse a secuencias más largas.

Longitud de la secuencia

El entrenamiento con secuencias más largas es computacionalmente desproporcionadamente costoso ya que la complejidad de la autoatención es cuadrática, es decir, el proceso de entrenamiento es lento. Para mitigar la lentitud del entrenamiento, introducimos el entrenamiento por etapas con una longitud de secuencia creciente. Primero, se observan 800B tokens con una secuencia de 2k tokens, luego 400B tokens con 4k, finalmente, 300B tokens con 8k de longitud.

Verificamos la adaptación a secuencias más largas calculando la perplejidad media en cada posición de token en un conjunto de validación retenido que contiene documentos de 8k de longitud de secuencia o más. Si el modelo aprende a utilizar la secuencia completa, es de esperar que la perplejidad disminuya con la longitud de la secuencia, ya que los tokens anteriores contienen información para el siguiente token previsto. Es decir, para una frase larga, cuanto más contexto se proporcione en forma de palabras anteriores, más fácil será adivinar la palabra siguiente. La figura anterior muestra que XGen, en cada etapa, aprende a utilizar contextos más largos, de hasta 8k de longitud de secuencia.

Resultados en Benchmarks Estándar

(i) MMLU

En primer lugar consideramos el benchmark Measuring Massive Multitask Language Understanding (ver ejemplos aquí), que es más reciente que otros, por lo que podría decirse que es menos susceptible a la contaminación de datos, como se ha informado en estudios recientes (véase la página 32 del documento GPT-4 y una discusión relacionada aquí), y se ha utilizado sistemáticamente como punto de referencia de evaluación. Recientemente, sin embargo, se ha informado de incoherencias en la notificación de las puntuaciones MMLU, lo que dio lugar a clasificaciones erróneas en la tabla de clasificación Open LLM de Hugginface; de hecho, Huggingface tuvo que escribir más tarde un blog para aclararlo. En nuestro trabajo, seguimos el estándar MMLU original, que es coherente con los resultados publicados (es decir, en LLaMA).

Resultados del aprendizaje en contexto de 5 disparos MMLU: En primer lugar, mostramos los resultados en el escenario de evaluación original (y recomendado) de 5 disparos, en el que el LLM recibe 5 demostraciones. XGen consigue los mejores resultados en la mayoría de las categorías, también en media ponderada.

Resultados de 0 disparos de la MMMLU: En MMLU tiro cero, de forma similar vemos buenos resultados aunque la diferencia con LLaMA es generalmente menor aquí.

(ii) General Zero-shot Results

A continuación, presentamos los resultados generales de zero-shot en tareas generales de PLN que implican razonamiento de sentido común y GC.

(iii) Resultados en la generación de código

Para evaluar la capacidad de generación de código de XGen a partir de instrucciones en lenguaje natural (docstrings), lo evaluamos en el conocido benchmark HumanEval. Fijamos la temperatura de muestreo en 0,2, p en 0,95 (para muestreo top-p), y num_samples_per_task (n) en 200. Informamos de los resultados estándar de cero disparos con la métrica pass@1.

Resultados en tareas de generación de secuencias largas

Para seguir evaluando nuestro modelo XGen-7b 8k en comparación con las líneas de base que se limitan a entradas de 2k, pasamos a la generación de diálogos largos, el resumen de textos y la garantía de calidad. Todas estas tareas se benefician del procesamiento y la comprensión de un contexto largo para generar una respuesta correcta. Hay que tener en cuenta que, para estas tareas, la mayoría de los modelos base preentrenados no consiguieron generar una respuesta plausible debido a la dificultad de la tarea. Por tanto, utilizamos modelos ajustados a las instrucciones.

Diálogo

Para evaluar las capacidades de comprensión y resumen de diálogos largos, presentamos los resultados de tres tareas de resumen de diálogos: AMI, ForeverDreaming (FD) y TVMegaSite (TMS). La longitud media de las fuentes de estos conjuntos de datos es de aproximadamente 5.570, 6.466 y 7.653, respectivamente. Evaluamos específicamente muestras de menos de 8K de longitud utilizando varios modelos ajustados a las instrucciones. En particular, cuando no se aplicó el truncamiento de la entrada, tanto MPT-7b-inst como Alpaca-inst no obtuvieron buenos resultados en este escenario. Nuestro modelo (XGen-7B-inst) obtuvo las puntuaciones ROUGE más altas en todas las métricas.

Control de calidad de forma larga

A continuación, evaluamos nuestro XGen-7b-inst en una tarea de control de calidad de forma larga que hemos diseñado internamente. Pedimos a ChatGPT que genere preguntas a partir de (a) documentos largos de Wikipedia que abarcan cuatro dominios: Física, Ingeniería, Historia y Entretenimiento, y (b) resúmenes de estos documentos. A continuación, consultamos los LLM para generar respuestas a estas preguntas. Las respuestas suelen tener hasta 256 tokens. Utilizamos GPT-4 para evaluar la calidad de las respuestas en términos de coherencia (estructura y organización) y relevancia (relevancia de la respuesta generada para la pregunta y el documento contextual) en una escala de 0 a 3. A partir de los resultados que se muestran a continuación, vemos que nuestro modelo obtiene puntuaciones más altas en diferentes aspectos en comparación con las líneas de base consideradas.

Resumen

Aquí evaluamos nuestro modelo en dos conjuntos de datos de resumen de texto incluidos en el SCROLLS Benchmark, a saber, QMSum y GovReport. Cubren dos ámbitos diferentes: conversaciones en reuniones e informes gubernamentales. Además, los datos de QMSum incluyen consultas específicas en lenguaje natural que indican al modelo los aspectos clave del documento fuente que deben incluirse en el resumen. Vemos que nuestro modelo XGen-7b supera a otras líneas de base en estas tareas.

Como vemos resultados alentadores de nuestros modelos XGen-7b en estas tareas de secuencias largas, nos gustaría señalar que como estos modelos no están entrenados con los mismos datos de instrucción, no son estrictamente comparables.

Nota sobre Riesgos Potenciales

Por último, a pesar de nuestro esfuerzo por abordar los riesgos de sesgo, toxicidad y alucinaciones tanto en las etapas de pre-entrenamiento como de ajuste fino, al igual que otros LLMs, los modelos XGen-7b no están libres de dichas limitaciones. Esperamos que nuestra base de código de código abierto ayude a otros investigadores a comprender mejor estos retos y a mejorar estas limitaciones clave para que la IA sea beneficiosa para todos.