Resumen Técnico: Cuando las Pérdidas se Alinean: Ponderación de Pérdidas Compuestas Basada en Gradientes para el Preentrenamiento Eficiente

Formulación del Problema

Los sistemas modernos de aprendizaje profundo emplean frecuentemente un paradigma de dos etapas: preentrenamiento sobre datos a gran escala, a menudo no etiquetados, seguido de adaptación a tareas posteriores. En muchas tuberías de preentrenamiento, la función objetivo no es una única pérdida, sino una combinación ponderada de múltiples términos de pérdida (por ejemplo, predicción del siguiente token, pérdidas de contraste u objetivos de regresión). Los pesos relativos de estos términos son hiperparámetros críticos que influyen significativamente en la calidad de las representaciones aprendidas y en el rendimiento posterior.

Actualmente, la selección de estos pesos depende del ajuste manual, la búsqueda aleatoria o la optimización bayesiana. Estos métodos son computacionalmente prohibitivos para el preentrenamiento a gran escala porque requieren numerosas ejecuciones de entrenamiento independientes. Además, los valores predeterminados ingenuos, como la ponderación uniforme, a menudo no tienen en cuenta las contribuciones variables y los conflictos entre los diferentes términos de pérdida. El artículo aborda la necesidad de un método eficiente para ajustar los pesos lineales de pérdida en objetivos de preentrenamiento compuestos para maximizar el rendimiento de la tarea posterior sin incurrir en el costo de ejecuciones completas de entrenamiento repetidas.

Metodología: Preentrenamiento Alineado por Gradientes (GraP)

Los autores proponen Preentrenamiento Alineado por Gradientes (GraP), un método de optimización multinivel basado en gradientes que aprende los pesos de pérdida en línea. La idea central es alinear el gradiente de preentrenamiento compuesto con el gradiente de un objetivo posterior.

1. Formulación de Optimización Multinivel

El problema se plantea como una tarea de optimización multinivel donde el bucle externo optimiza los pesos de pérdida $w$ para minimizar la pérdida posterior $L_{down}$, y el bucle interno optimiza los parámetros del modelo $\theta$ para minimizar la pérdida de preentrenamiento compuesta $L_{pre}(w)$.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{sujeto a} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. Regla de Actualización Basada en Gradientes

Utilizando un enfoque de retropropagación truncada (descenso de gradiente estocástico de un solo paso), los autores derivan una regla de actualización para los pesos. El gradiente de la pérdida posterior con respecto a los pesos es proporcional al producto interno entre el gradiente posterior y los gradientes individuales de preentrenamiento:
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
donde $g_{down}$ es el gradiente posterior y $g_k$ es el gradiente de la $k$-ésima pérdida de preentrenamiento. El objetivo es maximizar la alineación entre el gradiente compuesto $\sum w_k g_k$ y el gradiente posterior.

3. Aproximación en el Espacio de Incrustaciones

Una implementación ingenua requeriría calcular $K$ pasadas de retropropagación completas a través del modelo para obtener $g_k$. Para reducir el costo computacional, GraP calcula los gradientes con respecto a la incrustación de la columna vertebral compartida $z$ en lugar del espacio completo de parámetros.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
Dado que las pérdidas se aplican mediante cabezales separados sobre la misma incrustación, calcular estos gradientes es significativamente más barato que realizar pasadas de retropropagación completas. El objetivo de optimización se convierte en maximizar la alineación en este espacio de representación compartido:
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. Estrategia de Normalización

El artículo argumenta en contra de normalizar los pesos mismos (por ejemplo, $\sum w_k = 1$) o su norma $L_2$, ya que estos pueden distorsionar la dirección óptima del gradiente compuesto. En su lugar, GraP impone una restricción sobre la norma $L_2$ del propio gradiente compuesto ($\|\sum w_k g_k\| = 1$). Esto asegura que la optimización se centre en la dirección del vector de actualización en lugar de en su magnitud, recuperando la dirección de alineación óptima bajo el objetivo local.

5. Eficiencia del Algoritmo

El algoritmo resultante requiere solo una pasada hacia adelante y una pasada hacia atrás a través de la columna vertebral por iteración. Los gradientes calculados para los cabezales se reutilizan tanto para actualizar los pesos de pérdida como para actualizar los parámetros de la columna vertebral. Esto reduce la sobrecarga del ajuste de hiperparámetros a aproximadamente un 30% por encima de una sola ejecución de entrenamiento, en comparación con el costo de órdenes de magnitud de la optimización bayesiana.

Contribuciones Clave

1. Formulación del Problema: Los autores formulan la selección de pesos de pérdida de preentrenamiento como un problema de optimización multinivel que apunta explícitamente al rendimiento posterior.
2. Insight de Normalización: Analizan estrategias de normalización y demuestran que restringir la norma del vector de actualización (gradiente compuesto) es superior a restringir los pesos mismos, preservando la dirección de alineación óptima.
3. Algoritmo Eficiente: Proponen GraP, un algoritmo con complejidad $O(1)$ con respecto al número de pasos de avance-retroceso en relación con el número de hiperparámetros. Aprovecha la alineación de gradientes en el espacio de representación compartido para evitar múltiples pasadas de retropropagación.
4. Eficiencia Empírica: El método logra un rendimiento comparable o superior a la optimización estándar de hiperparámetros (por ejemplo, búsqueda bayesiana) mientras es casi 50$\times$ más eficiente, permitiendo el ajuste en configuraciones donde los métodos tradicionales son computacionalmente inviables.
5. Identificación de Redundancia: El enfoque demuestra la capacidad de identificar y reducir el peso de los términos de pérdida redundantes (por ejemplo, asignar un peso cercano a cero a pérdidas de contraste específicas en tareas de visión), actuando efectivamente como un mecanismo automático de selección de pérdidas.

Resultados Experimentales

El método fue evaluado en dos dominios: modelado de secuencias de eventos y visión por computadora auto-supervisada.

Modelado de Secuencias de Eventos

• Conjuntos de datos: Churn, AgePred, AlfaBattle, MIMIC-III y Taobao.
• Configuración: El preentrenamiento involucró un objetivo compuesto de predicción del siguiente token (NTP) a través de múltiples campos y una pérdida de contraste.
• Rendimiento: GraP igualó o mejoró las líneas base cuidadosamente ajustadas (incluyendo GradNorm, DWA, MGDA y PCGrad) y los pesos ajustados con Optuna. Consistentemente logró los mejores resultados en los conjuntos de datos Churn y MIMIC-III a través de columnas vertebrales RNN, Transformer y HT-Transformer.
• Eficiencia: En comparación con la optimización bayesiana (Optuna), que requiere docenas de ejecuciones independientes, GraP logró un rendimiento similar con una sola ejecución de entrenamiento y una sobrecarga moderada (ralentización promedio del 44% en el rendimiento).

Visión por Computadora Auto-Supervisada

• Configuración: Evaluado en CIFAR-10, CIFAR-100 e ImageNet-100 utilizando el marco All4One (combinando objetivos Barlow Twins, BYOL y NNCLR).
• Rendimiento: En esta configuración, donde los objetivos a menudo están bien equilibrados, GraP logró un rendimiento comparable a los parámetros predeterminados de All4One y otros métodos de aprendizaje multitarea.
• Selección de Pérdidas: Notablemente, GraP asignó un peso cercano a cero a la pérdida Attn-NNCLR. Eliminar manualmente esta pérdida produjo un rendimiento similar, validando la capacidad del método para identificar objetivos de bajo impacto o redundantes.

Estudios de Ablación

• Fuga de Etiquetas: Los experimentos que variaron la fracción de datos etiquetados utilizados para el ajuste de pesos mostraron robustez, sugiriendo que las etiquetas proporcionan una señal de guía débil en lugar de una supervisión directa para la columna vertebral.
• Normalización: La normalización del gradiente compuesto propuesta superó a la normalización estándar de la suma de pesos.
• Entrenamiento Auxiliar: El preentrenamiento con GraP seguido de un ajuste fino supervisado superó al entrenamiento auxiliar conjunto, probablemente debido a una reducción del sobreajuste en datos etiquetados limitados.

Significado y Afirmaciones

El artículo afirma que la optimización de pesos de pérdida basada en gradientes es una alternativa práctica y eficiente al aprendizaje multitarea clásico y a la optimización de hiperparámetros basada en búsqueda. Su significado principal radica en hacer viable el ajuste de hiperparámetros para escenarios de preentrenamiento a gran escala con múltiples términos de pérdida, donde los métodos tradicionales son computacionalmente prohibitivos.

Los autores enfatizan que GraP no requiere múltiples ejecuciones de entrenamiento, lo que lo hace adecuado para configuraciones con presupuestos computacionales limitados. Si bien el método asume cabezales separados para cada pérdida, maneja efectivamente las compensaciones entre objetivos conflictivos al alinear la dirección de actualización del preentrenamiento con el objetivo posterior. El trabajo sugiere que en tuberías de preentrenamiento complejas, la alineación automática y en línea de los pesos de pérdida puede producir representaciones que se transfieren tan bien, o mejor, que las obtenidas mediante un ajuste manual costoso o una búsqueda de caja negra.