Compresión y poda estática

CATEGORÍA

Rubén Rodríguez Abril

La ejecución de modelos de lenguaje requiere de enormes recursos computacionales, centralizando su avance en grandes estados y corporaciones. Para reducir esta dependencia, se emplean técnicas de compresión como poda estática (eliminación de pesos o capas), optimización aritmética (cuantización y funciones eficientes) y compresión de dimensionalidad. Estas métodos permiten modelos más ligeros y eficientes, manteniendo su rendimiento en dispositivos pequeños.

Introducción

En la última década, los requerimientos computacionales necesarios para entrenar a un modelo de lenguaje han crecido de tal manera que los modelos contemporáneos sólo pueden ser ejecutados, tanto en entrenamiento como en inferencia, en enormes centros de datos, dotados de una extraordinaria capacidad de procesamiento matemático y de memoria. Esta tendencia ha llevado a una centralización del desarrollo de la IA, concentrando su avance en manos de de hiperpotencias o de grandes corporaciones, las únicas entidades con capacidad económica suficiente para afrontar las cada vez mayores exigencias de hardware.

Para mitigar esta dependencia de infraestructuras masivas y permitir que los modelos de lenguaje sean ejecutables en pequeños dispositivos, se han ideado técnicas de compresión de modelos, que disminuyen su tamaño permitiendo el ahorro de memoria, de poder aritmético/computacional o de ambas cosas a la vez. En este artículo analizaremos varias de ellas, y en particular las técnicas de poda estática, optimización aritmética y compresión propiamente dicha.

Poda estática

La poda (pruning) consiste en la eliminación de determinadas partes del modelo, ya sean pesos sinápticos individuales, valores de activación o inlcuso capas enteras. Es una técnica que ha ganado recientemente mucha atención en relación con la hipótesis del billete de lotería en redes neuronales Frankle y Carbin, 2019), que afirma que a menudo las subredes de un modelo pueden ejecutar una tarea lingüística con la misma eficiencia que la red al completo.

La poda puede realizarse en el momento del entrenamiento (poda estática) o, con posterioridad, en la fase de inferencia (poda dinámica). Esta última será objeto del próximo artículo de esta serie.

Siguiendo el esquema del influyente artículo “Compressing Large-Scale Transformer-Based Models: A Case Study on BERT”, la poda puede clasificarse en estructurada y no estructurada, dependiendo de si afecta a pesos sinápticos aislados o a capas o subcapas enteras.

Poda no estructurada

La poda no estructurada consiste en la eliminación de pesos sinápticos individuales. Es la más sencilla de implementar, aunque la menos precisa en términos de rendimiento. El artículo seminal “Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding” propuso el siguiente procedimiento de poda basada en umbral, estructurado en tres fases:

-El entrenamiento del modelo comienza de manera convencional.

-A partir de cierto momento, los pesos sinápticos que no alcancen un determinado umbral mínimo son anulados, y quedan fuera de cálculos posteriores del entrenamiento.

-Las conexiones asociadas a estos pesos quedan excluidas tanto del flujo de datos de entrenamiento hacia la salida como de la retropropagación del gradiente.

El uso de este sistema logró una reducción masiva del número de de parámetros en modelos de visión por ordenador como AlexNet o VGG-16 con factores de compresión de 9x y de 13x respectivamente, sin apenas merma de su eficacia.

Entre otras variantes de poda no estructurada que refinan el proceso de eliminación de pesos sinápticos podemos citar las siguientes:

–Poda basada en movimiento. Elimina aquellos pesos sinápticos que con mayor velocidad tiendan a cero durante el entrenamiento.

–Poda proximal con reponderación. Usa un procedimiento iterativo consistente en fomentar la dispersión de los pesos y tras ello eliminar aquellos cuya magnitud se acerque a cero. Este método es más efectivo que el anterior, pero tiene un mayor costo computacional.

La poda no estructurada es útil en modelos como BERT, en el que abundan las capas densamente conectadas. Sin embargo, para que la poda efectivamente aligere los requisitos computacionales del modelo es necesario disponer de circuitos capaces de aprovechar la dispersión de las matrices (que mide la proporción de sus coeficientes que son iguales a cero), así como mecanismos que permitan una eficiente almacenamiento de las mismas. Tal es el caso de los Sparse Tensor Cores de las arquitecturas de NVIDIA posteriores a Ampère o de las MPUs (Matrix Processing Units) presentes en los modelos de Cerebras o Graphcore.

Poda estructurada

Es la eliminación de bloques estructurados de pesos o de partes enteras de un modelo de lenguaje. En el marco del modelo BERT, los autores del estudio Ganesh et al (2019) mencionaron los siguientes tipos de poda:

a) Poda de cabezas de atención. Supone la eliminación completa de algunas cabezas de un mecanismo de atención. Algunos experimentos sugieren que reteniendo apenas el 60% de las cabezas es posible preservar un alto rendimiento en BERT.

b) Poda de capas, que como su propio nombre indica implica el descarte de capas enteras. Algunas técnicas recurren a apagar aleatoriamente una parte de las mismas durante el entrenamiento, con el objeto de que en fase de inferencia pueda ser utilizado un submodelo de la longitud deseada. En los modelos residuales por lo general son descartadas aquellas capas con menores niveles de activación, ya que en este caso la información tiende a circular a través del atajo en que consiste la conexión residual.

c) Poda de embeddings, que elimina las dimensiones menos significativas y con menor carga semántica de los vectores de codificación de contenido.

Figura 1. Diferentes tipo de poda según Ganesh et al. A la izquierda se representan las podas de capas (a) y de embeddings (b). En el centro, la poda de cabezas de atención (c). A la derecha, la poda no estructurada, que afecta a capas densamente conectadas (d). Obsérvese la presencia de conexiones residuales, típicas del modelo BERT.

Optimización aritmética

La optimización aritmética abarca el conjunto de técnicas destinadas a simplificar de las operaciones aritméticas a realizar dentro de un modelo, tanto durante su entrenamiento como sobre todo en fase de inferencia. Pueden consistir en la modificación del formato numérico de pesos y activaciones, en la alteración de las funciones de activación, para hacerlas computacionalmente menos costosas y en la utilización de unidades de procesamiento que permitan aprovechar la esparsión de las matrices durante su multiplicación.

Modificación de formatos numéricos

La modificación de formatos numéricos consiste en la transformación de la representación de los pesos sinápticos, acortando, por lo general, su longitud en bits. Un caso paradigmático es la conversión valores en coma flotante (p.e. FP32) a números enteros de menor longitud en bits (p.e. INT 8), con el consiguiente ahorro tanto en memoria como en poder computacional. Si el resultado de la transformación es un número entero se habla de cuantización. Dependiendo del método aplicado, la cuantización puede ser de tres tipos: lineal, no lineal y entrenada.

Figura 2. Las ventajas de la optimización numérica. Los “Tensor Cores” (las partes de la unidad aritmético-lógica encargadas de realizar multiplicaciones matriciales) de la arquitectura Turing de NVIDIA trabajan 4 veces más rápido cuando se trata de números INT4 que en el caso de números FP16.. Fuente: NVIDIA.

La cuantización lineal es un simple reescalado del rango de los pesos sinápticos de una capa. Por ejemplo, si los valores INT8 oscilan entre 0,56 a 246,23, la amplitud de su rango sería 245,67. Para convertir estos valores a INT4, el reescalado se realiza mediante las ecuaciones siguientes:

rango = max-min

nuevo_valor = redondeo[255*[(valor_anterior-min)/rango]]

donde max y min son los valores máximo y mínimo de pesos sinápticos en el formato anterior.

La cuantización no lineal utiliza, como su propio nombre indica, una función no lineal para realizar el escalado. Puede tratarse de una función logarítmica o de la propia función logística. Esta última presenta la ventaja de que comprime los valores extremos de una manera más controlada.

Figura 3. En la imagen se representa la función logística cuyo rango ha sido reescalado a -8 a 7, que son los valores mínimo y máximo que puede tener un número INT4.

En la cuantización por agrupamiento, que es la más eficiente de todas, los pesos sinápticos se agrupan en clústers correspondientes a los posibles valores INT4 (16 en total). El procedimiento comienza con la selección aleatoria de 16 centroides (valores INT8), que formarán el núcleo de los clústers. Cada peso sináptico se asigna al clúster cuyo centroide esté más cercano. Tras ello, se recalculan los centroides haciendo la media de los pesos de su respectivo grupo. Finalmente, cada peso sináptico es sustituido por su centroide.

Si un centroide queda sin elementos puede reinicializarse en una nueva posición o bien combinarse con su vecino más cercano, distribuyéndose entre ellos los pesos.

Optimización de funciones de activación

Las funciones de activación que suelen aplicarse tras cada capa lineal en una red neuronal pueden reemplazarse por alternativas menos costosas, como por ejemplo ReLU6 en lugar de Swish o GeLU, que tan utilizadas son en el ámbito de los transformers.

Así mismo, su cálculo puede ser aligerado mediante los siguientes métodos:

–Tablas de consulta (look-up tables LUTs) en las que se almacenan valores precomputados a los que se accede durante el entrenamiento o la inferencia.

–Aproximaciones polinómicas para funciones como softmax o la función logística, entre otras.

-Empleo de circuitaje específico dentro de las GPUs, como los SFUs (special funcion units, SFUs) de las arquitecturas de NVIDIA. Aunque estos últimos se centran, fundamentalmente en el cálculo de funciones trigonométricas, exponenciales, logarítmicas, así como en interpolaciones lineales, no es descabellado que en el futuro se creen circuitos capaces de procesar funciones GeLU o softmax, que tanto uso tienen en el ámbito del procesamiento natural del lenguaje y en el del aprendizaje profundo en general.

Figura 4. Ejecución un un warp (conjunto de hilos) dentro de una GPU de NVIDIA. Las operaciones matriciales son realizadas por los núcleos CUDA (en verde), mientras que las funciones especiales son ejecutadas por las SFUs (en azul).

Compresión propiamente dicha

En esta sección quedan abarcadas técnicas que reducen el tamaño de los pesos sinápticos de las capas. En particular, en las capas densamente conectadas se aplica un mecanismo de reducción de dimensionalidad (por ejemplo, PCA) tanto a la propia capa como a los vectores de entrada.

De este modo, los cálculos propios de la capa tienen lugar en un espacio latente de menores dimensiones, con el consiguiente alivio para las áreas de cálculo de las GPUs. Tras la realización de los cálculos, el vector de salida experimenta una descompresión que restaura de nuevo sus dimensiones originales.

Existen dos enfoques principales para la compresión en redes neuronales:

–Compresión para almacenamiento sin impacto en el cálculo. La compresión se usa tan sólo para almacenar los pesos, pero las operaciones aritméticas se realizan sobre los datos descomprimidos. En este enfoque se incluyen los mecanismos de compresión basada en entropía, como la codificación Huffmann o la codificación aritmética.

–Compresión con operaciones en el espacio latente. La reducción de dimensionalidad no sólo se aplica a los pesos almacenados, sino que también afecta a la ejecución del modelo. Los cálculos de las capas tienen lugar en un espacio latente de menores dimensiones, con el consiguiente alivio para las áreas de cálculo de las GPUs. Una vez completados, el vector de salida experimenta una descompresión que restaura de nuevo la información a sus dimensiones originales. La realización de operaciones en espacios latentes, relativamente común en modelos de difusión, también ha encontrado un cierto uso en los modelos de lenguaje, como veremos en la sección siguiente.

Mecanismos de atención

Dada su importancia fundamental en el marco de los modelos de lenguaje, en este último apartado haremos referencia a dos grandes técnicas que afectan exclusivamente a los mecanismos de atención: La agrupación y compresión de claves y valores y la poda de cabezas de atención.

Tal y como se ha expuesto en otros artículos de esta misma serie, el transformer original de 2017 introdujo el sistema de la atención multicabeza (Multi-Head Attention), en cuya virtud la matriz de atención se divide en varias secciones (lonchas longitudinales) denominadas cabezas. A cada cabeza se le asigna una parte del hiperespacio semántico en el que se despliegan, como vectores, los tokens. Algunas de estas cabezas se especializan en detectar patrones gramaticales o semánticos, mientras que otras atienden las relaciones de naturaleza posicional. Cada cabeza está dotada de su propia submatriz de atención, que se calcula a partir de tres matrices diferentes, denominadas Q (query, consulta), K (key, clave) y V (value, valor).

En cada cabeza, las citadas matrices toman la forma:

Q_i = X_i^QQ

K_i = X_i^KK

V_i = X_i^VV

donde i es el número de cabeza y X_i^Q, X_i^K, X_i^V, las matrices de proyección, que crean versiones de las matrices consulta, clave y valor para cada cabeza.

Las nuevas técnicas de compresión en mecanismos de atención buscan reducir el número de vectores clave o consulta o bien disminuir las dimensiones de los mismos. Mostramos dos ejemplos de cada uno de los dos enfoques: Grouped Query Attention (GQA) y Multi-Head Latent Attention (MLA).

Agrupación de claves

En la atención multicabeza original hay versiones únicas de las matrices Q, K y V para cada cabeza. Todas ellas deben almacenarse en el caché de la GPU durante la ejecución del modelo, con la consiguiente sobrecarga de memoria. Por el contrario, en la atención multiconsulta (Multi-query attention), las dos matrices K y V (claves y valores) son compartidas por todas las cabezas. Hay un ahorro considerable de necesidad de almacenamiento, aunque a costa de disminuir las capacidades del modelo. La atención agrupada (Grouped-query attention) representa una solución intermedia entre ambos modelos, tal y como se muestra en la imagen.

Figura 5. En la atención multicabeza original (multi-head), hay versiones de claves y valores para cada cabeza. Todas esas matrices deben almacenarse en el caché. En la atención multiconsulta (multi-query), por el contrario, sólo hay dos únicas matrices, por medio de las cuales se calculan claves y valores en todas. La atención agrupada representa una solución intermedia. Fuente: DeepSeek V2.

Compresión KV

Los ingenieros que diseñaron DeepSeekV2 optaron por otra estrategia de optimización de almacenamiento consistente en comprimir claves y matrices en una única matriz latente KV. Otra matriz de descompresión/proyección era utilizada para obtener de nuevo claves y valores a partir de la información comprimido. Este procedimiento fue denominado atención latente multicabeza (Multi-Head Latent Attention, MLA). Las matrices de compresión y descompresión pueden ser incorporadas a las de atención y de embedding, con el consiguiente ahorro de gasto computacional.

Poda de cabezas de atención

La eliminación selectiva de algunas de las cabezas de una unidad de atención fue propuesta por Michel et al (2019). Su estudio se llevó a cabo en dos arquitecturas diferentes:

-El primero de ellos fue WMT, una modelo de traducción basado en la arquitectura del transformer primitivo del año 2017, entrenado en el corpus WMT2014 (inglés-francés) y evaluado sobre la métrica BLEU. Se analizaron configuraciones de sólo-codificador, sólo-descodificador y codificador-descodificador.

-El segundo fue BERT, que fue entrenado para realizar diversas tareas lingüísticas, evaluado sobre la métrica de la precisión (accuracy).

En ambos casos, la eliminación de cabezas apenas mostró un impacto mínimo rendimiento del sistema hasta un umbral del 60-80%. A partir de entonces la puntuación BLEU colapsó, aproximándose rápidamente a cero, como se puede ver en la imagen.

Figura 6. Rendimiento de las tres configuraciones de WMT. Dos de ellas podían seguir funcionando razonablemente bien incluso con una poda del 80% de sus cabezas de tención. Fuente: Michel et al (2019).

Se evaluó también el impacto en el rendimiento de los modelos al eliminar cabezas dentro de una única capa individual. En este caso, la eliminación de varias cabezas no afectó al rendimiento de los modelos, e incluso, en determinados casos, parecía mejorarla. Si bien en términos generales no hubo merma significativa en el rendimiento, en algunas capas críticas, como la última del codificador-descodificador WMT, sí que hubo una caída significativa en la puntuación BLEU. Por lo demás, las cabezas más importantes tendían a ser consistentes entre las diferentes bases de datos, lo que sugiere que algunas de ellas ejercen un papel clave dentro del modelo.