Two-Stage Regularization-Based Structured Pruning for LLMs

1. El Problema

La implementación de Grandes Modelos de Lenguaje (LLMs) en entornos reales se ve obstaculizada por su inmenso tamaño y el alto costo computacional. La poda estructurada (eliminar capas completas o canales) es una solución prometedora, pero los métodos existentes de poda por capas presentan limitaciones críticas:

• Pérdida de Conocimiento: Los métodos actuales calculan la importancia de las capas, eliminan las "menos importantes" y luego intentan recuperar el rendimiento mediante un reentrenamiento extensivo. Sin embargo, incluso las capas consideradas "poco importantes" pueden contener conocimiento valioso que se pierde irremediablemente al ser eliminadas directamente.
• Alto Costo de Reentrenamiento: La recuperación del rendimiento tras la poda requiere fine-tuning masivo, lo que genera una sobrecarga computacional significativa.

2. Metodología: TRSP (Two-Stage Regularization-Based Structured Pruning)

El autores proponen TRSP, un método novedoso que realiza la poda de capas en dos etapas de regularización antes de eliminar cualquier parámetro. El objetivo es redistribuir el conocimiento de las capas a eliminar hacia las capas preservadas antes de la poda, minimizando así la pérdida de información y evitando el reentrenamiento costoso.

El proceso consta de cuatro fases principales:

A. Preparación de Datos

Se utiliza una cantidad mínima de datos (ej. 128 muestras) extraídas de conjuntos de datos estándar (como WikiText-2, Alpaca, C4) para guiar el proceso de regularización. Esto reduce drásticamente la sobrecarga computacional.

B. Primera Etapa: Aprendizaje de Pesos de Capa (Regularización 1)

• Se asigna un peso aprendible inicial ($S[i]$) a cada capa del transformador.
• Se multiplica la salida de cada capa por su peso correspondiente.
• Se entrena el modelo iterativamente minimizando una función de pérdida que incluye:
  1. La pérdida de modelado del lenguaje ($L(W, X)$).
  2. La norma $\ell_1$ de los pesos de las capas como término de regularización.
• Estrategia Greedy Iterativa: En lugar de identificar todas las capas a eliminar de una sola vez (one-shot), el método identifica y oculta (mascara) la capa con el peso más bajo en cada iteración. Esto es crucial porque la importancia de una capa puede cambiar a medida que se eliminan otras.

C. Segunda Etapa: Transferencia de Conocimiento (Regularización 2)

• Una vez identificadas las capas candidatas a ser podadas (conjunto $P$), se aplica una segunda regularización.
• Se añade un término de pérdida basado en la diferencia entre la entrada ($X_{in}$) y la salida ($X_{out}$) de estas capas específicas.
• Objetivo: Forzar a que la salida de las capas a podar sea lo más similar posible a su entrada (función identidad). Esto obliga al modelo a transferir el conocimiento contenido en esas capas hacia las capas restantes, asegurando que la información no se pierda cuando las capas sean eliminadas.
• Se puede utilizar norma $\ell_1$ o $\ell_2$ para este término.

D. Poda

• Finalmente, se eliminan físicamente las capas que han sido regularizadas y tienen los pesos más bajos.
• El modelo resultante no requiere reentrenamiento (es retraining-free) y mantiene un alto rendimiento.

3. Contribuciones Clave

• Retención de Conocimiento: A diferencia de la eliminación directa, TRSP utiliza la regularización para redistribuir activamente la información de las capas a eliminar hacia las capas preservadas.
• Eficiencia y Sin Reentrenamiento: El método es "libre de reentrenamiento" después de la poda, eliminando la necesidad de costosos procesos de fine-tuning para recuperar el rendimiento.
• Bajo Costo Computacional: Solo requiere una pequeña fracción de datos para la regularización y evita el reentrenamiento masivo.
• Aceleración End-to-End: Al ser un método de poda por capas, ofrece aceleraciones significativas en la inferencia, reduciendo tanto la latencia como aumentando el rendimiento (throughput).

4. Resultados Experimentales

Los autores evaluaron TRSP en varios modelos (Phi-2, OPT, LLaMA2, LLaMA3) y tareas (generación y zero-shot).

• Rendimiento Superior: TRSP superó consistentemente a métodos de poda por capas de vanguardia (como SLEB, ShortGPT, LaCo).
  • En LLaMA2-7B con una poda del 25%, TRSP logró una perplejidad un 20% menor que ShortGPT.
  • En OPT-13B, la caída de precisión promedio fue de solo el 1% en comparación con el modelo denso.
• Aceleración:
  • En OPT-13B (50% de poda), se logró un aumento del 75% en el throughput y una reducción del 46% en la latencia.
  • En LLaMA2-13B (50% de poda), se obtuvo un aumento del 71% en el throughput.
• Robustez: El método mostró robustez ante diferentes ratios de poda (hasta 60%) y diferentes conjuntos de datos de calibración.
• Escalabilidad: Los experimentos adicionales en Llama-2-70B confirmaron que TRSP mantiene su ventaja sobre otros métodos incluso en modelos de gran escala.

5. Significado e Impacto

El trabajo de TRSP representa un avance significativo en la optimización de LLMs al abordar el problema fundamental de la pérdida de conocimiento durante la poda estructurada.

• Cambio de Paradigma: En lugar de simplemente "eliminar y arreglar" (poda + reentrenamiento), TRSP propone "preparar y eliminar" (redistribución de conocimiento + poda), lo que resulta en modelos más compactos sin sacrificar capacidades.
• Viabilidad de Despliegue: Al eliminar la necesidad de reentrenamiento costoso y ofrecer aceleraciones sustanciales, TRSP hace que el despliegue eficiente de LLMs en dispositivos con recursos limitados o en entornos de producción sea mucho más viable.
• Generalización: La técnica de regularización para transferir conocimiento podría inspirar futuras estrategias de compresión en otras arquitecturas de redes neuronales más allá de los transformadores.

En resumen, TRSP es una solución elegante y eficiente que maximiza la relación entre la reducción del tamaño del modelo y la preservación de su inteligencia, resolviendo uno de los cuellos de botella más grandes en la implementación práctica de la IA generativa.