A Training-Time Diagnostic for Generalization via the Log-Alignment Ratio

Título: Un diagnóstico en tiempo de entrenamiento para la generalización mediante la Razón de Alineación Logarítmica

Enunciado del Problema
A medida que escalan los modelos de redes neuronales y los conjuntos de datos, el costo computacional de las ejecuciones de entrenamiento aumenta, lo que hace que el estudio de la dinámica de entrenamiento y el diseño de herramientas de diagnóstico económicas sean cada vez más valiosos. Los profesionales requieren métricas que puedan señalar problemas como el sobreajuste, la inestabilidad o la saturación utilizando únicamente cantidades disponibles durante el entrenamiento, sin depender de costosos datos de validación retenidos. Aunque trabajos anteriores han explorado propiedades espectrales de las matrices de pesos (por ejemplo, rango efectivo, rango estable) o teorías de parametrización (por ejemplo, parametrización de actualización máxima, $\mu$P) para comprender la generalización, las métricas existentes a menudo se centran exclusivamente ya sea en los pesos o en las activaciones, o requieren operaciones computacionalmente costosas como la Descomposición en Valores Singulares (SVD) completa o el Análisis de Componentes Principales (PCA).

Metodología
Los autores investigan la Razón de Alineación Logarítmica (LAR), una métrica introducida originalmente en la teoría de parametrización por Everett et al. [2024] para medir la alineación entre la matriz de parámetros de una capa y sus activaciones de entrada. El artículo reformula teóricamente la LAR para revelar un segundo significado, previamente inexplorado: mide la superposición entre un espectro de pesos ($p$) y un espectro de activaciones ($q$).

• Definiciones:
  • Sea $W$ una matriz con valores singulares $s_i$. El espectro de pesos $p$ se define como los valores singulares al cuadrado normalizados: $p_i = s_i^2 / \sum s_j^2$.
  • Sea $X$ un lote de entradas. El espectro de activaciones $q$ se define como las proyecciones al cuadrado normalizadas de las entradas sobre las direcciones singulares de $W$: $q_i = \sum_{x \in B} (v_i^\top x)^2 / \sum_{x \in B} \|x\|^2$.
• Reformulación: Los autores derivan que la LAR puede expresarse como:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( \sum_{i} p_i q_i \right)$$
  Esta formulación muestra que la LAR es una medida de la superposición logarítmica entre las distribuciones $p$ y $q$. También puede escribirse en términos de covarianza y correlación de Pearson:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \cdot \text{cov}(p, q) + \frac{1}{n} \right)$$
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \sigma_p \sigma_q \text{corr}(p, q) + \frac{1}{n} \right)$$
  donde $\sigma_p$ y $\sigma_q$ representan la "dispersión" o concentración de las distribuciones.
• Eficiencia Computacional: La LAR es calculable a partir de cantidades disponibles durante el paso hacia adelante (normas de matrices) con una sobrecarga despreciable ($O(|B|m + |B|n + mn)$), sin requerir datos retenidos ni una SVD completa.

Contribuciones Clave

1. Reformulación Distribucional: El artículo demuestra que la LAR no es meramente una medida de alineación, sino un cuantificador de la superposición entre los espectros de pesos y activaciones. Proporciona formas estadísticas equivalentes (covarianza, correlación) que hacen que el comportamiento de la métrica sea interpretable en términos de cantidades elementales.
2. Predicción de Dimensión Efectiva: Los autores muestran que cuando una matriz de pesos concentra energía en $k$ direcciones singulares y las activaciones se alinean con ellas, la LAR se relaciona con la dimensión efectiva $k$ mediante $k \approx n^2(1 - \text{LAR})$. Esta fórmula se verifica empíricamente en experimentos de "grokking" (comprensión repentina).
3. Diagnóstico de Sobreajuste a Gran Escala: En el preentrenamiento a gran escala (modelos de 3 mil millones de parámetros), el artículo establece que la LAR de desencaje (unembedding) se estabiliza durante la generalización pero disminuye bruscamente a medida que se acerca el sobreajuste. La desviación de la LAR respecto a una línea base sin sobreajuste rastrea la brecha de generalización.
4. Proxy para la Dispersión Espectral: El estudio encuentra que la optimización impulsa a $p$ y $q$ a estar altamente correlacionadas. En consecuencia, la LAR actúa como un proxy para la varianza (dispersión) de estas distribuciones. En el preentrenamiento, la dirección singular superior domina, lo que hace que la contribución del componente superior ($L_1$) sea un indicador principal de la firma previa al sobreajuste.

Resultados Experimentales
Los autores validan estos hallazgos en dos configuraciones distintas:

1. Experimentos de Grokking (Tareas Algorítmicas Pequeñas):

   • Entrenaron Transformers de 2 capas en 12 tareas de operaciones binarias.
   • Observación: Durante la fase de memorización, la LAR es más baja (indicando $p$ y $q$ difusas). A medida que el modelo transita hacia el "grokking" (generalización), la LAR aumenta a medida que $p$ y $q$ se concentran en un conjunto compartido de direcciones.
   • Dimensionalidad: La dimensión predicha por la LAR, $k_{pred} = n^2(1 - \text{LAR})$, coincide estrechamente con $k_{95}$ (el número de componentes principales que explican el 95% de la varianza) en los puntos de control finales, con un Error Cuadrático Medio (RMSE) bajo.
   • Diferenciación: Los modelos que logran "grokking" alcanzan valores finales de LAR más altos ($>0.65$) en comparación con los modelos que no logran "grokking" ($\le 0.65$).

2. Preentrenamiento a Gran Escala (Modelos de Lenguaje de 3B):

   • Entrenaron modelos basados en Gemma-2 (3 mil millones de parámetros) utilizando los optimizadores Adam y Muon en conjuntos de datos que van desde 200 millones hasta 100 mil millones de tokens.
   • Generalización vs. Sobreajuste: En regímenes sin sobreajuste (por ejemplo, 100 mil millones de tokens), la LAR sube a un pico y luego disminuye lentamente, estabilizándose. En regímenes de sobreajuste (conjuntos de datos más pequeños), la LAR disminuye bruscamente a medida que se acerca el paso de sobreajuste.
   • Brecha de Generalización: La diferencia entre la LAR de una ejecución específica y la ejecución base de 100 mil millones (diferencial de LAR) refleja estrechamente la brecha de generalización (pérdida de validación menos pérdida de entrenamiento).
   • Mecanismo: La disminución de la LAR antes del sobreajuste es impulsada principalmente por el componente singular superior ($L_1$). Con Adam, esto es impulsado por la difusión de las activaciones ($q_1$); con Muon, implica tanto pesos ($p_1$) como activaciones.

Significado y Afirmaciones
El artículo afirma que la LAR ofrece un diagnóstico práctico, computacionalmente económico y libre de datos para la generalización.

• Eficiencia: A diferencia del rango efectivo u otras medidas espectrales que requieren SVD, la LAR puede calcularse con una sobrecarga despreciable durante el paso hacia adelante.
• Poder Predictivo: Rastrea con éxito la transición de la memorización a la generalización en el "grokking" y sirve como señal de advertencia temprana para el sobreajuste en el preentrenamiento a gran escala, potencialmente habilitando métodos de entrenamiento sin evaluación.
• Interpretabilidad: Al vincular la LAR con la dispersión de los espectros de pesos y activaciones, la métrica proporciona una visión unificada de cómo la dinámica de optimización (concentración vs. difusión de espectros) se relaciona con la generalización del modelo.

Los autores reconocen limitaciones, señalando que su evidencia es correlacional, que se centran principalmente en la capa de desencaje (unembedding) y que aún no han probado si regularizar directamente la dispersión espectral evitaría el sobreajuste. Sugieren que trabajos futuros deberían explorar la LAR para otros tipos de matrices y su comportamiento bajo tamaños de lote variables.