Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

Resumen Técnico: Más allá de la Perplejidad: Un Estudio Geométrico y Espectral del Pre-entrenamiento de Bajo Rango

Enunciado del Problema
El pre-entrenamiento de Modelos de Lenguaje Grandes (LLM) está limitado por los costos de memoria asociados al almacenamiento de pesos, gradientes y estados del optimizador de rango completo. El pre-entrenamiento de bajo rango ha surgido como una solución para reducir estos costos, con métodos que han evolucionado hacia tres categorías: basados en optimizadores (por ejemplo, GaLore, Fira), basados en arquitectura (por ejemplo, CoLA, SLTrain) y basados en adaptadores (por ejemplo, ReLoRA). Sin embargo, la evaluación de estos métodos sigue siendo limitada. Las prácticas actuales dependen casi exclusivamente de la perplejidad de validación obtenida en ejecuciones de una sola semilla. Los autores argumentan que la perplejidad es un mal proxy para la calidad de la solución; modelos con perplejidad idéntica pueden converger hacia regiones fundamentalmente diferentes del paisaje de pérdida con representaciones internas distintas. Además, las comparaciones existentes a menudo no logran distinguir entre qué tan bien un método ajusta fino un modelo pre-entrenado y qué tan bien un modelo pre-entrenado con ese método se generaliza aguas abajo.

Metodología
Los autores introducen un marco diagnóstico para caracterizar las soluciones encontradas por cinco métodos de pre-entrenamiento de bajo rango (GaLore, Fira, CoLA, SLTrain, ReLoRA) frente al entrenamiento de rango completo en tres escalas de modelo (60M, 130M, 350M). El estudio evalúa estos métodos utilizando 16 métricas en cuatro dimensiones:

1. Paisaje de Pérdida 1-D (Direcciones Aleatorias): Los autores perturban los parámetros del modelo a lo largo de 100 direcciones aleatorias para medir la media y la varianza de la pérdida de validación. Derivan métricas escalares para Agudeza (la tasa de crecimiento de la pérdida alejada del punto de control) y Varianza de Dirección (anisotropía de la superficie de pérdida).
2. Paisaje de Pérdida 1-D (Direcciones Top-K de PCA): Se aplican perturbaciones a lo largo de las direcciones singulares principales de las matrices de pesos para analizar la geometría a lo largo de los ejes más significativos.
3. Interpolación 1-D: Los autores miden la pérdida a lo largo de trayectorias lineales entre puntos de control. Introducen dos métricas específicas:
   • Altura de la Barrera entre Puntos de Control Consecutivos (CCBH): Mide la barrera entre puntos de control secuenciales del mismo método.
   • Altura de la Barrera Inter-Método (IMBH): Una métrica novedosa que cuantifica la separación geométrica entre diferentes métodos de entrenamiento en el mismo paso de entrenamiento. Esto revela si diferentes métodos convergen hacia la misma cuenca o hacia soluciones geométricamente distintas.
4. Métricas Espectrales y de Rango: El estudio rastrea la estructura espectral de los pesos y las actualizaciones, incluyendo Rango Efectivo, Rango Estable, Brecha Espectral y Rango Umbral.
5. Similitud de Activaciones: Los autores comparan los estados ocultos de los modelos de bajo rango frente a las líneas base de rango completo utilizando Distancia L2 de Activación, Similitud Coseno y CKA Lineal (Alineación de Kernel Centrada).

Finalmente, los autores prueban si estas características geométricas y espectrales pueden predecir el rendimiento de cero disparos aguas abajo mejor que la perplejidad por sí sola, utilizando regresión lineal con validación cruzada dejando una escala fuera (LOSO) y dejando un método fuera (LOMO).

Contribuciones Clave

1. Marco Diagnóstico: El artículo propone una suite de evaluación integral que incluye la métrica novedosa IMBH, la cual cuantifica la separación geométrica entre algoritmos de entrenamiento en lugar de solo entre semillas o épocas.
2. Hallazgos Empíricos Trans-Familiares: El estudio establece que los métodos de bajo rango no son equivalentes al entrenamiento de rango completo ni entre sí. Los hallazgos clave incluyen:
   • Separación Geométrica: Los métodos convergen hacia cuencas geométricamente distintas, con la IMBH creciendo monótonamente a medida que avanza el entrenamiento.
   • Firmas Espectrales: GaLore coincide estrechamente con la estructura espectral de las actualizaciones de rango completo, mientras que CoLA no lo hace.
   • Divergencia de Representaciones: Las activaciones de bajo rango divergen cada vez más de las de rango completo en las capas posteriores. CoLA exhibe activaciones ortogonales, mientras que Fira recupera la alineación solo en la capa final.
3. Mejora en la Predicción Aguas Abajo: La perplejidad de validación por sí sola es un predictor incompleto del rendimiento aguas abajo. Aumentar la perplejidad con características geométricas y espectrales mejora significativamente la precisión de la predicción (la correlación de Pearson de LOSO aumenta de 0.873 a 0.913).
4. Base de Código Unificada: Los autores lanzan una pipeline de código abierto que integra el pre-entrenamiento, el cálculo de métricas y la evaluación aguas abajo para facilitar la investigación futura.

Resultados

• Geometría del Paisaje de Pérdida: El entrenamiento de rango completo se asienta en cuencas más agudas a lo largo de direcciones aleatorias, pero en cuencas más planas a lo largo de la dirección Top-1 de PCA. Por el contrario, los métodos de bajo rango como Fira y SLTrain convergen hacia cuencas más planas a lo largo de direcciones aleatorias, mientras que CoLA y ReLoRA alcanzan las cuencas más agudas. GaLore exhibe un perfil mixto: moderadamente plano a lo largo de direcciones aleatorias pero agudo a lo largo de su dirección Top-1 de PCA.
• Barreras Inter-Método: Todos los métodos ocupan cuencas geométricamente separadas. En la escala de 60M, la barrera entre rango completo y cualquier método de bajo rango es más alta que las barreras entre métodos de bajo rango. Fira converge hacia la cuenca geométricamente más cercana al rango completo, mientras que SLTrain permanece como el más aislado.
• Análisis Espectral: El rango completo domina todas las métricas de rango. GaLore rastrea al rango completo estrechamente en el rango efectivo y estable, pero exhibe una brecha espectral más pequeña. Fira y GaLore reproducen con mayor fidelidad la distribución espectral de las actualizaciones de pesos de rango completo.
• Alineación de Activaciones: ReLoRA rastrea las representaciones de rango completo con mayor cercanía en todas las escalas. CoLA muestra la mayor desviación, con activaciones que se vuelven ortogonales a la línea base en las capas posteriores. Fira muestra una deriva promediada por capas alta, pero sorprendentemente baja desviación en la capa final.
• Rendimiento Aguas Abajo: La perplejidad de validación no predice consistentemente la clasificación aguas abajo. Por ejemplo, Fira logra la perplejidad más baja en todas las escalas, sin embargo, CoLA lidera en múltiples categorías aguas abajo en 130M y 350M a pesar de no clasificarse entre los dos primeros en perplejidad.

Significado
El artículo concluye que los métodos de pre-entrenamiento de bajo rango no son aproximaciones intercambiables del entrenamiento de rango completo; alteran fundamentalmente las soluciones alcanzadas. El estudio demuestra que la perplejidad de validación proporciona una imagen incompleta de la calidad del modelo. Al incorporar métricas geométricas y espectrales, los investigadores pueden predecir mejor el rendimiento aguas abajo y comprender los sesgos inductivos distintos introducidos por diferentes estrategias de bajo rango. El trabajo destaca que, mientras algunos métodos (como Fira) convergen hacia cuencas geométricamente cercanas al rango completo, otros (como CoLA) ocupan valles distintos, lo que sugiere que la elección del método de bajo rango dicta la región específica del paisaje de pérdida en la que habita el modelo.