Resumo Técnico: Quando as Perdas se Alinham: Ponderação de Perda Composta Baseada em Gradiente para Pré-treinamento Eficiente

Formulação do Problema

Sistemas modernos de aprendizado profundo frequentemente empregam um paradigma de duas etapas: pré-treinamento em dados de grande escala, muitas vezes não rotulados, seguido por adaptação a tarefas downstream. Em muitos pipelines de pré-treinamento, a função objetivo não é uma única perda, mas uma combinação ponderada de múltiplos termos de perda (por exemplo, previsão do próximo token, perdas contrastivas ou objetivos de regressão). Os pesos relativos desses termos são hiperparâmetros críticos que influenciam significativamente a qualidade das representações aprendidas e o desempenho downstream.

Atualmente, a seleção desses pesos depende de ajuste manual, busca aleatória ou otimização bayesiana. Esses métodos são computacionalmente proibitivos para pré-treinamento em grande escala, pois exigem inúmeras execuções de treinamento independentes. Além disso, padrões ingênuos, como ponderação uniforme, frequentemente falham em contabilizar as contribuições variáveis e os conflitos entre diferentes termos de perda. O artigo aborda a necessidade de um método eficiente para ajustar pesos de perda lineares em objetivos de pré-treinamento compostos, a fim de maximizar o desempenho em tarefas downstream sem incorrer no custo de execuções de treinamento completas repetidas.

Metodologia: Pré-treinamento Alinhado a Gradientes (GraP)

Os autores propõem o Pré-treinamento Alinhado a Gradientes (GraP), um método de otimização bilevel baseado em gradiente que aprende pesos de perda online. A ideia central é alinhar o gradiente composto de pré-treinamento com o gradiente de um objetivo downstream.

1. Formulação de Otimização Bilevel

O problema é formulado como uma tarefa de otimização bilevel, onde o loop externo otimiza os pesos de perda $w$ para minimizar a perda downstream $L_{down}$, e o loop interno otimiza os parâmetros do modelo $\theta$ para minimizar a perda de pré-treinamento composta $L_{pre}(w)$.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{sujeito a} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. Regra de Atualização Baseada em Gradiente

Usando uma abordagem de retropropagação truncada (SGD de passo único), os autores derivam uma regra de atualização para os pesos. O gradiente da perda downstream em relação aos pesos é proporcional ao produto interno entre o gradiente downstream e os gradientes individuais de pré-treinamento:
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
onde $g_{down}$ é o gradiente downstream e $g_k$ é o gradiente da $k$-ésima perda de pré-treinamento. O objetivo é maximizar o alinhamento entre o gradiente composto $\sum w_k g_k$ e o gradiente downstream.

3. Aproximação no Espaço de Embedding

Uma implementação ingênua exigiria calcular $K$ passadas de retropropagação completas através do modelo para obter $g_k$. Para reduzir o custo computacional, o GraP calcula gradientes em relação ao embedding de backbone compartilhado $z$ em vez do espaço completo de parâmetros.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
Como as perdas são aplicadas via cabeças separadas sobre o mesmo embedding, calcular esses gradientes é significativamente mais barato do que passadas de retropropagação completas. O objetivo de otimização torna-se maximizar o alinhamento neste espaço de representação compartilhado:
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. Estratégia de Normalização

O artigo argumenta contra a normalização dos próprios pesos (por exemplo, $\sum w_k = 1$) ou de sua norma $L_2$, pois essas podem distorcer a direção ótima do gradiente composto. Em vez disso, o GraP impõe uma restrição sobre a norma $L_2$ do próprio gradiente composto ($\|\sum w_k g_k\| = 1$). Isso garante que a otimização se concentre na direção do vetor de atualização em vez de sua magnitude, recuperando a direção de alinhamento ótima sob o objetivo local.

5. Eficiência do Algoritmo

O algoritmo resultante requer apenas uma passada forward e uma passada backward através do backbone por iteração. Os gradientes computados para as cabeças são reutilizados tanto para atualizar os pesos de perda quanto para atualizar os parâmetros do backbone. Isso reduz a sobrecarga do ajuste de hiperparâmetros para aproximadamente 30% acima de uma única execução de treinamento, comparado aos custos de ordens de magnitude da otimização bayesiana.

Contribuições Principais

1. Formulação do Problema: Os autores formulam a seleção de pesos de perda de pré-treinamento como um problema de otimização bilevel visando explicitamente o desempenho downstream.
2. Insight de Normalização: Eles analisam estratégias de normalização e demonstram que restringir a norma do vetor de atualização (gradiente composto) é superior a restringir os próprios pesos, preservando a direção de alinhamento ótima.
3. Algoritmo Eficiente: Eles propõem o GraP, um algoritmo com complexidade $O(1)$ em relação ao número de passos forward-backward em relação ao número de hiperparâmetros. Ele aproveita o alinhamento de gradientes no espaço de representação compartilhado para evitar múltiplas passadas de retropropagação.
4. Eficiência Empírica: O método alcança desempenho comparável ou superior à otimização padrão de hiperparâmetros (por exemplo, busca bayesiana) enquanto é quase 50$\times$ mais eficiente, permitindo o ajuste em cenários onde métodos tradicionais são computacionalmente inviáveis.
5. Identificação de Redundância: A abordagem demonstra a capacidade de identificar e reduzir o peso de termos de perda redundantes (por exemplo, atribuir peso próximo de zero a perdas contrastivas específicas em tarefas de visão), atuando efetivamente como um mecanismo automático de seleção de perdas.

Resultados Experimentais

O método foi avaliado em dois domínios: modelagem de sequências de eventos e visão computacional auto-supervisionada.

Modelagem de Sequências de Eventos

• Conjuntos de Dados: Churn, AgePred, AlfaBattle, MIMIC-III e Taobao.
• Configuração: O pré-treinamento envolveu um objetivo composto de previsão do próximo token (NTP) em múltiplos campos e uma perda contrastiva.
• Desempenho: O GraP igualou ou melhorou baselines cuidadosamente ajustadas (incluindo GradNorm, DWA, MGDA e PCGrad) e pesos ajustados pelo Optuna. Consistentemente alcançou os melhores resultados nos conjuntos de dados Churn e MIMIC-III em backbones RNN, Transformer e HT-Transformer.
• Eficiência: Comparado à otimização bayesiana (Optuna), que requer dezenas de execuções independentes, o GraP alcançou desempenho similar com uma única execução de treinamento e uma sobrecarga moderada (desaceleração média de 44% na vazão).

Visão Computacional Auto-Supervisionada

• Configuração: Avaliado em CIFAR-10, CIFAR-100 e ImageNet-100 usando o framework All4One (combinando objetivos Barlow Twins, BYOL e NNCLR).
• Desempenho: Neste cenário, onde os objetivos são frequentemente bem equilibrados, o GraP alcançou desempenho comparável aos parâmetros padrão do All4One e a outros métodos de aprendizado multi-tarefa.
• Seleção de Perda: Notavelmente, o GraP atribuiu peso próximo de zero à perda Attn-NNCLR. Remover manualmente essa perda resultou em desempenho similar, validando a capacidade do método de identificar objetivos de baixo impacto ou redundantes.

Estudos de Ablação

• Vazamento de Rótulos: Experimentos variando a fração de dados rotulados usados para ajuste de pesos mostraram robustez, sugerindo que os rótulos fornecem um sinal de orientação fraco em vez de supervisão direta para o backbone.
• Normalização: A normalização de gradiente composto proposta superou a normalização padrão de soma de pesos.
• Treinamento Auxiliar: O pré-treinamento com GraP seguido de ajuste fino supervisionado superou o treinamento auxiliar conjunto, provavelmente devido à redução do overfitting em dados rotulados limitados.

Significado e Alegações

O artigo alega que a otimização de pesos de perda baseada em gradiente é uma alternativa prática e eficiente ao aprendizado multi-tarefa clássico e à otimização de hiperparâmetros baseada em busca. Seu significado principal reside em tornar o ajuste de hiperparâmetros viável para cenários de pré-treinamento em grande escala com múltiplos termos de perda, onde métodos tradicionais são computacionalmente proibitivos.

Os autores enfatizam que o GraP não requer múltiplas execuções de treinamento, tornando-o adequado para cenários com orçamentos computacionais limitados. Embora o método assuma cabeças separadas para cada perda, ele lida efetivamente com os compromissos entre objetivos conflitantes ao alinhar a direção de atualização do pré-treinamento com o objetivo downstream. O trabalho sugere que, em pipelines de pré-treinamento complexos, o alinhamento automático e online de pesos de perda pode produzir representações que transferem tão bem, ou melhor, do que aquelas obtidas por ajuste manual caro ou busca de caixa preta.