Learning What to Predict: Downstream-Guided Task Design for Continued Pretraining

Resumo Técnico: Aprendendo o que Predizer: Design de Tarefas Guiado por Downstream para Pré-treinamento Continuado

1. Declaração do Problema

O pré-treinamento continuado é um método padrão para adaptar modelos de fundação a novos domínios (ex: raciocínio matemático ou predição densa de visão). No entanto, a prática atual depende de um ciclo de feedback grosseiro. A unidade de otimização é uma única atualização autossupervisionada (ex: predição do próximo token ou reconstrução mascarada), mas a unidade de seleção é toda a execução do pré-treinamento avaliada em benchmarks de downstream.

Praticantes tipicamente avaliam checkpoints, revisam manualmente misturas de dados ou objetivos e reiniciam execuções. Esse processo torna as atualizações individuais do pré-treinamento "cegas" para saber se elas contribuem para capacidades específicas de downstream. Embora o ajuste fino (SFT, RL) utilize rótulos de downstream para atualizar diretamente o aprendiz, isso substitui a atualização autossupervisionada em vez de guiá-la. O artigo questiona: Pode um pequeno conjunto de exemplos de downstream verificáveis fornecer feedback ao nível de passo durante o pré-treinamento continuado para melhorar capacidades alvo sem supervisionar diretamente o aprendiz ou prejudicar a generalização?

2. Metodologia: V-pretraining

Os autores introduzem o V-pretraining (pré-treinamento baseado em valor), um framework que separa o processo de aprendizado em dois componentes:

1. O Aprendiz (Learner): Treinado exclusivamente via uma perda autossupervisionada em um fluxo não rotulado. Seus parâmetros ($\theta$) nunca são atualizados por gradientes de downstream.
2. O Designer de Tarefas (Task Designer): Um módulo leve com parâmetros ($\phi$) que constrói dinamicamente a tarefa autossupervisionada (ex: selecionando alvos ou visões) para o lote (batch) não rotulado.

Mecanismo Central:
O designer utiliza um pequeno conjunto de feedback ($D_{fb}$) contendo rótulos de downstream para estimar o "valor" de uma candidata atualização autossupervisionada.

• Alinhamento de Gradiente: Seja $g_{pre}$ o gradiente da proposta de atualização autossupervisionada e $g_{down}$ o gradiente da perda de downstream no conjunto de feedback.
• Objetivo de Valor: O artigo postula que a mudança de primeira ordem na perda de downstream é aproximada pelo produto escalar $V = g_{down}^\top g_{pre}$.
• Ciclo de Treinamento:
  1. O designer constró de uma tarefa (ex: uma distribuição de alvos suaves ou uma visão específica) para o lote não rotulado atual.
  2. O designer é atualizado para maximizar $V$ (ou seja, para alinhar o gradiente autossupervisionado com a direção que reduz a perda de downstream).
  3. Os alvos ou visões construídos são desconectados (aplicação de stop-gradient).
  4. O aprendiz é atualizado usando a perda autossupervisionada padrão com esses alvos desconectados.

Instanciações:

• Linguagem: O designer aprende alvos suaves top-K adaptativos para a predição do próximo token. Em vez de um alvo one-hot rígido, o designer produz uma distribuição de probabilidade sobre um pequeno conjunto de candidatos (incluindo o token real e os $K-1$ principais tokens previstos pelo modelo), ponderados por uma constante de portão (gating).
• Visão: O designer aprende visões ou máscaras por instância para o aprendizado autossupervisionado do tipo DINO, modificando os cortes (crops) ou aumentações de entrada para melhor alinhar com tarefas de predição densa de downstream (segmentação, estimativa de profundidade).

3. Principais Contribuições

• Conceitual: O artigo reformula o pré-treinamento continuado guiado por downstream como construção de tarefa aprendível. Ele introduz um canal para orientação que é indireto para o aprendiz (via design de tarefa), mas direto para o designer, evitando a substituição de atualizações autossupervisionadas por supervisionadas.
• Algorítmica: Propõe um objetivo de nível de passo computacionalmente eficiente ($g_{down}^\top g_{pre}$) que serve como um surrolato de primeira ordem para a melhoria de downstream. Isso evita a intratabilidade de diferenciar trajetórias completas de pré-treinamento (otimização bilevel).
• Empírica: Experimentos extensos pareados por tempo de execução (wall-clock-matched) demonstram que o V-pretraining melhora as capacidades alvo sem colapsar as representações de propósito geral.

4. Resultados Experimentais

Os autores conduziram experimentos pareados por computação (mesmo tempo de execução, hardware e condições iniciais) através de modalidades de linguagem e visão.

Linguagem (Raciocínio Matemático):

• Configuração: Pré-treinamento continuado de modelos Qwen no NuminaMath-CoT usando 1.024 exemplos de GSM8K como feedback.
• Resultados: O V-pretraining melhorou significativamente o Pass@1 do GSM8K.
  • Qwen2.5-0.5B: Ganho de execução única de +7,4 pontos (22,20 $\to$ 29,60), uma melhoria relativa de 33%.
  • Qwen1.5-4B: Melhorou de 56,48 para 58,98.
• Generalização: O método não prejudicou as capacidades amplas. As pontuações MMLU permaneceram estáveis para os modelos de 4B e 7B, embora o modelo de 0,5B tenha mostrado uma leve queda, sugerindo que aprendizes pequenos são mais sensíveis a feedbacks estreitos.

Visão (Predição Densa):

• Configuração: Treinamento autossupervisionado continuado de backbones DINOv3 no ImageNet-1K usando ADE20K (segmentação) e NYUv2 (profundidade) como feedback.
• Resultados:
  • ADE20K: O mIoU melhorou de 51,33 para 52,40 para o ViT-L.
  • NYUv2: O RMSE diminuiu de 0,5752 para 0,5522 para o ViT-L.
  • Generalização: A acurácia de classificação linear no ImageNet foi preservada ou levemente melhorada, indicando que o modelo não se especializou excessivamente nas tarefas densas em detrimento do reconhecimento global.
  • Transferência: A recuperação de instâncias nos datasets Oxford/Paris mostrou melhorias na maioria dos protocolos, confirmando que as representações permaneceram robustas.

Ablações e Controles:

• Descontaminação: Remover duplicatas próximas de GSM8K/MATH do fluxo de treinamento não eliminou os ganhos, descartando a simples memorização.
• Eliminação de Atalhos: Os ganhos não foram devidos a suavização genérica, autodestilação ou ruído aleatório. Variantes usando feedback aleatório ou suavização uniforme tiveram desempenho inferior ao V-pretraining alinhado.
• Sobrecarga (Overhead): O designer de tarefas introduz uma redução de ~9,4% no throughput de tokens e um aumento de ~16,4% na memória de pico, mas esses custos são contabilizados nas comparações pareadas por tempo de execução.

5. Significância e Alegações

O artigo afirma que o V-pretraining oferece um complemento computacionalmente eficiente aos métodos diretos de pós-treinamento (como SFT ou RL).

• Eficiência: Permite a injeção de sinais de capacidade de downstream durante a fase de pré-treinamento continuado, potencialmente reduzindo a necessidade de ajuste fino de downstream caro posteriormente.
• Mecanismo: Demonstra que o feedback de downstream pode direcionar a construção de tarefas autossupervisionadas (o que o modelo prediz) em vez do próprio aprendiz (como o modelo atualiza).
• Modéstia: Os autores declaram explicitamente que o objetivo de valor fornece uma garantia local de um passo (maximizando um limite inferior para a melhoria imediata de downstream), em vez de um teorema de convergência de longo horizonte. Eles reconhecem que o método pode levar modelos pequenos ao sobreajuste (over-steer) para tarefas estreitas e que o sinal de alinhamento pode ser ruidoso.

Em resumo, o V-pretraining estabelece que um pequeno conjunto de exemplos de downstream pode guiar efetivamente o aspecto de "o que predizer" do pré-treinamento continuado, melhorando capacidades específicas enquanto mantém as propriedades de generalização do modelo de fundação.