Resumen Técnico: TraceNAS

Definición del Problema

El despliegue de Modelos de Lenguaje de Gran Escala (LLMs) está limitado por el masivo conteo de parámetros y los requisitos computacionales. Si bien la poda estructurada ofrece una vía para el despliegue eficiente mediante la eliminación de unidades arquitectónicas (por ejemplo, cabezales de atención, bloques transformer), los métodos existentes enfrentan limitaciones significativas:

1. Dependencias Locales vs. Globales: La mayoría de los métodos de poda evalúan la importancia de capas, cabezales o canales de forma aislada utilizando heurísticas locales (por ejemplo, magnitudes ponderadas por activación o reconstrucción basada en el Hessiano). Estos enfoques ignoran las complejas dependencias estructurales globales y la propagación de errores inducidos por la poda a través del modelo.
2. Costo Computacional de la Búsqueda Consciente del Entrenamiento: Los métodos que intentan capturar dependencias globales suelen depender de la Búsqueda de Arquitectura Neuronal (NAS) consciente del entrenamiento. Estos enfoques requieren actualizaciones de gradiente iterativas y flujos continuos de tokens no vistos para prevenir el sobreajuste, lo que hace que el proceso de búsqueda sea tan costoso computacionalmente como el propio entrenamiento de recuperación post-poda.
3. Inadecuación de los Proxies de Cero Disparos Existentes: Los proxies actuales libres de entrenamiento están diseñados para identificar arquitecturas entrenables y con inicialización aleatoria centrándose en la entrenabilidad (por ejemplo, estabilidad del gradiente, suavidad del paisaje de pérdida). Estos fallan al evaluar la "herencia funcional": la capacidad de un modelo podado para recuperar el rendimiento preservando el paisaje de pérdida preentrenado.

Metodología

Se propone TraceNAS como un marco de NAS libre de entrenamiento diseñado para optimizar conjuntamente la profundidad y el ancho de los LLMs sin entrenamiento durante la búsqueda.

1. Espacio de Búsqueda y Codificación

El marco trata al LLM preentrenado como una super-red. El espacio de búsqueda $G$ se define como una distribución conjunta sobre:

• Profundidad ($d$): Una máscara discreta sobre $L$ bloques transformer, permitiendo una reducción de profundidad no uniforme manteniendo la continuidad del flujo residual.
• Ancho ($\kappa$): Razones de retención de parámetros para los sub-bloques de atención y MLP, permitiendo una dispersión heterogénea a través del modelo.

2. Realización Arquitectónica In-Situ

Para evitar la sobrecarga de memoria de instanciar miles de modelos candidatos, TraceNAS emplea una estrategia de enmascaramiento in-situ:

• Enmascaramiento de Ancho: Utiliza un producto ponderado por activación (similar a Wanda) para generar máscaras binarias para los canales de peso. Este heurístico $O(d^2)$ identifica canales funcionalmente críticos mientras preserva las características atípicas (outliers).
• Enmascaramiento de Profundidad: Omite bloques enteros si son desactivados, enrutando las activaciones a través de conexiones residuales.
• Ejecución: El marco modifica temporalmente las matrices de peso del modelo base para calcular gradientes sobre un conjunto de calibración mínimo, luego restaura los pesos originales, preservando el modelo base para iteraciones subsecuentes.

3. El Proxy de Cero Disparos: Correlación de Trazas de Gradiente

La innovación central es un proxy $\Phi$ invariante de escala que mide la alineación funcional entre un candidato podado ($M_{sub}$) y la base preentrenada ($M_{base}$).

• Trazas de Gradiente: Los gradientes se calculan sobre un conjunto de calibración. Para asegurar la escalabilidad, estos se proyectan en un subespacio de Bajo Rango (LoRA), aprovechando la baja dimensionalidad intrínseca de los LLMs para reducir la sobrecarga de memoria en órdenes de magnitud.
• Métrica de Alineación: El proxy calcula el Coeficiente de Correlación de Pearson entre las trazas de gradiente estandarizadas de los modelos podado y base para cada capa. Esto asegura la invarianza de escala, enfocándose en la alineación direccional en lugar de los cambios de magnitud.
• Agregación Ponderada por Dispersión: La puntuación final $\Phi$ es una suma ponderada de las correlaciones por capa, donde los pesos corresponden a las razones de retención de parámetros. Esto ancla la puntuación global en regiones de alta capacidad, evitando que el ruido de los sub-bloques agresivamente podados sesgue la clasificación.

4. Búsqueda Evolutiva

TraceNAS emplea una estrategia evolutiva para navegar el espacio de búsqueda discreto:

• Inicialización: Se genera una población basada en prioridades de importancia global.
• Evolución: La población experimenta cruce y mutación. La evolución de la profundidad utiliza mutaciones de cambio de bit y cruce de rango contiguo, mientras que la evolución del ancho utiliza interpolación y jitter Gaussiano.
• Selección: Los candidatos se clasifican únicamente por el proxy $\Phi$, seleccionando modelos con el mayor potencial de recuperación de rendimiento durante el preentrenamiento continuo (CPT) post-poda.

Contribuciones Clave

1. NAS Escalable Libre de Entrenamiento: TraceNAS proporciona un marco unificado para la poda conjunta de profundidad y ancho, permitiendo el descubrimiento de arquitecturas no uniformes sin entrenamiento durante la búsqueda.
2. Proxy de Herencia Funcional: Introduce un nuevo proxy basado en gradientes e invariante de escala que evalúa la alineación de las trazas de gradiente entre los modelos podados y preentrenados. Logra una alta correlación con la precisión en tareas posteriores (Spearman $\rho = 0.94$, Kendall $\tau = 0.82$).
3. Eficiencia: El marco reduce la sobrecarga de búsqueda mediante el uso de subespacios de gradiente de bajo rango y enmascaramiento in-situ, permitiendo el descubrimiento de modelos de alta fidelidad en una sola GPU.

Resultados Experimentales

TraceNAS fue evaluado en Llama-2-7B, Llama-3.1-8B y Qwen-2.5-14B, podándolos a 2.7B, 4.6B y 8.4B de parámetros respectivamente.

• Eficiencia de Búsqueda: TraceNAS identificó arquitecturas óptimas en 8.5 horas de GPU (en una A100), lo que representa una reducción de 10× en el tiempo de búsqueda comparado con baselines conscientes del entrenamiento como DarwinLM. Requirió solo 98M de tokens de búsqueda, una reducción de 10× respecto a DarwinLM.
• Rendimiento:
  • Llama-2-7B: TraceNAS logró la mayor precisión promedio (62.81%) a través de ocho benchmarks de razonamiento, superando a DarwinLM (62.57%) y ShearedLLaMA (62.63%). Notablemente, el rendimiento de ShearedLLaMA cayó significativamente (a 60.86%) cuando se restringió al mismo presupuesto de tokens usado por TraceNAS.
  • Llama-3.1-8B y Qwen-2.5-14B: TraceNAS superó a los baselines de poda uniforme y a los métodos conscientes del entrenamiento (DarwinLM, E3-Pruner) en tareas de razonamiento complejo (ej. ARC-Challenge, BoolQ) mientras requería significativamente menos tokens de recuperación (20B vs. 94B para Minitron).
• Validación del Proxy: El proxy demostró una alta estabilidad a través de diversos hiperparámetros (tamaño de muestra, longitud de contexto y rango de LoRA), confirmando que la alineación de la traza de gradiente captura efectivamente la herencia del modelo incluso bajo configuraciones restringidas.

Significado y Reivindicaciones

El artículo afirma que TraceNAS aborda el cuello de botella crítico del costo computacional en la poda estructurada. Al cambiar el enfoque de la "entrenabilidad" (aprendizaje desde cero) a la "herencia" (preservación del paisaje de pérdida preentrenado), TraceNAS permite el descubrimiento de modelos podados no uniformes de alto rendimiento con un gasto mínimo de recursos.

Los autores enfatizan que su método:

• Elimina el costo prohibitivo del entrenamiento durante la búsqueda, haciendo que la compresión de modelos a gran escala sea accesible.
• Captura dependencias estructurales globales que las heurísticas locales pasan por alto, asegurando la preservación de capacidades de razonamiento complejas.
• Reduce la huella de carbono de la compresión de modelos al disminuir drásticamente las horas-GPU requeridas para la búsqueda de arquitectura.

El trabajo posiciona a TraceNAS como una base escalable para el despliegue de LLMs eficientes en entornos con recursos limitados, manteniendo la integridad funcional de los modelos densos preentrenados.