Model Pruning

La poda de modelos elimina parámetros que apenas contribuyen a la salida de una red neuronal entrenada, desde pesos individuales hasta filtros, cabezas de atención o capas completas, para que el modelo sea más pequeño y barato de ejecutar. La mayoría de las redes tienen demasiados parámetros, por lo que se puede eliminar una gran parte de sus pesos sin apenas perder precisión, sobre todo si después se ajusta el modelo. Un método de poda y reentrenamiento de 2015 redujo AlexNet de 61 millones a 6,7 millones de parámetros y VGG-16 de 138 millones a 10,3 millones sin perder precisión en ImageNet. La poda es un paso habitual para preparar modelos para teléfonos, cámaras y otros equipos de IA en el borde, y ahora también se aplica a los modelos de lenguaje grandes (LLM).

Cómo funciona la poda de modelos#

La poda puntúa cada peso o estructura candidato según su importancia, elimina la fracción con menor puntuación y ajusta los pesos restantes del modelo para recuperar precisión. El criterio más sencillo y aún más habitual es la magnitud: los pesos cercanos a cero o los filtros con una norma L1 o L2 pequeña son los que menos alteran la salida al eliminarlos. Otros criterios usan gradientes o activaciones para estimar cuánto aumentaría la pérdida al retirar cada elemento. Eliminar de golpe una fracción grande suele perjudicar la precisión, por lo que la poda normalmente se realiza por rondas, cada una seguida de un ajuste fino, hasta alcanzar el tamaño objetivo del modelo o hasta que la precisión caiga por debajo del umbral aceptable:

La poda también cambió la forma en que los investigadores conciben el tamaño de las redes. La hipótesis del billete de lotería propone que una red densa contiene subredes pequeñas que, entrenadas de forma aislada desde sus pesos iniciales originales, alcanzan la precisión de la red completa; sus autores encontraron sistemáticamente subredes de menos del 10-20 % del tamaño original en MNIST y CIFAR-10.

Poda no estructurada frente a estructurada#

La poda no estructurada pone a cero pesos individuales en cualquier parte de un tensor. Alcanza una gran dispersión con poca pérdida de precisión, pero los tensores mantienen su forma, por lo que los núcleos densos estándar realizan el mismo trabajo y las aceleraciones requieren núcleos dispersos o compatibilidad de hardware. Las GPU NVIDIA desde la generación Ampere ejecutan uno de estos patrones, la dispersión 2:4 con dos ceros en cada grupo de cuatro pesos consecutivos, a una velocidad de cálculo denso de hasta el doble. La poda estructurada elimina filtros, canales, cabezas o capas completos y deja un modelo denso más pequeño que se ejecuta más rápido en cualquier procesador, aunque con mayor coste en precisión si se elimina el mismo número de pesos. Al eliminar un filtro también se elimina el canal de entrada correspondiente de la capa siguiente; bibliotecas como Torch-Pruning gestionan automáticamente estas dependencias. Las utilidades de poda de PyTorch muestran la diferencia entre ambos métodos:

import torch
from torch.nn.utils import prune

# Two identical convolution layers, one for each pruning style
torch.manual_seed(0)
a = torch.nn.Conv2d(16, 32, kernel_size=3)
b = torch.nn.Conv2d(16, 32, kernel_size=3)

# Unstructured: zero the 50% of individual weights with the smallest magnitude
prune.l1_unstructured(a, name="weight", amount=0.5)

# Structured: zero the 50% of output filters with the smallest L2 norm
prune.ln_structured(b, name="weight", amount=0.5, n=2, dim=0)

for name, layer in [("unstructured", a), ("structured", b)]:
    zeros = (layer.weight == 0).float().mean().item()
    empty = (layer.weight.flatten(1) == 0).all(dim=1).sum().item()
    print(f"{name}: {zeros:.0%} of weights zero, {empty} of 32 filters empty")

Ambas capas acaban con el 50 % de sus pesos a cero. Solo la capa estructurada tiene filtros vacíos —16 de 32— y, al eliminarlos, queda una capa densa con la mitad de filtros, lo que reduce a la mitad sus cálculos en hardware corriente. El hardware diseñado para la dispersión también se beneficia de la poda no estructurada: la integración de CVflow de Ambarella para Ultralytics YOLO, mantenida por Ambarella en su propia bifurcación del paquete, poda los pesos progresivamente durante el entrenamiento de YOLO y mantiene esa dispersión durante la exportación a ONNX para compilar el modelo para sus chips de IA en el borde.

La poda frente a técnicas relacionadas#

• Cuantización de modelos: Reduce la precisión de cada peso, por ejemplo, de números de coma flotante de 32 bits a enteros de 8 bits, y los conserva todos. La poda y la cuantización suelen combinarse: la poda reduce el número de operaciones y la cuantización abarata cada una.
• Destilación de conocimiento: Entrena un modelo estudiante más pequeño para imitar a un modelo profesor más grande y, a menudo, se aplica después de la poda estructurada para recuperar la precisión que se pierde al eliminar filtros o capas completos.
• Dropout: Pone temporalmente a cero activaciones aleatorias durante el entrenamiento como forma de regularización; en inferencia se usa la red completa.

Poda de modelos de lenguaje grandes#

Reentrenar un LLM después de cada ronda de poda cuesta demasiado, así que los métodos de una sola pasada lo podan de una vez usando un pequeño conjunto de calibración. SparseGPT redujo los modelos OPT de 175.000 millones de parámetros y BLOOM de 176.000 millones a una dispersión no estructurada del 60 % en menos de 4,5 horas, con un aumento insignificante de la perplejidad; también admite el patrón 2:4, que las GPU aceleran. La poda estructurada combinada con destilación de conocimiento produce modelos densos más pequeños: Minitron de NVIDIA obtuvo modelos 8B y 4B a partir de un modelo preentrenado 15B mediante la poda de canales de incrustación, neuronas MLP y cabezas de atención; después los reentrenó con hasta 40 veces menos tokens que los necesarios para entrenarlos desde cero.