PanGu, la propuesta de Huawei

Rubén Rodríguez Abril

 Pangu, el gigante primordial de la mitología china, da su nombre a la serie de modelos de lenguaje de Huawei, cuyo ejemplar más avanzado, PanGu-Σ, es el más grande de los modelos no multimodales existentes hasta la fecha.

Durante la década de los 2010, el desarrollo de los grandes modelos de lenguaje fue liderado por compañías estadounidenses como Open AI o Google, quienes desarrollaron innovaciones clave como los algoritmos seq2seq o la arquitectura de transformers, entre muchas otras. En este tiempo, los esfuerzos de China se enfocaban en los algoritmos de visión artificial y en el uso industrial de la IA. Para paliar su retraso en procesamiento de lenguaje natural, en el año 2021 Huawei inauguró una serie de modelos de lenguaje preentrenados a la que llamó PanGu, en honor al gigante primordial de la mitología china, de cuyo cuerpo emergió todo el Cosmos.

En este artículo analizaremos dos de los modelos de la serie: PanGu-α y PanGu-Σ.

PanGu-α

PanGu-α, es un modelo de lenguaje autorregresivo, preentrenado sobre el lenguaje chino y basado en la predicción del siguiente token.

Se trata de un codificador de transformer, muy similar a la arquitectura original de 2017, pero con una leve modificación: a las capas de transformers se les añade una nueva, llamada de consulta (query layer), que es como las capas anteriores pero ligeramente modificada.

Como el lector recordará de artículos anteriores, cada capa de transfomers se componía de una unidad de atención y una mini-red orientada hacia adelante. Dentro de la unidad de atención, la matriz de atención A se calcula multiplicando la matriz consulta (query) Q por la matriz clave (key) K:

Ah = QhVh

Qh = Hh-1Wq

Kh = Hh-1Wk

donde h es la capa actual, Hh-1 las activaciones de la capa anterior, y Wq y Wk matrices de pesos sinápticos aprendidos.
En las capas de consulta (query layers) de PanGu el cálculo de la matriz consulta es levemente modificado:

Qh = phWq

donde pn es un vector que expresa la codificación posicional n. El lector recordará que en la capa inicial de cualquier transformers se producen dos codificaciones iniciales de la matriz de entrada: la codificación de contenido (embedding) y la codificación de posición. pn es la fila de la matriz de posición que corresponde al token n, tal y como se muestra en la imagen.

Figura 1. Diagrama de flujo de información entre las diferentes capas de PanGu-α, cuando éste está imprimiendo el token 6. Abajo del todo se representa la capa de embedding y codificación de posición, cuyo resultado, sumado, es transmitido a los módulos de transformer (en verde). El rectángulo azul representa la capa de consulta, que recibe como entrada al vector con la codificación posicional del token 6.

La sustitución Hh-1 por ph produce una reducción de la dimensionalidad. La salida de la capa de consulta ya no es una matriz sino un vector con el que se alimenta la última capa del modelo, consistente en una capa lineal (cuyos pesos y sesgos son aprendidos) y una función softmax.

En total, se entrenaron tres versiones del modelo con diferente magnitud de parámetros, PanGu-α 2,6B, PanGu-α 13B y PanGu-α 200B, sobre un corpus textual chino de 1,1TB.

Figura 2. Diferentes modelos de PanGu-α. Las columnas describen el número de parámetros, de capas, el tamaño del modelo, el de la red orientada hacia adelante así como el número de cabezas de las unidades de atención.

En el momento en que se diseñó el modelo incial de PanGu-α, existían tres grandes corpus digitales de texto chino: CLUECorpus2020, extraído de Common Crawl y de 100GB de tamaño; la base de datos del modelo multimodal M6 (con 300 GB de texto); y WuDaoCorpus6 (con 300 GB de datos de texto). En total, su tamaño era insuficiente para entrenar las diferentes versiones de PanGu, particularmente la de 200B parámetros. Por ello, los investigadores de Huawei decidieron crear una nueva base de datos de texto chinos del tamaño de 1,1TB, siguiendo los pasos establecidos en la imagen.

Figura 3: Fases de la construcción de la base de datos de PanGu-α. Las fuentes de datos fueron los “crawls” de Google y Baidu, enciclopedias como Baidu Baike o Soigu Baike, e-books en lengua china y noticias desde 1992 a 2011, entre muchas otras. Todas ellas fueron sometidas a un proceso de “desduplicación”, con el objeto de eliminar los textos repetidos, así como a repetidos procesos de limpieza y filtrado, tendentes a eliminar los textos de contenido dañino o de baja calidad.

El framework MindSpore

Dado que en la mayoría de arquitecturas de Aprendizaje Profundo los parámetros vienen constituidos por números de 16 o 32 bits, para almacenar los parámetros de PanGu-α 200B sería necesarios unos 750 GB, lo cual está fuera de las capacidades de memoria de cualquier chip singular actual. Para ejecutar en fase de inferencia y (aún más) para entrenar a un modelo es necesario emplear centenares o miles de aceleradores. En el caso de PanGu-α 200B fueron utilizados 2048 chips Ascend 910 de Huawei coordinados a través del framework de MindSpore.

El módulo Auto-parallel de MindSpore distribuye, como su propio nombre indica, automáticamente la carga de trabajo entre los diferentes chips de un centro de datos, sin que el diseñador del modelo de lenguaje tenga que penetrar en estos detalles.
La paralelización puede realizarse de diferentes modos

Paralelización de datos (data parallelism): El lote (batch) se divide en varias partes y cada chip es encargado de realizar los cómputos relativos a cada una de ellas.

Paralelización de modelo (model parallelism), que implica que cada chip se encarga de gestionar una parte diferente del modelo. Se subdivide a su vez en paralelismo de operandos (op-level paralelism), que divide a cada capa en varias secciones que se asignan a diferentes dispositivos; de procedimiento escalonado (pipeline model parallelism), en el que cada capa es gestionada por un chip; y de optimización (optimizer model parallelism), en el que el cómputo de los gradientes se hace de un modo distribuido.

Rematerialización (rematerialization). Los valores de activación son necesarios para repropagar el gradiente de la función de pérdida. Sin embargo, su almacenamiento ocupa mucho espacio. Por este motivo, en aquellas capas cuyo cálculo no sea computacionalmente costoso (p.e. ReLU), los valores de activación son borrados tras el paso de la información a la siguiente capa y recalculados en el momento de la retropropagación.

PanGu-Σ

La versión más avanzada de la serie PanGu, PanGu-Σ, es el más grande modelo de lenguaje (en términos de número de parámetros) existente hasta la fecha.

Las principales diferencias con PanGu-α son las siguientes:

-En PanGu-Σ no hay una única capa de codificación de contenido (embedding), sino varias. Y cada una de ellas se encarga de un dominio diferente, como textos en chino, inglés, polaco, fórmulas matemáticas o código informático.

Figura 4. Diferentes dominios de PanGu-Σ. A cada uno de ellos se asigna un conjunto de expertos.

-Por lo que se refiere al cuerpo principal del transformer, en PanGu-Σ se divide en dos grandes secciones. La primera sección está formada por M capas ordinarias de transformers. La segunda, por N capas de sistemas de expertos. El modelo puede funcionar en tres modos: modo mixto (mixed mode), cuando M y N son diferentes a 0, modo disperso (sparse mode), cuando M es nulo, o modo denso (dense mode), cuando N es nulo. En este último caso, el funcionamiento es similar al de PanGu-α.

-Las capas de expertos se denominan Random Routed Experts (RRE). La unidad de atención es común a todas estas capas. Los expertos sustituyen a la red orientada hacia adelante Como se muestra en la imagen, los expertos se agrupan por dominios. Y el enrutamiento se realiza en una doble fase. La primera fase es por dominio. La segunda, dentro de cada dominio, es aleatorio. Los enrutadores no aprenden durante el entrenamiento.

Figura 5. Estructura de PanGu-Σ. Cada color representa un dominio diferente (matemáticas, textos en chino, textos bilingües, etc…). En las capas de expertos, la red orientada hacia adelante (feed-forward) es sustituida por una red de expertos. Del modelo general se extrae el submodelo que se ejecutará en fase de inferencia.

Para ahorrar poder computacional los investigadores de Huawei elaboraron un mecanismo de herencia entre las diferentes versiones del modelo: los pesos sinápticos de PanGu-Σ no se inicializan aleatoriamente sino que son heredados de PanGu-α, en el modo esquematizado en la ilustración.

Figura 6. La inicialización de los pesos sinápticos de PanGu-Σ se hace a partir de los PanGu-α. Así, las diferentes capas de “embedding” del primero se inicializan a partir de la única capa del segundo. Lo mismo hacen los sistemas de expertos a partir de la capa feed-forward.

PanGu-Σ fue entrenado con nuevas bases de datos de texto, no empleadas para versiones anteriores del modelo, como WuDaoCorpora 2.0 y ClueCORPUS2020, que suman más de 300 GB de datos en chino. En lengua inglesa fueron empleadas The Pile y C4, con un total de 1550 GB en tokens. Y finalmente, en el ámbito del código se usaron 147 GB de código Python y 161 GB de código Java.

Evaluación

PanGu-Σ fue sometido a diversas pruebas de comprensión lectora, inferencia de lenguaje natural, clasificación textual, similaridad semántica, desafíos de esquemas de Winograd, preguntas Cloze (completar frases), traducción automática o generación de código. El modelo mostró una abierta progresión en tareas zero-shot (es decir, sin haber visto ejemplos en el entrenamiento) respecto a otras arquitecturas chinas como ERNIE 3.0 Titan.

A pesar de todo, los investigadores de Huawei consideraron que aun existen ciertos retos en el ámbito de los modelos de lenguaje que deberán ser afrontados en el futuro, como la profundización en el principio de esparsión, la necesidad de que el modelo de lenguaje reciba retroalimentación directamente del mundo físico, su integración en sistemas multimodales, la necesidad de implementar modelos ligeros que no requieran demasiado poder computacional o necesidad de almacenamiento en fase de inferencia o la adaptación online del conocimiento de los mismos.

Lecturas Recomendadas

– PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation.

– PanGu-Σ: Towards Trillion Parameter Language Model with Sparse Heterogeneous Computing.