Pruning

La poda es una técnica de optimización de modelos que elimina pesos, neuronas, filtros, cabezales de atención o capas completas que apenas contribuyen a la salida de una red neuronal entrenada, dando lugar a un modelo más pequeño que necesita menos memoria, capacidad de cálculo y energía. En aprendizaje profundo, el término es sinónimo de poda de modelos. La mayoría de las redes están muy sobredimensionadas en parámetros: Han et al. redujeron AlexNet de 61 millones a 6,7 millones de parámetros y VGG-16 de 138 millones a 10,3 millones en ImageNet, sin pérdida de precisión, mediante entrenamiento, eliminación de conexiones de baja magnitud y reentrenamiento de las restantes. La poda es un paso habitual antes del despliegue de IA en el edge, donde los presupuestos de memoria, potencia y latencia de inferencia son ajustados.

Cómo funciona la poda#

Cada método de poda decide qué parámetros eliminar y cuándo hacerlo. La puntuación de importancia suele basarse en la magnitud del peso, ya que los pesos cercanos a cero son los que menos modifican la salida, mientras que las puntuaciones basadas en gradientes estiman de forma más directa el efecto de cada parámetro en la función de pérdida. Eliminar de golpe una gran proporción suele costar más precisión que aplicar una poda iterativa: se elimina una pequeña proporción, se ajustan los pesos restantes para recuperarse y se repite el proceso hasta alcanzar el objetivo de tamaño o velocidad. La hipótesis del billete de lotería añade que las redes densas contienen subredes dispersas que, al restablecerse a su inicialización original, se entrenan por sí solas y alcanzan una precisión comparable a la de la red completa.

Poda no estructurada, estructurada y semiestructurada#

La poda no estructurada pone a cero pesos individuales en cualquier posición de un tensor. Alcanza la mayor dispersión posible para una precisión determinada, pero produce una matriz dispersa irregular que los kernels densos de GPU y CPU siguen procesando con el coste completo, salvo que un entorno de ejecución disperso omita los ceros. La poda estructurada elimina filtros, canales, cabezales de atención o capas enteros de una red neuronal convolucional o un transformer, de modo que cada capa podada se convierte en una capa densa más pequeña que funciona más rápido en cualquier hardware, a costa de perder más precisión por parámetro eliminado. Las utilidades de poda de PyTorch aplican ambos tipos como máscaras sobre una capa existente:

import copy

from torch import nn
from torch.nn.utils import prune

# Two copies of one convolution layer with 32 filters of shape 16x3x3
unstructured = nn.Conv2d(16, 32, kernel_size=3)
structured = copy.deepcopy(unstructured)

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

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

for name, layer in [("Unstructured", unstructured), ("Structured", structured)]:
    zeros = (layer.weight == 0).float().mean()
    empty_filters = (layer.weight.flatten(1).abs().sum(dim=1) == 0).sum()
    print(f"{name}: {zeros:.0%} of weights zeroed, {empty_filters}/32 filters entirely empty")

Ambas capas pierden el 50 % de sus pesos, pero solo la capa estructurada acaba con 16 de sus 32 filtros completamente vacíos, que se pueden eliminar físicamente para dejar una convolución densa de 16 filtros.

La poda semiestructurada se sitúa entre ambas. El patrón de dispersión 2:4 de NVIDIA pone a cero dos de cada cuatro pesos consecutivos; los Sparse Tensor Cores de Ampere y las GPU posteriores omiten esos pesos para realizar el mismo cálculo matricial en la mitad de tiempo. NVIDIA midió mejoras de hasta un 20 % de extremo a extremo para ResNeXt-101 disperso con TensorRT en una A100. En hardware denso, los ceros dispersos por sí solos no aportan velocidad: podar Ultralytics YOLOv5x hasta alcanzar un 30 % de dispersión global no estructurada con su utilidad prune(), sin ajuste fino, redujo el COCO mAP_50-95 de 0.507 a 0.489 y no modificó el tiempo de inferencia en una V100, que se mantuvo en 5,2 ms por imagen.

Aplicaciones en el mundo real#

La poda es más importante cuando un modelo debe ejecutarse en hardware con recursos limitados. Las redes de detección de objetos podadas por canales se ajustan a los presupuestos de memoria y potencia de drones, cámaras inteligentes y teléfonos, lo que permite la inferencia en tiempo real sin conexión a la nube y mantiene las imágenes en el dispositivo para proteger la privacidad de los datos. En la inspección de fabricación, la misma compresión permite ejecutar un detector de defectos en un módulo edge de bajo coste junto a la línea de producción.

También se podan los modelos de lenguaje grandes. SparseGPT poda OPT-175B y BLOOM-176B hasta alcanzar una dispersión no estructurada del 60 % en una sola operación, sin reentrenamiento y con un aumento insignificante de la perplejidad. El trabajo Minitron de NVIDIA podó Llama 3.1 8B para convertirlo en un modelo 4B mediante la eliminación de capas o dimensiones ocultas y, después, recuperó calidad mediante la destilación de conocimiento del modelo original.

La poda frente a técnicas relacionadas#

La poda suele combinarse con otros métodos de compresión, y el término también tiene un significado anterior ajeno a las redes neuronales.

• Cuantización de modelos: La cuantización conserva todos los pesos y reduce su precisión numérica, por ejemplo, pasando de números en coma flotante de 32 bits a enteros de 8 bits. A menudo se cuantizan los modelos podados para ahorrar aún más recursos.
• Destilación de conocimiento: La destilación entrena un alumno compacto independiente para imitar a un profesor de mayor tamaño y es una forma habitual de recuperar la precisión que pierde un modelo muy podado.
• Capa Dropout: Dropout pone a cero activaciones aleatorias durante un único paso de entrenamiento como técnica de regularización y mantiene intacto el número de parámetros, mientras que la poda elimina parámetros de forma permanente.
• Poda de árboles de decisión: En el aprendizaje automático clásico, la poda elimina ramas de un árbol de decisión ajustado para reducir el sobreajuste, como ocurre con la poda por complejidad de coste de scikit-learn, que elimina los subárboles más débiles a medida que aumenta su parámetro ccp_alpha.