Conservation Laws for Modern Neural Architectures

Resumen Técnico: Leyes de Conservación para Arquitecturas Neuronales Modernas

Planteamiento del Problema
Comprender la dinámica de entrenamiento de las redes neuronales sobreparametrizadas es fundamental para explicar su éxito. Un mecanismo clave en esta comprensión es el concepto de leyes de conservación: cantidades que permanecen invariantes a lo largo de las trayectorias de optimización (específicamente, el flujo de gradiente). Si bien estas leyes están bien caracterizadas para redes lineales y ReLU, permanecen en gran medida inexploradas para las arquitecturas modernas. Los marcos de trabajo existentes luchan por proporcionar caracterizaciones completas para modelos contemporáneos que utilizan activaciones suaves (por ejemplo, GELU, SiLU), mecanismos de atención complejos (Atención Multicabezal con codificaciones posicionales) y diseños de Mezcla de Expertos (MoE). Además, el trabajo previo sobre la Atención Multicabezal (MHA) se limitaba a configuraciones de una sola cabeza, dejando el caso de múltiples cabezas como un problema abierto.

Metodología
El artículo desarrolla un marco teórico unificado para caracterizar las leyes de conservación para modelos de aprendizaje profundo modernos. El enfoque se basa en el trabajo de Marcotte et al. (2023; 2025), pero diverge en su estrategia analítica:

1. Formulación del Problema: Los autores definen una ley de conservación $h(\theta)$ como una función invariante bajo el flujo de gradiente euclidiano $\dot{\theta} = -\nabla L_D(\theta)$. Para hacer el problema tratable, adoptan el supuesto de que las leyes de conservación deben cumplirse para cualquier conjunto de datos, aislando el sesgo implícito inducido únicamente por la arquitectura y la dinámica de optimización.
2. Reducción a EDPs: En lugar de depender de la maquinaria de la teoría de Lie, los autores reducen la condición de conservación a un sistema de ecuaciones diferenciales parciales (EDP). Específicamente, requieren que el gradiente de la ley de conservación $\nabla_\theta h$ sea ortogonal al espacio generado por los vectores Jacobianos del modelo para todas las entradas.
3. Técnicas de Análisis Complejo: La estrategia central de la demostración implica ver las condiciones de ortogonalidad como funciones de valor complejo. Al restringir estas funciones a subespacios específicos y utilizar el teorema de la identidad para funciones meromorfas, los autores derivan restricciones independientes de la entrada sobre los gradientes de $h$.
4. Caracterización mediante Invariantes: Las restricciones de las EDP resultantes se resuelven para identificar "invariantes característicos". Se demuestra que las leyes de conservación son constantes en los componentes conectados de los conjuntos de nivel definidos por estos invariantes.

Contribuciones Clave y Resultados
El artículo proporciona caracterizaciones completas de las leyes de conservación para tres clases primarias de arquitecturas modernas:

1. Redes de Propagación hacia Adelante (FFNs):

   • GELU y SiLU: Los autores demuestran que para las FFNs que utilizan activaciones GELU o SiLU, las únicas leyes de conservación son las funciones constantes. No existen invariantes no triviales.
   • SwiGLU: En contraste, las arquitecturas SwiGLU (que involucran una interacción multiplicativa entre matrices de parámetros) exhiben leyes de conservación no triviales. Los invariantes están determinados por las cantidades $\|A_{:,i}\|^2 - \|C_{i,:}\|^2$ para cada unidad oculta $i$, donde $A$ y $C$ son las matrices de pesos involucradas en el mecanismo de compuerta (gating).

2. Atención Multicabezal (MHA):

   • MHA Vanilla: Los autores resuelven el problema abierto relativo a la atención de múltiples cabezas, confirmando que las leyes de conservación están caracterizadas por los invariantes $Q_i^\top Q_i - K_i^\top K_i$ y $V_i^\top V_i - O_i^\top O_i$ para cada cabeza $i$.
   • Codificación Posicional Sinusoidal (PE): Dado que la PE sinusoidal actúa como un desplazamiento biyectivo fijo al la entrada, la estructura de conservación permanece idéntica al caso de la MHA vanilla.
   • Codificación Posicional Rotatoria (RoPE): RoPE altera fundamentalmente la estructura interna de la atención. Los autores derivan un conjunto distinto de invariantes: $\|Q_i^{(j)}\|_F^2 - \|K_i^{(j)}\|_F^2$ (donde $Q$ y $K$ se descomponen en bloques de $2 \times 2$) y $V_i^\top V_i - O_i^\top O_i$. Esto demuestra que RoPE introduce una clase diferente de leyes de conservación en comparación con las codificaciones posicionales absolutas.

3. Mezcla de Expertos (MoE):

   • MoE Densa y Dispersa: El estudio cubre tanto la compuerta densa como la MoE dispersa (SMoE) con enrutamiento Top-$k$. Los resultados muestran que las leyes de conservación para SMoE coinciden con las de su correspondiente MoE densa.
   • Mecanismos de Compuerta (Gating): Los invariantes se localizan al nivel de experto (siguiendo los invariantes de SwiGLU para cada experto) e incluyen un invariante adicional para los pesos de la compuerta: la suma de los vectores de peso de la compuerta $\sum W_i$.
   • Compuerta Sigmoide: Los autores demuestran que reemplazar la compuerta softmax con una compuerta sigmoide normalizada no altera el conjunto de leyes de conservación; los invariantes siguen siendo idénticos al caso con compuerta softmax.

Validación Experimental
Los hallazgos teóricos se validan empíricamente a través de tareas de modelado de lenguaje (WikiText-103, Penn Treebank) y visión computacional (CIFAR-10, ImageNet-1K).

• Configuración: Los experimentos utilizan descenso de gradiente de lote completo (para aproximar el flujo de tiempo continuo) y SGD de mini-lote en conjuntos de datos de pequeña y gran escala.
• Métricas: Los autores rastrean el "error de conservación" (desviación de los invariantes de sus valores iniciales) a lo largo de los pasos de entrenamiento.
• Hallazgos: Los experimentos confirman que las leyes de conservación permanecen acotadas durante el entrenamiento. Específicamente, el error escala con la tasa de aprendizaje $\tau$ y el número de pasos $k$ según predice la teoría ($O(\tau^2 k)$ para tamaños de paso constantes). Los resultados se mantienen consistentes a través de MHA, SwiGLU FFNs, RoPE y mecanismos de compuerta de MoE, independientemente del conjunto de datos o del régimen de entrenamiento.

Significancia y Reivindicaciones
El artículo afirma avanzar en la comprensión teórica de la dinámica de entrenamiento al proporcionar la primera caracterización completa de las leyes de conservación para un amplio espectro de arquitecturas de aprendizaje profundo modernas.

• Unificación Teórica: Establece un marco unificado que extiende los resultados previos de redes lineales/ReLU a activaciones suaves y mecanismos de atención complejos.
• Resolución de Problemas Abiertos: Resuelve explícitamente el problema abierto de caracterizar las leyes de conservación para la atención de múltiples cabezas y extiende el análisis para incluir codificaciones posicionales y compuertas de MoE.
• Sesgo Implícito: Al identificar estos invariantes, el trabajo arroja luz sobre las restricciones geométricas inherentes a las arquitecturas modernas, revelando propiedades que persisten desde la inicialización hasta la convergencia.
• Limitaciones: Los autores señalan modestamente que su dependencia de técnicas de análisis complejo (extensiones meromorfas) limita actualmente la aplicación del marco a arquitecturas que involucran componentes no meromorfos, tales como la normalización de capa (layer normalization) o composiciones profundas que involucren funciones como $\sqrt{x}$ o $e^{1/x}$. También reconocen que el supuesto de independencia del conjunto de datos, aunque es un principio para aislar el sesgo arquitectónico, puede ser visto como restrictivo en contextos más amplios.