A Tale of Two Problems: Multi-Task Bilevel Learning Meets Equality Constrained Multi-Objective Optimization

Resumen Técnico: Un Cuento de Dos Problemas: Aprendizaje Bilevel Multi-Tarea Encuentra Optimización Multi-Objetivo con Restricciones de Igualdad

Enunciado del Problema
El artículo aborda el desafío del Aprendizaje Bilevel Multi-Tarea (MTBL) bajo Convexidad General en el Nivel Inferior (LLGC). Si bien la optimización bilevel (BLO) se utiliza ampliamente en el aprendizaje automático (por ejemplo, ajuste de hiperparámetros, meta-aprendizaje y alineación de LLM), los marcos teóricos y algoritmos existentes dependen predominantemente de la suposición de Convexidad Fuerte en el Nivel Inferior (LLSC). Esta suposición garantiza una solución única en el nivel inferior y un hipergradiente bien definido, pero a menudo se viola en problemas modernos y complejos de aprendizaje automático donde la función objetivo del nivel inferior es meramente convexa.

Además, los trabajos existentes sobre MTBL se limitan en gran medida a configuraciones de una sola tarea o dependen de LLSC. Al extender MTBL al escenario LLGC, surgen dos problemas críticos:

1. El hipergradiente se vuelve indefinido debido a la falta de convexidad fuerte.
2. La presencia de múltiples objetivos en el nivel superior, potencialmente conflictivos, requiere una interpretación de equilibrio de Pareto, lo que complica los criterios de optimalidad.

Metodología
Los autores proponen un marco novedoso que reformula el problema LLGC-MTBL en un problema de Optimización Multi-Objetivo con Restricciones de Igualdad (ECMO). Los pasos metodológicos centrales son:

1. Reformulación a ECMO:
   En lugar de resolver directamente el problema bilevel anidado, los autores observan que para una función convexa en el nivel inferior $g(x, y)$, la condición de optimalidad $\nabla_y g(x, y) = 0$ es tanto necesaria como suficiente. Reformulan el problema MTBL como:
   $$\min_{z} F(z) = [f_1(z), \dots, f_S(z)]^\top \quad \text{s.t.} \quad h_i(z) = 0, \forall i \in [q]$$
   donde $z = [x^\top, y^\top]^\top$ y $h_i(z)$ representa la $i$-ésima componente del gradiente en el nivel inferior $\nabla_y g(x, y)$.

2. Caracterización de la Estacionariedad de Pareto para ECMO:
   El artículo identifica que las definiciones estándar de estacionariedad de Pareto para la optimización multi-objetivo (MOO) sin restricciones son insuficientes para ECMO debido a las restricciones de factibilidad. Los autores establecen un nuevo criterio de estacionariedad de Pareto basado en KKT. Aprovechan la escalarización Chebyshev Ponderada (WC) para crear una correspondencia uno a uno entre la frontera de Pareto del problema ECMO y las soluciones de un problema escalarizado. Al analizar las condiciones Karush-Kuhn-Tucker (KKT) del problema escalarizado WC, definen una métrica de convergencia rigurosa basada en la norma del sistema residual de KKT.

3. Diseño del Algoritmo (WC-Penalty):
   Para resolver el problema ECMO, los autores proponen el algoritmo WC-Penalty. Este método:

   • Reformula el problema escalarizado WC introduciendo variables de holgura para manejar la minimización de la norma $\ell_\infty$.
   • Convierte las restricciones de igualdad (tanto las restricciones originales como las de las variables de holgura) en términos de penalización dentro de la función objetivo.
   • Resuelve el problema resultante de un solo nivel, con restricciones y un solo objetivo, utilizando Descenso de Gradiente Proyectado (GD).
   • Opera tanto en entornos deterministas como estocásticos.

Contribuciones Clave

• Extensión Teórica: Este trabajo es el primero en extender la optimización bilevel al escenario multi-tarea bajo la suposición relajada de LLGC.
• Nueva Formulación del Problema: Introduce el problema ECMO como una formulación distinta que no había sido estudiada teóricamente antes de este trabajo, abordando específicamente la falta de garantías de convergencia en tiempo finito en la literatura existente.
• Nueva Métrica de Convergencia: Los autores definen una nueva métrica de estacionariedad de Pareto basada en KKT para ECMO, resolviendo la ambigüedad de las condiciones de optimalidad en configuraciones multi-objetivo con restricciones.
• Algoritmo y Tasa de Convergencia: Proponen el algoritmo WC-Penalty y demuestran una tasa de convergencia en tiempo finito de $O(S T^{-1/2})$ hacia la estacionariedad de Pareto basada en KKT, donde $S$ es el número de objetivos y $T$ es el número total de iteraciones. Esta es la primera garantía de este tipo para problemas ECMO.
• Exploración de la Frontera de Pareto: Al variar el vector de preferencias sobre el simplex, el método explora sistemáticamente toda la frontera de Pareto.

Resultados Experimentales
Los autores validan su enfoque en dos tareas de ponderación de datos multi-objetivo:

1. Entrenamiento de Modelos de Recompensa para RLHF: Utilizando el conjunto de datos HelpSteer, el algoritmo optimiza los pesos de los datos para equilibrar múltiples métricas alineadas con humanos (utilidad, corrección, etc.). Los resultados muestran que el algoritmo WC-Penalty cubre una porción significativamente mayor de la frontera de Pareto en comparación con líneas base como MOML, MoCo y FORUM.
2. Alineación Multi-Objetivo de LLM: El método se aplica al ajuste fino de Llama-3.2-1B-Instruct (y escalado a modelos de 3B y 8B) para minimizar múltiples pérdidas de validación. El algoritmo demuestra un rendimiento superior en comparaciones de volumen hipervolumen y métrica $\epsilon$ frente a líneas base existentes de MTBL y bilevel.
3. Ejemplos Sintéticos: Experimentos sintéticos adicionales confirman la capacidad del algoritmo para recuperar la frontera de Pareto completa tanto en configuraciones ECMO como LLGC-MTBL.

Significado y Afirmaciones
El artículo afirma cerrar el ciclo entre marcos de optimización fundamentales al traducir soluciones del problema ECMO directamente de vuelta al problema MTBL original. Los autores enfatizan que su trabajo proporciona la primera garantía de convergencia en tiempo finito para problemas ECMO, una brecha que anteriormente no se había abordado en la literatura. Al relajar la suposición de LLSC, el marco propuesto extiende la aplicabilidad del aprendizaje bilevel a escenarios de aprendizaje automático más complejos y del mundo real donde no se puede asumir la convexidad fuerte. Los resultados sugieren que el método WC-Penalty ofrece tanto rigor teórico como eficiencia práctica en la exploración de soluciones óptimas de Pareto para el aprendizaje bilevel multi-tarea.