Resumen Técnico: Transferencia de Hiperparámetros para Memorias Asociativas Densas

Enunciado del Problema

El aprendizaje profundo moderno depende en gran medida de la transferencia de hiperparámetros (HP), donde se identifican configuraciones óptimas (por ejemplo, escalas de inicialización, tasas de aprendizaje) en modelos y conjuntos de datos pequeños y luego se transfieren a escalas mayores para evitar un ajuste prohibitivamente costoso. Aunque existen marcos teóricos para la transferencia de HP en redes de alimentación hacia adelante (por ejemplo, basados en la heurística $\mu$P), no se han desarrollado para Memorias Asociativas Densas (DenseAMs).

Las DenseAMs presentan desafíos únicos para los métodos de transferencia existentes debido a tres características arquitectónicas específicas:

1. Compartición de Pesos: La función de energía requiere que los pesos se compartan entre capas y dentro de las operaciones, a diferencia de las redes de alimentación hacia adelante estándar.
2. Dinámica Residual: Las ecuaciones dinámicas que describen el descenso de energía contienen nativamente conexiones residuales.
3. Activaciones Empinadas: Para lograr leyes de escalado superiores en el almacenamiento de información, las DenseAMs utilizan funciones de activación que alcanzan picos rápidamente (por ejemplo, ReLUs de alta potencia, softmax), las cuales son raras en las arquitecturas de alimentación hacia adelante estándar.

Estos factores crean un paisaje complejo donde el análisis dimensional estándar falla al predecir cómo deben escalar los hiperparámetros con el tamaño del modelo ($N, K$), el tamaño del conjunto de datos ($P$) y el tamaño del lote ($B$).

Metodología

Los autores inician el estudio de la transferencia de HP tipo $\mu$P para DenseAMs derivando prescripciones explícitas para escalar hiperparámetros en el régimen de escalado proporcional, definido como:
$$N \to \infty, \quad K = \kappa N, \quad P = \rho N, \quad B = \beta P$$
donde $\kappa, \rho, \beta$ son constantes fijas.

Definición del Modelo

El estudio se centra en una regla de actualización de DenseAM superficial que realiza un solo paso de descenso de gradiente en un paisaje de energía:
$$f(x) = s_2 W^\top \sigma_C [s_1 W g(x) + b] + c$$
donde $W \in \mathbb{R}^{K \times N}$ son pesos atados, $g$ es una no linealidad fija y $\sigma_C$ es una función de activación potencialmente centrada. El objetivo de entrenamiento es una pérdida de eliminación de ruido:
$$L(W) = \frac{1}{2B} \sum_{\mu=1}^B \| f_W(x_\mu + \epsilon_\mu) - x_\mu \|^2$$

Modificaciones Técnicas Clave

Los autores identifican que las parametrizaciones estándar fallan para DenseAMs con activaciones no lineales debido a "picos inducidos por la media" en la matriz Gram de las activaciones. Para abordar esto, introducen la centrado:

• Operador de Centrado: Se aplica una matriz de proyección $C = I_K - \frac{1}{K}\mathbf{1}_K\mathbf{1}_K^\top$ a las pre-activaciones y post-activaciones.
• Necesidad: Para activaciones no lineales (ReLU, Softmax) con medias no nulas, la falta de centrado conduce a un pico de rango uno en la matriz Gram de activación ($S^\top S$) que escala de manera diferente a los autovalores principales. Esta discrepanza causa inestabilidad numérica e impide la transferencia de HP. El centrado elimina este pico, asegurando que todas las estadísticas internas escalen de manera consistente.

Enfoque de Derivación

Los autores adaptan la heurística $\mu$P realizando un análisis dimensional para satisfacer tres desiderata:

1. Estabilidad: Las estadísticas de inicialización (media/varianza) de las pre-activaciones ($Z$) y las salidas ($F$) deben permanecer $O(1)$ a medida que $N, K \to \infty$.
2. Maximalidad: Las actualizaciones de parámetros ($\Delta Z, \Delta F$) después de $T$ pasos deben permanecer $O(1)$.
3. Equilibrio: Los diversos términos que contribuyen a las actualizaciones (que surgen de las dos apariciones de los pesos atados $W$) deben estar equilibrados en magnitud.

Al analizar los gradientes de la pérdida MSE y las actualizaciones resultantes para los optimizadores SGD y Adam, los autores derivan leyes de escalado específicas para las escalas de inicialización ($s_1, s_2$) y las tasas de aprendizaje ($\eta_W, \eta_b, \eta_c$).

Contribuciones y Resultados Clave

1. Prescripciones de Escalado Explícitas

El artículo proporciona una tabla de reglas de escalado (Tabla 1) para diferentes funciones de activación y optimizadores. Los hallazgos clave incluyen:

• Activaciones Lineales: En el régimen proporcional, $s_1 \sim N^{-1/2}$, $s_2 \sim K^{-1/2}$ y $\eta_W \sim K$ (para SGD).
• Activaciones tipo ReLU: El centrado es crítico. Sin centrado, la transferencia de HP falla debido a la inestabilidad. Con centrado, las prescripciones lineales se mantienen para SGD y Adam. Adam permanece estable incluso para ReLUs de alta potencia ($p=10$), mientras que SGD se vuelve inestable con tasas de aprendizaje grandes.
• Activaciones Softmax:
  • Inestabilidad de SGD: SGD sufre inestabilidad numérica a grandes escalas debido a la "retroalimentación de amplificación por filas", donde las filas con pesos softmax más grandes reciben actualizaciones desproporcionadamente grandes, lo que lleva a la localización y divergencia. El artículo no proporciona una prescripción estable de SGD para softmax.
  • Éxito de Adam: Adam es invariante a la escala y exhibe una transferencia exitosa de HP para DenseAMs de softmax centrado con $s_1 \sim N^{-1/2}$, $s_2 \sim K^{1/2}$ y $\eta_W \sim 1$.

2. Validación Empírica

Los autores validan estos hallazgos teóricos en múltiples configuraciones:

• Escalado Proporcional: Los experimentos con datos gaussianos isotrópicos y MNIST (mediante muestreo de plaquetas) demuestran que la dinámica de entrenamiento (curvas de pérdida, covarianzas de características internas) colapsa a través de escalas cuando se utilizan las prescripciones derivadas.
• Datos Anisotrópicos: Las prescripciones funcionan sorprendentemente bien incluso para datos anisotrópicos y realistas (MNIST), a pesar de que las derivaciones teóricas asumen entradas isotrópicas.
• Centrado vs. No Centrado: Los resultados empíricos (Figuras 7, 8, 10, 11, 13) muestran claramente que los modelos no centrados divergen o fallan al transferir hiperparámetros, mientras que los modelos centrados exhiben dinámicas consistentes y transferencia exitosa.

3. Colapso Dinámico

Más allá de la simple transferencia de pérdida, el artículo demuestra que la parametrización propuesta conduce a un colapso de la dinámica de entrenamiento. Las estadísticas internas, como el RMS de las actualizaciones de pesos y las covarianzas de características, permanecen constantes a través de diferentes tamaños de modelo ($N, K$) cuando se aplican las leyes de escalado.

Significado y Afirmaciones

El artículo afirma ser el primer trabajo que inicia el estudio de la transferencia de hiperparámetros para Memorias Asociativas Densas. Su importancia radica en:

• Cerrar una Brecha: Extiende la teoría de la transferencia de hiperparámetros desde las redes de alimentación hacia adelante estándar hasta modelos basados en energía con pesos atados y activaciones empinadas.
• Utilidad Práctica: Proporciona reglas accionables y fundamentadas teóricamente para escalar DenseAMs, permitiendo a los practicantes entrenar modelos a gran escala (por ejemplo, Transformers de Energía, NRGPT) sin un ajuste costoso de hiperparámetros a gran escala.
• Perspectiva Arquitectónica: Identifica el centrado como un requisito fundamental para la estabilidad de las DenseAMs con activaciones no lineales, un matiz previamente no abordado en el contexto de la transferencia de HP.

Los autores mantienen modestia respecto al alcance, señalando que su análisis asume datos isotrópicos para las derivaciones (aunque empíricamente robusto de lo contrario) y que una descripción analítica completa de la dinámica de entrenamiento a escala infinita para DenseAMs no lineales sigue siendo un tema para trabajos futuros. No afirman resolver la transferencia de HP para todas las arquitecturas modernas, sino específicamente para la familia DenseAM.