BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization

Resumen Técnico: BROS: Subespacios Aleatorizados Corregidos por Sesgo para Optimización Bilevel de Bucle Único Eficiente en Memoria

Enunciado del Problema
La optimización bilevel estocástica (SBO) es un marco fundamental para aplicaciones de aprendizaje profundo como la optimización de hiperparámetros, la reponderación de datos, el aprendizaje de representaciones y la limpieza de hiperdatos. La formulación estándar involucra una variable de nivel superior $x$ y una variable de nivel inferior $Y$ (típicamente los parámetros de una red neuronal), donde el objetivo superior depende de la solución del problema de optimización inferior.

Los métodos existentes de SBO de bucle único exactos (por ejemplo, MA-SOBA) proporcionan garantías fuertes de convergencia (complejidad de muestra $O(\epsilon^{-2})$) pero sufren de cuellos de botella severos de memoria cuando la red de nivel inferior es grande. Estos métodos dependen de productos vectoriales de Hessiano en el espacio completo (HVP) y productos vectoriales de Jacobiano (JVP) para mantener variables auxiliares, requiriendo memoria proporcional al conteo total de parámetros. Por el contrario, los métodos sustitutos eficientes en memoria (por ejemplo, métodos de penalización, enfoques de orden cero) reducen el uso de memoria pero a menudo sacrifican garantías de convergencia exacta, requieren suposiciones más fuertes u optimizan un problema sustituto en lugar del objetivo bilevel original.

Metodología: BROS
Los autores proponen BROS (Optimización Bilevel mediante Subespacios Aleatorizados), un método de SBO de bucle único eficiente en memoria diseñado para problemas parametrizados por matrices multicapa. BROS aborda el cuello de botella de memoria realizando actualizaciones de nivel inferior y de nivel auxiliar dentro de subespacios aleatorizados de baja dimensión, preservando al mismo tiempo la estructura exacta del hipergradiente de bucle único.

Los componentes centrales de BROS son:

1. Descomposición en Subespacios Aleatorizados:
   En lugar de actualizar los parámetros completos de nivel inferior $Y \in \mathbb{R}^{m \times n}$, BROS muestrea un proyector Haar escalado $P \in \mathbb{R}^{m \times r}$ (donde $r \ll m$) en cada iteración. La actualización de nivel inferior se calcula en el subespacio reducido $B \in \mathbb{R}^{r \times n}$ mediante el objetivo proyectado $\tilde{g}(x, B) = g(x, Y + PB)$. La variable de espacio completo se actualiza como $Y^{k+1} = Y^k + P B^k$. Esto reduce la huella de memoria de los gradientes, las activaciones y las direcciones auxiliares de $O(mn)$a$O(rn)$.

2. Actualizaciones Auxiliares Corregidas por Sesgo:
   Un enfoque ingenuo de "boceto-luego-elevación" para la variable auxiliar $Z$ (que resuelve el sistema lineal que involucra el Hessiano) introduce sesgo porque la acción del Hessiano elevada $P \tilde{H} [P^\top Z]$ no es un estimador insesgado de la acción del Hessiano de espacio completo $H[Z]$. Este sesgo distorsiona el hipergradiente y puede conducir a límites no estacionarios.
   Para resolver esto, BROS introduce una corrección de doble sonda Rademacher. El algoritmo muestrea vectores Rademacher independientes (vectores de signo) $u$ y $\xi$ para estimar términos de molestia tipo traza y tipo transpuesta inducidos por la proyección. Estas sondas permiten construir un estimador de HVP corregido, denotado como $\widehat{H}[Z]$, que es condicionalmente insesgado para la acción del Hessiano de espacio completo sin requerir la formación de matrices de espacio completo.

3. Recursión de Bucle Único:
   BROS mantiene una recursión acoplada de bucle único sobre la variable superior $x$, la variable inferior $Y$, la variable auxiliar $Z$ y un estimador de hipergradiente de media móvil $h$. Las actualizaciones utilizan el gradiente inferior proyectado y la dirección auxiliar corregida por sesgo.

Contribuciones Clave

• Recursión Eficiente en Memoria: BROS descompone las recursiones de nivel inferior y auxiliar en subespacios aleatorizados, reduciendo significativamente la dimensionalidad de los gradientes almacenados, las activaciones y los cálculos de HVP/JVP.
• Recuperación de Hipergradiente Insesgado: El artículo identifica el sesgo específico inducido por la proyección en las acciones del Hessiano auxiliar elevadas y propone una corrección de doble sonda Rademacher por capa. Esto asegura que el estimador permanezca insesgado para la acción del Hessiano multicapa completo sin imponer suposiciones estructurales adicionales (por ejemplo, Hessianos diagonales por bloques).
• Garantías Teóricas: Bajo suposiciones estándar (suavidad, convexidad fuerte del nivel inferior), los autores demuestran que BROS preserva la complejidad de muestra $O(\epsilon^{-2})$ de los métodos exactos de bucle único (específicamente MA-SOBA) para encontrar un punto $\epsilon$-estacionario. La tasa de convergencia depende de la relación de proyección $\omega_P = \max(m_\ell/r_\ell)$, recuperando la tasa de espacio completo cuando $r=m$.

Resultados Experimentales
Los autores evalúan BROS en cuatro tareas de aprendizaje bilevel:

1. Limpieza de Hiperdatos MNIST: BROS logra una precisión de prueba del 84.30%, coincidiendo estrechamente con la línea base de espacio completo MA-SOBA (84.46%) y superando a los métodos basados en penalización y de orden cero.
2. Aprendizaje de Mezcla de Datos (modelos estilo GPT de 280M): BROS iguala el rendimiento de MA-SOBA (perplejidad logarítmica final 2.8109 vs. 2.8098) mientras reduce la memoria máxima de entrenamiento de proxy en aproximadamente 44.9% en comparación con MA-SOBA y un 27% en comparación con los métodos de Penalización.
3. Aprendizaje de Hiperrepresentación CIFAR-10: BROS logra una precisión de prueba final del 78.03%, comparable a MA-SOBA (78.57%).
4. Reponderación de Muestras ViT CIFAR-100: BROS logra la mejor precisión de validación (49.88%) y prueba (49.05%) entre los métodos comparados, utilizando solo 6.42 GB de memoria máxima en comparación con más de 23 GB para las líneas base de espacio completo.

Significado y Afirmaciones
El artículo afirma que BROS cierra exitosamente la brecha entre la eficiencia de memoria y la convergencia teórica en la optimización bilevel estocástica. A diferencia de los métodos sustitutos que intercambian exactitud por eficiencia, BROS mantiene la estructura exacta de actualización de hipergradiente de bucle único y su orden de convergencia asociado $O(\epsilon^{-2})$. Al utilizar subespacios aleatorizados y un mecanismo novedoso de corrección de sesgo, BROS habilita tareas de optimización bilevel (como la reponderación de datos y el ajuste de hiperparámetros) en redes neuronales grandes que de otro modo serían inviables debido a restricciones de memoria, reduciendo el uso de memoria máxima hasta en un 44.9% mientras coincide estrechamente con el rendimiento de espacio completo.