DarwinLM: Evolutionary Structured Pruning of Large Language Models
DarwinLM es un método de poda estructurada sensible al entrenamiento que emplea un proceso de búsqueda evolutiva con un entrenamiento ligero de múltiples pasos para identificar subestructuras de modelos no uniformes óptimas, logrando un rendimiento de vanguardia en varios modelos de lenguaje extensos al tiempo que reduce significativamente los requisitos de datos de entrenamiento post-compresión.
Autores originales:Shengkun Tang, Oliver Sieberling, Eldar Kurtic, Zhiqiang Shen, Dan Alistarh
Imagina que tienes un cerebro de robot gigante y superinteligente que puede escribir historias, resolver problemas matemáticos y charlar como un humano. Esto es lo que los científicos llaman un Modelo de Lenguaje Extenso (LLM). Estos cerebros son asombrosos, pero también son increíblemente pesados y hambrientos de electricidad, como un dragón que necesita todo el oro de un castillo solo para despertar. Debido a que son tan grandes, son difíciles de transportar o de ejecutar en computadoras regulares. Para solucionar esto, los investigadores han estado intentando "encoger" estos cerebros. Una forma popular de hacerlo es la "poda". Piensa en la poda como recortar un bonsái: cortas las ramas que no necesitas para hacer el árbol más pequeño y rápido, con la esperanza de que siga viéndose y actuando igual. Pero aquí está la parte difícil: no todas las ramas son iguales. Algunas son vitales para la forma del árbol, mientras que otras son solo hojas extra. Si cortas las equivocadas, el árbol muere. Si cortas las correctas, obtienes un árbol diminuto y rápido que aún da frutos. La gran pregunta es: ¿cómo sabes exactamente qué ramas cortar sin matar accidentalmente al genio que hay dentro?
Aquí entra DarwinLM, un nuevo método que actúa como un naturalista digital intentando evolucionar el cerebro diminuto perfecto. En lugar de simplemente adivinar qué partes cortar, los investigadores dejan que el modelo "evolucione" a través de un proceso inspirado en cómo la naturaleza selecciona a los animales más aptos. Comienzan con un modelo grande y crean muchas versiones de "descendencia" ligeramente diferentes, cada una con un patrón de cortes único. Luego, someten a estos descendientes a una prueba de entrenamiento rápida y ligera —como un sprint— para ver cuáles pueden seguir corriendo rápido y pensando con claridad. Los ganadores sobreviven para formar la siguiente generación, mientras que los perdedores son descartados. Este ciclo se repite, refinando lentamente el modelo hasta que encuentra el equilibrio perfecto entre tamaño e inteligencia. El artículo muestra que este enfoque evolutivo es mucho mejor que simplemente cortar las cosas de manera uniforme (como recortar la misma cantidad de cada capa). De hecho, DarwinLM puede encoger modelos como Llama-2 y Llama-3.1 a la mitad de su tamaño manteniendo su inteligencia casi intacta, y lo hace utilizando mucha menos información de entrenamiento que los métodos anteriores. Por ejemplo, mientras que otro método famoso necesitó 50 mil millones de palabras para entrenar su modelo pequeño, DarwinLM logró resultados incluso mejores con solo 10 mil millones. Incluso funciona en modelos complejos de "Mezcla de Expertos" (Mixture of Experts), demostrando que puedes recortar la grasa sin perder el músculo, haciendo que la IA poderosa sea accesible en dispositivos cotidianos.
Resumen Técnico: DarwinLM
Declaración del Problema Los Modelos de Lenguaje de Gran Escala (LLMs) han alcanzado un éxito significativo, pero enfrentan costos computacionales masivos que dificultan su despliegue, particularmente en aplicaciones de tiempo real. Si bien la poda estructurada ofrece una solución al comprimir los modelos para proporcionar mejoras de velocidad de extremo a extremo en hardware convencional, los métodos existentes a menudo no tienen en cuenta dos factores críticos:
Sensibilidad No Uniforme: Diferentes componentes del modelo (capas, cabezales de atención, dimensiones de MLP) exhiben sensibilidades variables a la poda. La compresión uniforme suele provocar un colapso en el rendimiento.
Entrenamiento Post-Compresión: Un método de poda no solo debe identificar una subestructura capaz, sino también asegurar que la estructura sea apta para la recuperación mediante el ajuste fino (fine-tuning) post-compresión. Los métodos previos, como ZipLM, a menudo optimizan errores locales por capa o el rendimiento de un solo paso (one-shot), lo cual no se correlaciona bien con el rendimiento en tareas posteriores o la recuperación mediante ajuste fino.
Metodología Los autores proponen DarwinLM, un método de poda estructurada consciente del entrenamiento que utiliza un proceso de búsqueda evolutiva para encontrar asignaciones de dispersión (sparsity) no uniformes óptimas. El flujo de trabajo consta de dos etapas principales:
Búsqueda Evolutiva con Selección Consciente del Entrenamiento:
Inicialización: La búsqueda comienza con un modelo "padre" generado mediante una poda estructurada de un solo paso utilizando información de segundo orden (formulación de Optimal Brain Surgeon).
Mutación: En cada generación, el algoritmo genera modelos "descendientes" copiando al padre y aplicando una mutación de "cambio de nivel" (level switch). Esto implica desplazar los niveles de dispersión entre subbloques (por ejemplo, aumentando la dispersión en una capa mientras se disminuye en otra) para mantener la restricción de dispersión global o de aceleración.
Selección Consciente del Entrenamiento: Una innovación central es la integración de un ajuste fino ligero en el proceso de selección. En lugar de seleccionar descendientes basándose únicamente en métricas de un solo paso (como la divergencia KL en un conjunto de calibración pequeño), el método emplea un proceso de selección de múltiples pasos. Los descendientes se entrenan en presupuestos de tokens progresivamente mayores (por ejemplo, 10K → 50K → 200K tokens). Los modelos con bajo rendimiento son eliminados en cada etapa. Esto asegura que el modelo "más apto" seleccionado no sea solo preciso en un estado estático, sino que también sea robusto y recuperable tras un entrenamiento continuo.
Terminación: Una vez que la búsqueda converge, el candidato final se somete a una etapa de ajuste fino más extensa (por ejemplo, 10B tokens) antes de la evaluación final.
Base de Datos de Capas Podadas:
Para facilitar una búsqueda eficiente, el método precalcula una base de datos de subbloques dispersos (módulos MLP y de Atención) en varios niveles de dispersión discretizados.
Para los módulos de Atención, la poda se realiza por cabezal. Para los MLPs, las dimensiones se podan en múltiplos de 32.
Para modelos con Atención de Consulta por Grupo (GQA), las matrices K y V no se podan directamente; en su lugar, los cabezales correspondientes se eliminan durante el pase hacia adelante (forward pass) para mantener la compatibilidad computacional.
Extensión a Arquitecturas MoE:
El método se extiende a modelos de Mezcla de Expertos (MoE) (por ejemplo, Qwen3-30B-A3B). En este entorno, la búsqueda optimiza la dispersión dentro de los MLPs de los expertos manteniendo una dispersión uniforme a través de los bloques MoE. Los módulos de Atención generalmente se dejan sin podar, ya que los expertos contienen la mayoría de los parámetros.
Contribuciones Clave
Búsqueda Evolutiva Consciente del Entrenamiento: DarwinLM introduce un nuevo marco evolutivo que considera explícitamente la capacidad del modelo para recuperar el rendimiento durante el ajuste fino, en lugar de optimizar únicamente la precisión de un solo paso.
Poda Estructurada No Uniforme: El método permite una compresión no uniforme y de grano fino de los LLMs, aprovechando las variaciones de sensibilidad de las diferentes capas para lograr mayores ratios de compresión sin pérdida de precisión.
Adaptación a MoE: Este trabajo representa la primera exploración de la poda estructurada no uniforme aplicada específicamente a arquitecturas MoE, demostrando que la dispersión estructurada es efectiva incluso en modelos complejos basados en expertos.
Eficiencia: El proceso de búsqueda es altamente eficiente, completándose en aproximadamente 8 horas en 4 GPUs de consumo, mientras que la etapa final de ajuste fino toma aproximadamente medio día en un clúster estándar.
Resultos Experimentales Los autores validan DarwinLM en Llama-2-7B, Llama-3.1-8B, Qwen-2.5-14B-Instruct y Qwen3-30B-A3B (MoE).
Rendimiento: DarwinLM logra un rendimiento de vanguardia (state-of-the-art) en la poda estructurada de un solo paso. Por ejemplo, al podar Llama-3.1-8B a la mitad de su tamaño, alcanza un 5.9% más de precisión promedio en zero-shot comparado con el mejor método previo (ZipLM).
Eficiencia de Datos: El método reduce significativamente el presupuesto de entrenamiento requerido para la recuperación. DarwinLM supera a ShearedLlama (que utiliza 50B tokens para el ajuste fino) requiriendo solo 10B de tokens. En el mismo escenario de 10B de tokens, DarwinLM supera a ShearedLlama.
Resultados de MoE: En el modelo Qwen3-30B-A3B, DarwinLM produce una variante 16B-A2B que retiene ≥90% de la precisión del modelo original tras 10B de tokens de ajuste fino.
Eficiencia de Hardware: Los modelos podados demuestran aceleraciones directas y reducciones de memoria. Por ejemplo, el modelo DarwinLM de 2.7B logra 1.98× de throughput y utiliza 2.43× menos memoria en comparación con la línea base densa de 7B en GPUs L40s.
Comparación: El método supera a las líneas base de poda uniforme y a otros métodos de poda estructurada (incluyendo ShortGPT, Shortened-Llama y EvoPress) en varios niveles de dispersión, particularmente en tasas de compresión altas donde otros métodos se degradan rápidamente.
Significancia El artículo afirma que DarwinLM representa un avance significativo en la compresión de LLMs al cerrar la brecha entre la búsqueda estructural y la recuperación post-entrenamiento. Al demostrar que la poda estructurada no uniforme puede aplicarse eficazmente tanto a arquitecturas densas como MoE con un mínimo de datos de entrenamiento, este trabajo desafía la noción de que la compresión de alto rendimiento requiere presupuestos masivos de reentrenamiento. Los autores posicionan a DarwinLM como una solución práctica y agnóstica al hardware que permite el despliegue de LLMs eficientes sin sacrificar la precisión.