The Geometric Wall: Manifold Structure Predicts Layerwise Sparse Autoencoder Scaling Laws

Resumen Técnico: La Pared Geométrica: La Estructura de Variedad Predice las Leyes de Escalamiento por Capas de los Autoencoders Esparsos

Enunciado del Problema
Los Autoencoders Esparsos (SAE) operacionalizan la Hipótesis de Representación Lineal (LRH) reconstruyendo las activaciones del modelo como combinaciones lineales esparsas de átomos del diccionario. Este enfoque asume implícitamente que el espacio de activaciones está bien aproximado por una estructura globalmente lineal. Sin embargo, las leyes de escalamiento existentes para la pérdida de reconstrucción de los SAE, específicamente la ley de potencias conjunta que relaciona la pérdida con el ancho del diccionario ($n$) y la esparsidad ($k$), se han ajustado únicamente en capas individuales (por ejemplo, capa 5/6 de GPT-4) y no tienen en cuenta la variación aguda en el error de reconstrucción observada entre diferentes capas de una red.

El artículo argumenta que esta variación no es meramente una limitación de recursos, sino una huella empírica de una discrepancia geométrica. Mientras que los SAE asumen una base plana y globalmente lineal, las variedades de activación de los modelos de lenguaje grandes son curvas, y su dimensión intrínseca varía sistemáticamente con la profundidad. En consecuencia, ningún diccionario lineal esparso único puede igualar uniformemente estas variedades, lo que hace que el escalamiento ancho-esparsidad del SAE se convierta en una función dependiente de la capa de la geometría de la variedad, en lugar de una ley universal.

Metodología
Los autores realizan el primer estudio de escalamiento transversal a capas utilizando 844 checkpoints de SAE de la corriente residual de la familia Gemma Scope, que cubren todas las 26 capas de Gemma 2 2B y todas las 42 capas de Gemma 2 9B. El análisis procede en dos etapas:

1. Etapa 1: Ajuste de Superficie de Escalamiento por Capa

   • Para cada capa, los autores ajustan una superficie de escalamiento log-lineal "sin suelo" a la pérdida de reconstrucción ($L$) en función del ancho del diccionario ($n$) y la esparsidad ($k$).
   • Debido a la cobertura limitada de anchos en la mayoría de las capas (solo disponibles dos anchos de columna vertebral), ajustan una superficie de 4 parámetros: $\log L = a_0 + \beta_n \log n + \beta_k \log k + \gamma \log n \log k$.
   • A partir de esto, derivan el exponente de escalamiento de ancho por capa $\alpha_\ell(k) = -(\beta_{n,\ell} + \gamma_\ell \log k)$.
   • En seis capas "de muestra" con rejillas de anchos más ricas ($\ge3$ anchos), realizan un ajuste no lineal conjunto de una superficie de 6 parámetros "con suelo" ($L = A(k)n^{-\alpha(k)} + B(k)$) para identificar el suelo de reconstrucción asintótico $B(k)$.

2. Etapa 2: Regresión Geométrica Transversal a Capas

   • Los autores calculan cuatro resúmenes geométricos por capa de la variedad de activación utilizando estimadores euclidianos extrínsecos (TWO-NN, PCA multi-escala, etc.):
     • Dimensión intrínseca ($d_{int}$): Grados de libertad locales.
     • Curvatura multi-escala ($\kappa_{ms}$): Desviación promedio del plano tangente a través de radios variables.
     • Variación del plano tangente ($\kappa_{tv}$): Tasa de rotación del espacio tangente.
     • Heterogeneidad ($\nu$): Variación local en la dimensión intrínseca.
   • Regresan las salidas de la Etapa 1 (específicamente el exponente de ancho $\alpha_\ell(k)$ y los coeficientes $\beta_n, \gamma$) contra estos resúmenes geométricos utilizando Mínimos Cuadrados Ordinarios (OLS).
   • La prueba de hipótesis compara una nulidad invariante a la geometría ($H_0$) contra modelos donde el objetivo depende de características individuales, pares de características de baja correlación o las cuatro características.
   • Transferencia Trans-modelo: Los coeficientes aprendidos en un modelo (por ejemplo, 2B) se utilizan para predecir los exponentes por capa del otro modelo (9B) para probar la existencia de una ley geométrica transferible.

Contribuciones Clave

1. Marco Geométrico para Residuos de SAE: El artículo proporciona una cuenta de información geométrica de pullback que explica por qué los espacios de activación pueden no admitir códigos lineales esparsos globalmente eficientes. Plantea que la métrica "natural" en el espacio de activación es la métrica de Fisher-Rao de pullback, mientras que los SAE optimizan la distancia euclidiana ambiental, creando una discrepancia base-variedad.
2. Ley de Escalamiento Condicionada por Geometría: Los autores extienden la ley de escalamiento estándar ancho-esparsidad a una cuenta condicionada por la geometría. Demuestran que los parámetros de la ley de escalamiento (específicamente el exponente de ancho) son funciones por capa de la geometría de la variedad de activación.
3. Ley Geométrica Trans-modelo y Acoplamiento de Suelo:
   • Muestran que los coeficientes de regresión aprendidos en un modelo predicen con éxito los exponentes de escalamiento de ancho por capa de otro modelo, indicando una ley geométrica transferible.
   • En las capas de muestra, el suelo de reconstrucción asintótico ($B$) sigue el orden geométrico: las capas con mayor curvatura y dimensión intrínseca exhiben suelos irreducibles más altos.

Resultados

• Predicción de Exponentes de Escalamiento: La geometría de la variedad predice fuertemente el exponente de escalamiento de ancho por capa $\alpha_\ell(k=50)$. La curvatura multi-escala ($\kappa_{ms}$) es la característica individual dominante, explicando casi toda la variación transversal a capas en el modelo 2B y la mayoría en el modelo 9B.
• Transferibilidad: La ley geométrica es transferible. Los coeficientes de regresión ajustados en Gemma 2 2B predicen los exponentes 9B con alta fidelidad ($R^2 \approx 0.92-0.93$), y viceversa, lo que sugiere que la relación no es un artefacto de un tamaño de modelo específico.
• La Pared Geométrica: En las capas donde el suelo asintótico es identificable, la magnitud del suelo se correlaciona con la complejidad geométrica. Una mayor curvatura y dimensión intrínseca corresponden a suelos más altos. Esto se interpreta como un residuo de segundo orden irreducible: un diccionario lineal esparso plano no puede aproximar perfectamente una variedad curva, dejando un error residual que no desaparece incluso cuando el ancho del diccionario tiende a infinito.
• Interpretación Mecanística: La "pared geométrica" no es un techo de recursos finitos, sino un límite dependiente de la geometría. Una mayor curvatura aumenta la desviación de segundo orden del plano tangente, mientras que una mayor dimensión intrínseca reduce la porción del espacio tangente local cubrible por un presupuesto finito de átomos.

Significado y Afirmaciones
El artículo afirma replantear el límite de los SAE como una pared dependiente de la geometría en lugar de un techo de recursos. Argumenta que la variación en la fidelidad de los SAE a través de las capas está condicionada por la estructura geométrica de las activaciones, que las prácticas actuales de escalamiento de una sola capa dejan sin modelar.

Los autores no afirman que los SAE sean inadecuados para la interpretabilidad. En cambio, afirman que el comportamiento de escalamiento por capa está fundamentalmente restringido por la estructura de la variedad que el SAE intenta reconstruir. Los hallazgos sugieren que la "hipótesis de representación lineal" es una aproximación dependiente de la métrica; las "direcciones lineales significativas" no son libres de métrica, y la alineación entre la métrica euclidiana ambiental (utilizada por los SAE) y la métrica de pullback predictiva varía sistemáticamente con la profundidad. El trabajo proporciona un marco diagnóstico para comprender dónde y por qué falla la reconstrucción de SAE, vinculándolo directamente con la geometría intrínseca de los estados ocultos de la red neuronal.